#!/usr/bin/env python3
"""gpt-image-2.5 文生图 / 图生图工具。

支持：
- AI / Agent CLI 调用（参数或 JSON）
- 流程回调（HTTP webhook / 状态文件 / NDJSON 事件流）
- 指定本地生成保存路径

接口文档: Doc/gpt-image_doc.md
"""

from __future__ import annotations

import argparse
import base64
import json
import mimetypes
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Union

DEFAULT_MODEL = "gpt-image-2.5"
VIP_MODEL = "gpt-image-2-vip"
DEFAULT_API_KEY = "sk-b69877bd82c6412ca4872351f428d580"
DEFAULT_BASE_URL = "https://grsai.dakka.com.cn"  # 国内节点；全球: https://grsaiapi.com
DEFAULT_ASPECT_RATIO = "1024x1024"
DEFAULT_QUALITY = "auto"
DEFAULT_REPLY_TYPE = "json"
POLL_INTERVAL_SEC = 2.0
POLL_TIMEOUT_SEC = 300.0
# 标准模型仅支持约 1K；任一边超过该阈值需改用 VIP 像素尺寸
MAX_STANDARD_EDGE = 1024
VIP_MODELS = frozenset(
    {
        "gpt-image-2-vip",
        "gpt-image-2.5-flare",
        "gpt-image-2.5-sunburst",
    }
)


class GenerateError(RuntimeError):
    """生成失败。"""


def _now_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _resolve_api_key(explicit: Optional[str] = None) -> str:
    return (
        explicit
        or os.environ.get("GRSAI_API_KEY")
        or os.environ.get("GPT_IMAGE_API_KEY")
        or DEFAULT_API_KEY
    )


def _normalize_base_url(base_url: str) -> str:
    return base_url.rstrip("/")


def _guess_mime(path: Path) -> str:
    mime, _ = mimetypes.guess_type(str(path))
    return mime or "application/octet-stream"


def _file_to_data_url(path: Path) -> str:
    data = path.read_bytes()
    b64 = base64.b64encode(data).decode("ascii")
    return f"data:{_guess_mime(path)};base64,{b64}"


def parse_pixel_size(aspect_ratio: str) -> Optional[tuple[int, int]]:
    """解析 '2048x2048' / '2048X1152'；比例字符串（如 16:9）返回 None。"""
    text = (aspect_ratio or "").strip().lower().replace(" ", "")
    if "x" not in text:
        return None
    left, _, right = text.partition("x")
    if not left.isdigit() or not right.isdigit():
        return None
    return int(left), int(right)


def needs_vip_for_size(aspect_ratio: str) -> bool:
    """任一边 >1024 的像素尺寸需要 VIP 模型。"""
    size = parse_pixel_size(aspect_ratio)
    if size is None:
        return False
    w, h = size
    return w > MAX_STANDARD_EDGE or h > MAX_STANDARD_EDGE


def resolve_model_for_size(model: str, aspect_ratio: str) -> str:
    """大于 1K 时：已是 VIP 系列则保留；否则自动升级到 gpt-image-2-vip。"""
    if not needs_vip_for_size(aspect_ratio):
        return model
    if model in VIP_MODELS:
        return model
    return VIP_MODEL


def normalize_image_ref(ref: str) -> str:
    """将本地路径转为 data URL；URL / base64 / data URL 原样返回。"""
    text = ref.strip()
    if not text:
        raise GenerateError("空的参考图路径")
    if text.startswith(("http://", "https://", "data:")):
        return text
    # 纯 base64（无 data: 前缀）
    if len(text) > 256 and all(
        c.isalnum() or c in "+/=\n\r" for c in text[:200]
    ) and not Path(text).exists():
        return text

    path = Path(text).expanduser().resolve()
    if not path.is_file():
        raise GenerateError(f"参考图不存在: {path}")
    return _file_to_data_url(path)


def http_json(
    method: str,
    url: str,
    *,
    headers: Optional[Dict[str, str]] = None,
    body: Optional[Dict[str, Any]] = None,
    timeout: float = 300.0,
) -> Dict[str, Any]:
    data = None
    req_headers = {"Accept": "application/json", **(headers or {})}
    if body is not None:
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
        req_headers["Content-Type"] = "application/json"

    request = urllib.request.Request(url, data=data, headers=req_headers, method=method)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            if not raw.strip():
                return {}
            return json.loads(raw)
    except urllib.error.HTTPError as exc:
        err_body = exc.read().decode("utf-8", errors="replace")
        try:
            payload = json.loads(err_body) if err_body else {}
        except json.JSONDecodeError:
            payload = {"error": err_body}
        message = payload.get("error") or err_body or str(exc)
        raise GenerateError(f"HTTP {exc.code}: {message}") from exc
    except urllib.error.URLError as exc:
        raise GenerateError(f"网络错误: {exc}") from exc


def download_file(url: str, dest: Path, timeout: float = 120.0) -> Path:
    dest.parent.mkdir(parents=True, exist_ok=True)
    request = urllib.request.Request(url, headers={"User-Agent": "AIGame-gpt-image/1.0"})
    with urllib.request.urlopen(request, timeout=timeout) as resp:
        dest.write_bytes(resp.read())
    return dest


class CallbackHub:
    """向 stdout / 文件 / HTTP webhook 广播流程事件。"""

    def __init__(
        self,
        *,
        callback_url: Optional[str] = None,
        callback_file: Optional[Union[str, Path]] = None,
        quiet: bool = False,
    ) -> None:
        self.callback_url = callback_url
        self.callback_file = Path(callback_file) if callback_file else None
        self.quiet = quiet
        self.events: List[Dict[str, Any]] = []

    def emit(self, event: str, **payload: Any) -> Dict[str, Any]:
        record = {
            "event": event,
            "ts": _now_iso(),
            **payload,
        }
        self.events.append(record)

        if not self.quiet:
            # NDJSON，便于 AI / 流程引擎按行解析
            print(json.dumps(record, ensure_ascii=False), flush=True)

        if self.callback_file:
            self.callback_file.parent.mkdir(parents=True, exist_ok=True)
            self.callback_file.write_text(
                json.dumps(
                    {"latest": record, "history": self.events},
                    ensure_ascii=False,
                    indent=2,
                ),
                encoding="utf-8",
            )

        if self.callback_url:
            try:
                http_json("POST", self.callback_url, body=record, timeout=30.0)
            except Exception as exc:  # noqa: BLE001 — 回调失败不阻断主流程
                if not self.quiet:
                    warn = {
                        "event": "callback_error",
                        "ts": _now_iso(),
                        "error": str(exc),
                        "callback_url": self.callback_url,
                    }
                    print(json.dumps(warn, ensure_ascii=False), flush=True)

        return record


def query_result(
    base_url: str,
    task_id: str,
    api_key: str,
    timeout: float = 60.0,
) -> Dict[str, Any]:
    url = f"{_normalize_base_url(base_url)}/v1/api/result?{urllib.parse.urlencode({'id': task_id})}"
    return http_json(
        "GET",
        url,
        headers={"Authorization": f"Bearer {api_key}"},
        timeout=timeout,
    )


def submit_generate(
    *,
    prompt: str,
    api_key: str,
    base_url: str = DEFAULT_BASE_URL,
    model: str = DEFAULT_MODEL,
    images: Optional[Sequence[str]] = None,
    aspect_ratio: str = DEFAULT_ASPECT_RATIO,
    quality: str = DEFAULT_QUALITY,
    reply_type: str = DEFAULT_REPLY_TYPE,
    background: Optional[str] = None,
    mask: Optional[str] = None,
) -> Dict[str, Any]:
    image_refs = [normalize_image_ref(x) for x in (images or [])]
    body: Dict[str, Any] = {
        "model": model,
        "prompt": prompt,
        "images": image_refs,
        "aspectRatio": aspect_ratio,
        "quality": quality,
        "replyType": reply_type,
    }
    if background:
        body["background"] = background
    if mask:
        body["mask"] = normalize_image_ref(mask)

    url = f"{_normalize_base_url(base_url)}/v1/api/generate"
    return http_json(
        "POST",
        url,
        headers={"Authorization": f"Bearer {api_key}"},
        body=body,
        timeout=300.0,
    )


def wait_until_done(
    *,
    base_url: str,
    api_key: str,
    initial: Dict[str, Any],
    reply_type: str,
    callbacks: CallbackHub,
    poll_interval: float = POLL_INTERVAL_SEC,
    poll_timeout: float = POLL_TIMEOUT_SEC,
) -> Dict[str, Any]:
    status = (initial.get("status") or "").lower()
    if status in {"succeeded", "failed", "violation"} or (
        reply_type == "json" and initial.get("results")
    ):
        return initial

    task_id = initial.get("id")
    if not task_id:
        raise GenerateError(f"响应缺少任务 id: {initial}")

    deadline = time.time() + poll_timeout
    last_progress = initial.get("progress")
    current = initial

    while time.time() < deadline:
        current = query_result(base_url, task_id, api_key)
        status = (current.get("status") or "").lower()
        progress = current.get("progress")
        if progress != last_progress:
            callbacks.emit(
                "progress",
                id=task_id,
                status=status,
                progress=progress,
            )
            last_progress = progress

        if status == "succeeded" or current.get("results"):
            return current
        if status in {"failed", "violation"}:
            raise GenerateError(current.get("error") or f"任务失败: {current}")

        time.sleep(poll_interval)

    raise GenerateError(f"等待超时({poll_timeout}s)，最后状态: {current}")


def resolve_output_paths(
    output: Optional[str],
    count: int,
) -> List[Path]:
    if not output:
        stamp = time.strftime("%Y%m%d_%H%M%S")
        output = str(Path("output") / f"gpt_image_{stamp}.png")

    path = Path(output).expanduser()
    if count <= 1:
        path.parent.mkdir(parents=True, exist_ok=True)
        return [path]

    # 多图：若给的是文件路径，则追加 _1/_2；若是目录则编号保存
    if path.suffix:
        stem, suffix = path.stem, path.suffix
        parent = path.parent
        parent.mkdir(parents=True, exist_ok=True)
        return [parent / f"{stem}_{i + 1}{suffix}" for i in range(count)]

    path.mkdir(parents=True, exist_ok=True)
    return [path / f"image_{i + 1}.png" for i in range(count)]


def save_results(
    result: Dict[str, Any],
    output: Optional[str],
) -> List[Dict[str, Any]]:
    results = result.get("results") or []
    if not results:
        raise GenerateError(f"成功响应但无 results: {result}")

    paths = resolve_output_paths(output, len(results))
    saved: List[Dict[str, Any]] = []
    for item, dest in zip(results, paths):
        url = item.get("url")
        if not url:
            raise GenerateError(f"结果缺少 url: {item}")
        download_file(url, dest)
        saved.append({"url": url, "path": str(dest.resolve())})
    return saved


def generate(
    *,
    prompt: str,
    output: Optional[str] = None,
    images: Optional[Sequence[str]] = None,
    model: str = DEFAULT_MODEL,
    api_key: Optional[str] = None,
    base_url: str = DEFAULT_BASE_URL,
    aspect_ratio: str = DEFAULT_ASPECT_RATIO,
    quality: str = DEFAULT_QUALITY,
    reply_type: str = DEFAULT_REPLY_TYPE,
    background: Optional[str] = None,
    mask: Optional[str] = None,
    callback_url: Optional[str] = None,
    callback_file: Optional[str] = None,
    poll_interval: float = POLL_INTERVAL_SEC,
    poll_timeout: float = POLL_TIMEOUT_SEC,
    quiet: bool = False,
    on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
) -> Dict[str, Any]:
    """执行一次文/图生图，返回结构化结果（可供流程编排直接使用）。"""
    key = _resolve_api_key(api_key)
    model = resolve_model_for_size(model, aspect_ratio)
    callbacks = CallbackHub(
        callback_url=callback_url,
        callback_file=callback_file,
        quiet=quiet,
    )
    if on_event:
        original_emit = callbacks.emit

        def emit_with_hook(event: str, **payload: Any) -> Dict[str, Any]:
            record = original_emit(event, **payload)
            on_event(record)
            return record

        callbacks.emit = emit_with_hook  # type: ignore[method-assign]

    callbacks.emit(
        "started",
        model=model,
        prompt=prompt,
        output=output,
        aspect_ratio=aspect_ratio,
        image_count=len(images or []),
        reply_type=reply_type,
    )

    try:
        submitted = submit_generate(
            prompt=prompt,
            api_key=key,
            base_url=base_url,
            model=model,
            images=images,
            aspect_ratio=aspect_ratio,
            quality=quality,
            reply_type=reply_type,
            background=background,
            mask=mask,
        )
        callbacks.emit(
            "submitted",
            id=submitted.get("id"),
            status=submitted.get("status"),
            progress=submitted.get("progress"),
        )

        # async / 未完成时轮询
        effective_reply = reply_type
        if (submitted.get("status") or "").lower() == "running" and not submitted.get(
            "results"
        ):
            effective_reply = "async"

        final = wait_until_done(
            base_url=base_url,
            api_key=key,
            initial=submitted,
            reply_type=effective_reply,
            callbacks=callbacks,
            poll_interval=poll_interval,
            poll_timeout=poll_timeout,
        )

        status = (final.get("status") or "").lower()
        if status in {"failed", "violation"}:
            raise GenerateError(final.get("error") or f"任务失败: {final}")

        saved = save_results(final, output)
        result = {
            "ok": True,
            "id": final.get("id"),
            "status": final.get("status") or "succeeded",
            "progress": final.get("progress", 100),
            "model": model,
            "prompt": prompt,
            "files": saved,
            "raw": final,
        }
        callbacks.emit("succeeded", **{k: v for k, v in result.items() if k != "raw"})
        return result
    except Exception as exc:
        err = {"ok": False, "error": str(exc)}
        callbacks.emit("failed", **err)
        raise


def _load_json_arg(raw: Optional[str]) -> Dict[str, Any]:
    if not raw:
        return {}
    if raw == "-":
        return json.load(sys.stdin)
    path = Path(raw)
    if path.is_file():
        return json.loads(path.read_text(encoding="utf-8"))
    return json.loads(raw)


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="gpt_image_generate",
        description="gpt-image-2.5 文生图/图生图（支持 AI 调用与流程回调）",
    )
    p.add_argument("--prompt", "-p", help="提示词")
    p.add_argument(
        "--image",
        "-i",
        action="append",
        dest="images",
        default=[],
        help="参考图（可多次；本地路径 / URL / base64）",
    )
    p.add_argument(
        "--output",
        "-o",
        help="生成图保存路径（文件或目录）。多图时自动编号",
    )
    p.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"模型，默认 {DEFAULT_MODEL}；像素任一边>1024 时自动升为 {VIP_MODEL}",
    )
    p.add_argument("--api-key", default=None, help="API Key，默认读环境变量或内置配置")
    p.add_argument(
        "--base-url",
        default=DEFAULT_BASE_URL,
        help="API 节点，默认国内节点",
    )
    p.add_argument(
        "--aspect-ratio",
        default=DEFAULT_ASPECT_RATIO,
        help='比例或像素，如 "16:9" / "1024x1024"；>1K 请用像素如 "2048x2048"（自动 VIP）',
    )
    p.add_argument("--quality", default=DEFAULT_QUALITY, help="质量: auto/low/medium/high/...")
    p.add_argument(
        "--reply-type",
        default=DEFAULT_REPLY_TYPE,
        choices=["json", "stream", "async"],
        help="回复类型",
    )
    p.add_argument("--background", default=None, help="背景参数，如 transparent（VIP 模型）")
    p.add_argument("--mask", default=None, help="mask 图路径或 URL")
    p.add_argument(
        "--callback-url",
        default=None,
        help="流程回调 URL：每个事件 POST JSON",
    )
    p.add_argument(
        "--callback-file",
        default=None,
        help="流程回调状态文件（持续覆盖写入 latest/history）",
    )
    p.add_argument("--poll-interval", type=float, default=POLL_INTERVAL_SEC)
    p.add_argument("--poll-timeout", type=float, default=POLL_TIMEOUT_SEC)
    p.add_argument(
        "--json",
        dest="json_input",
        default=None,
        help='JSON 参数（对象字符串 / 文件路径 / "-" 读 stdin）',
    )
    p.add_argument(
        "--result-only",
        action="store_true",
        help="只输出最终结果 JSON（抑制中间 NDJSON 事件到 stdout，仍写回调）",
    )
    p.add_argument(
        "--quiet",
        action="store_true",
        help="不向 stdout 打印事件流（最终结果仍打印）",
    )
    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    payload = _load_json_arg(args.json_input)

    prompt = payload.get("prompt") or args.prompt
    if not prompt:
        parser.error("必须提供 --prompt 或 JSON.prompt")

    images = payload.get("images") or args.images or []
    if isinstance(images, str):
        images = [images]

    output = payload.get("output") or args.output
    quiet_events = bool(args.quiet or args.result_only or payload.get("quiet"))

    try:
        result = generate(
            prompt=prompt,
            output=output,
            images=images,
            model=payload.get("model") or args.model,
            api_key=payload.get("api_key") or args.api_key,
            base_url=payload.get("base_url") or args.base_url,
            aspect_ratio=payload.get("aspectRatio")
            or payload.get("aspect_ratio")
            or args.aspect_ratio,
            quality=payload.get("quality") or args.quality,
            reply_type=payload.get("replyType")
            or payload.get("reply_type")
            or args.reply_type,
            background=payload.get("background") or args.background,
            mask=payload.get("mask") or args.mask,
            callback_url=payload.get("callback_url") or args.callback_url,
            callback_file=payload.get("callback_file") or args.callback_file,
            poll_interval=float(payload.get("poll_interval") or args.poll_interval),
            poll_timeout=float(payload.get("poll_timeout") or args.poll_timeout),
            quiet=quiet_events,
        )
    except GenerateError as exc:
        err = {"ok": False, "error": str(exc)}
        print(json.dumps(err, ensure_ascii=False), flush=True)
        return 1
    except Exception as exc:  # noqa: BLE001
        err = {"ok": False, "error": str(exc)}
        print(json.dumps(err, ensure_ascii=False), flush=True)
        return 2

    # 最终结果单独一行，AI / 流程可取最后一行或 --result-only 模式
    print(json.dumps(result, ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())