#!/usr/bin/env python3
"""Call the EAPIL gpt-image-2 image API and save decoded images.

The base URL and model are fixed by the skill. The API key may be supplied
temporarily with --api-key, or through EAPIL_API_KEY / OPENAI_API_KEY.
"""

from __future__ import annotations

import argparse
import base64
import json
import mimetypes
import os
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path
from typing import Any


BASE_URL = "https://api-n-cd.playones.com/v1"
MODEL = "gpt-image-2"
API_KEY_ENV = "EAPIL_API_KEY"
OPENAI_API_KEY_ENV = "OPENAI_API_KEY"


def fail(message: str, code: int = 1) -> None:
    print(f"error: {message}", file=sys.stderr)
    raise SystemExit(code)


def resolve_api_key(cli_value: str | None) -> str:
    for value in (
        cli_value,
        os.environ.get(API_KEY_ENV),
        os.environ.get(OPENAI_API_KEY_ENV),
    ):
        if value and value.strip():
            return value.strip()
    fail(
        f"provide --api-key or set {API_KEY_ENV} / {OPENAI_API_KEY_ENV} before running this script"
    )


def compact(values: dict[str, Any]) -> dict[str, Any]:
    return {key: value for key, value in values.items() if value is not None}


def output_path(value: str | None, default_name: str) -> Path:
    path = Path(value or default_name)
    if not path.is_absolute():
        path = Path.cwd() / path
    path.parent.mkdir(parents=True, exist_ok=True)
    return path


def request_json(path: str, payload: dict[str, Any], key: str) -> dict[str, Any]:
    body = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        f"{BASE_URL}{path}",
        data=body,
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    return read_json_response(request)


def multipart_body(fields: dict[str, str], files: list[tuple[str, Path]]) -> tuple[bytes, str]:
    boundary = f"----eapil-gpt2-image-{uuid.uuid4().hex}"
    chunks: list[bytes] = []

    def add(line: str) -> None:
        chunks.append(line.encode("utf-8"))

    for name, value in fields.items():
        add(f"--{boundary}\r\n")
        add(f'Content-Disposition: form-data; name="{name}"\r\n\r\n')
        add(f"{value}\r\n")

    for name, path in files:
        if not path.exists():
            fail(f"file not found: {path}")
        content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
        add(f"--{boundary}\r\n")
        add(
            f'Content-Disposition: form-data; name="{name}"; filename="{path.name}"\r\n'
        )
        add(f"Content-Type: {content_type}\r\n\r\n")
        chunks.append(path.read_bytes())
        add("\r\n")

    add(f"--{boundary}--\r\n")
    return b"".join(chunks), boundary


def request_multipart(
    path: str, fields: dict[str, str], files: list[tuple[str, Path]], key: str
) -> dict[str, Any]:
    body, boundary = multipart_body(fields, files)
    request = urllib.request.Request(
        f"{BASE_URL}{path}",
        data=body,
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": f"multipart/form-data; boundary={boundary}",
            "Accept": "application/json",
        },
        method="POST",
    )
    return read_json_response(request)


def read_json_response(request: urllib.request.Request) -> dict[str, Any]:
    try:
        with urllib.request.urlopen(request, timeout=300) as response:
            raw = response.read()
    except urllib.error.HTTPError as err:
        raw = err.read().decode("utf-8", "replace")
        try:
            parsed = json.loads(raw)
            message = parsed.get("error", {}).get("message") or raw[:500]
        except json.JSONDecodeError:
            message = raw[:500]
        fail(f"HTTP {err.code}: {message}")
    except urllib.error.URLError as err:
        fail(f"request failed: {err.reason}")

    try:
        return json.loads(raw)
    except json.JSONDecodeError as err:
        fail(f"response was not JSON: {err}")


def extract_image_bytes_from_item(item: Any) -> bytes:
    if not isinstance(item, dict):
        fail("response data item was not an object")

    b64_value = item.get("b64_json")
    if isinstance(b64_value, str) and b64_value.strip():
        return base64.b64decode(b64_value)

    url_value = item.get("url")
    if isinstance(url_value, str) and url_value.startswith("data:"):
        _, _, payload = url_value.partition(",")
        if payload:
            return base64.b64decode(payload)

    fail("response item did not contain b64_json or a data URL")


def extract_image_bytes(response: dict[str, Any]) -> bytes:
    data = response.get("data")
    if not isinstance(data, list) or not data:
        fail("response did not contain data[0]")

    return extract_image_bytes_from_item(data[0])


def numbered_output_path(path: Path, index: int) -> Path:
    return path.parent / f"{path.stem}_{index}{path.suffix}"


def save_response_images(response: dict[str, Any], path: Path, n: int) -> list[Path]:
    data = response.get("data")
    if not isinstance(data, list) or not data:
        fail("response did not contain data")

    written: list[Path] = []
    byte_counts: list[int] = []
    use_suffix = n != 1 or len(data) != 1

    path.parent.mkdir(parents=True, exist_ok=True)
    for index, item in enumerate(data):
        image_bytes = extract_image_bytes_from_item(item)
        target = numbered_output_path(path, index) if use_suffix else path
        target.write_bytes(image_bytes)
        written.append(target)
        byte_counts.append(len(image_bytes))

    print(
        json.dumps(
            {
                "outputs": [str(item) for item in written],
                "bytes": byte_counts,
                "data_items": len(data),
            },
            ensure_ascii=False,
        )
    )
    return written


def save_response_image(response: dict[str, Any], path: Path) -> None:
    save_response_images(response, path, 1)


def command_generate(args: argparse.Namespace) -> None:
    key = resolve_api_key(args.api_key)
    payload = compact(
        {
            "model": MODEL,
            "prompt": args.prompt,
            "size": args.size,
            "quality": args.quality,
            "n": args.n,
            "background": args.background,
            "moderation": args.moderation,
            "output_format": args.output_format,
            "output_compression": args.output_compression,
            "user": args.user,
            "response_format": "b64_json",
        }
    )
    response = request_json("/images/generations", payload, key)
    default_name = f"generated-image.{args.output_format or 'png'}"
    save_response_images(response, output_path(args.output, default_name), args.n)


def command_edit(args: argparse.Namespace) -> None:
    key = resolve_api_key(args.api_key)
    fields = compact(
        {
            "model": MODEL,
            "prompt": args.prompt,
            "size": args.size,
            "quality": args.quality,
            "n": args.n,
            "background": args.background,
            "output_format": args.output_format,
            "output_compression": args.output_compression,
            "user": args.user,
            "response_format": "b64_json",
        }
    )
    files = [("image", Path(image)) for image in args.image]
    if args.mask:
        files.append(("mask", Path(args.mask)))
    response = request_multipart("/images/edits", fields, files, key)
    default_name = f"edited-image.{args.output_format or 'png'}"
    save_response_images(response, output_path(args.output, default_name), args.n)


def add_common_options(
    command: argparse.ArgumentParser, *, include_moderation: bool = False
) -> None:
    command.add_argument(
        "--api-key",
        help="Temporary API key. Prefer EAPIL_API_KEY for repeated use; OPENAI_API_KEY is also accepted.",
    )
    command.add_argument("--size", default="1024x1024")
    command.add_argument("--quality", choices=["auto", "low", "medium", "high"])
    command.add_argument("--n", type=int, default=1)
    command.add_argument("--background")
    if include_moderation:
        command.add_argument("--moderation")
    command.add_argument("--output-format", choices=["png", "jpeg", "webp"])
    command.add_argument("--output-compression", type=int)
    command.add_argument("--user")


def parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate or edit images with the EAPIL gpt-image-2 API."
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    generate = subparsers.add_parser("generate", help="generate a new image")
    generate.add_argument("--prompt", required=True)
    generate.add_argument("--output")
    add_common_options(generate, include_moderation=True)
    generate.set_defaults(func=command_generate)

    edit = subparsers.add_parser("edit", help="edit an existing image")
    edit.add_argument("--prompt", required=True)
    edit.add_argument("--image", action="append", required=True)
    edit.add_argument("--mask")
    edit.add_argument("--output")
    add_common_options(edit)
    edit.set_defaults(func=command_edit)

    return parser


def main() -> None:
    started = time.time()
    args = parser().parse_args()
    args.func(args)
    print(json.dumps({"elapsed_seconds": round(time.time() - started, 2)}))


if __name__ == "__main__":
    main()
