"""gpt_image: 调 unifyllm 网关的 OpenAI Images API 生图 / 改图。

第二图像后端(第一个是豆包 seedream):模型 ID + 单价全在 `config/media/unifyllm.yaml`,
本 tool 只装配:
- 文生图走 /images/generations JSON,gpt-image-2 支持自定义尺寸和质量档;
- 单图改图走 /images/edits multipart/form-data;
- 响应直接返 b64_json,无需二次下载;
- 复杂图片可能需 ~2min。
完成后:
- 图片落 `<working_dir>/figures/<YYYYMMDD-HHMMSS>-<rand6>.png` + 同名 `.meta.json`
- usage_events 写 kind="image" 一行(model_profile="unifyllm.<variant>",
  usage tokens 记进 units → 网关价目公布后可回填对账)
"""
from __future__ import annotations

import base64
import json
import re
import struct
import time
from datetime import datetime
from pathlib import Path
from typing import Optional
from uuid import UUID

from core.ark_client import ArkClient, ArkConfig, ArkError
from core.artifacts import ArtifactRef, ToolExecutionResult, resolve_artifact_path
from core.storage.usage import record_image_usage

from .base import Tool
from .image_ref import load_image_as_data_url
from .media_common import quota_gate, record_usage_safe, stamped_path, write_meta


class GptImageTool(Tool):
    name = "gpt_image"
    description = (
        "Generate an image via GPT Image, saved to working_dir/figures/. Supports custom size "
        "and quality, plus single-reference image editing. For editing, pass reference_images "
        "with the existing image path; the original is preserved and the result is saved as a "
        "new image. Complex images may take up to ~2 minutes. Don't generate decoratively — "
        "only when the user actually wants an image. Returns the saved relative path."
    )
    parameters = {
        "type": "object",
        "properties": {
            "prompt": {
                "type": "string",
                "description": "中文或英文都行,详尽描述画面(主体/风格/光线/构图)。",
            },
            "reference_images": {
                "type": "array",
                "items": {"type": "string"},
                "description": (
                    "改图(image-to-image):传 1 张 task_dir 内已存在图片的相对路径。"
                    "不传 = 从零文生图;当前仅支持单张参考图。"
                ),
            },
            "size": {
                "type": "string",
                "description": (
                    "输出尺寸。默认 auto;也可传 WIDTHxHEIGHT,如 1024x1024、1536x864、"
                    "2048x1152。宽高须为 16 的倍数,比例不超过 3:1,总像素 "
                    "655360-8294400,单边不超过 3840。"
                ),
            },
            "quality": {
                "type": "string",
                "enum": ["auto", "low", "medium", "high"],
                "description": "生成质量。默认 auto;草稿用 low,正式产物用 medium 或 high。",
            },
        },
        "required": ["prompt"],
    }

    def __init__(
        self,
        *,
        gw_cfg: ArkConfig,
        image_variant_cfg: dict,
        variant_key: str,
        working_dir: Path,
        task_id: UUID,
        user_id: UUID,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
        daily_limit: int = 0,
    ) -> None:
        super().__init__(base_dir, user_root=user_root)
        self.gw_cfg = gw_cfg
        self.cfg = image_variant_cfg
        self.variant_key = variant_key  # 'gpt_image' → usage_events.model_profile = "unifyllm.gpt_image"
        self.working_dir = Path(working_dir)
        self.task_id = task_id
        self.user_id = user_id
        self.daily_limit = int(daily_limit)  # 0 / 负 = 不限;与 seedream 共享 kind="image" 每日配额

    def execute(
        self,
        prompt: str,
        reference_images: list | None = None,
        size: Optional[str] = None,
        quality: Optional[str] = None,
    ) -> str:
        if not (prompt or "").strip():
            return "[Error] prompt 不能为空"

        refs = [str(r).strip() for r in (reference_images or []) if str(r).strip()]
        if len(refs) > 1:
            return (
                "[Error] reference_images 当前仅支持单张参考图(传了 "
                f"{len(refs)} 张)。请只传 1 张。"
            )
        ref_bytes = b""
        ref_mime = ""
        ref_disp = ""
        if refs:
            data_url, ref_disp, ref_err = load_image_as_data_url(
                refs[0],
                working_dir=self.working_dir,
                user_root=self.user_root,
                display_fn=self._display,
            )
            if ref_err:
                return ref_err
            header, encoded = data_url.split(",", 1)
            ref_mime = header.removeprefix("data:").removesuffix(";base64")
            ref_bytes = base64.b64decode(encoded)
        is_i2i = bool(ref_bytes)

        cfg = self.cfg
        chosen_size, size_err = self._normalize_size(
            size if size is not None else cfg.get("default_size", "auto")
        )
        if size_err:
            return size_err
        chosen_quality = str(
            quality if quality is not None else cfg.get("default_quality", "auto")
        ).strip().lower()
        if chosen_quality not in {"auto", "low", "medium", "high"}:
            return "[Error] quality 必须是 auto / low / medium / high"

        # 每账号每日配额(kind="image" 与 seedream 同口径合计)
        quota_err = quota_gate(
            self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="张",
        )
        if quota_err:
            return quota_err

        model_id = cfg["model_id"]
        endpoint = (
            cfg.get("edit_endpoint", "/images/edits")
            if is_i2i
            else cfg.get("endpoint", "/images/generations")
        )
        timeout_s = float(cfg.get("request_timeout_s", 300))
        price = float(cfg.get("price_cny_per_image", 0))

        body = {
            "model": model_id,
            "prompt": prompt,
            "n": 1,
            "size": chosen_size,
            "quality": chosen_quality,
        }

        t0 = time.monotonic()
        try:
            with ArkClient(self.gw_cfg, timeout_s=timeout_s) as client:
                if is_i2i:
                    form = {
                        key: str(value).lower() if isinstance(value, bool) else str(value)
                        for key, value in body.items()
                    }
                    form["response_format"] = "b64_json"
                    resp = client.post_multipart(
                        endpoint,
                        form,
                        {"image": (Path(refs[0]).name or "reference.png", ref_bytes, ref_mime)},
                        timeout_s=timeout_s,
                    )
                else:
                    resp = client.post_json(endpoint, body, timeout_s=timeout_s)
        except ArkError as e:
            return f"[Error] gpt_image API: {e}"

        data = resp.get("data")
        b64 = ""
        if isinstance(data, list) and data and isinstance(data[0], dict):
            b64 = data[0].get("b64_json") or ""
            if not b64:
                image_url = str(data[0].get("url") or "")
                if image_url.startswith("data:") and ";base64," in image_url:
                    b64 = image_url.split(",", 1)[1]
        if not b64:
            return f"[Error] gpt_image response 缺 b64_json: {json.dumps(resp, ensure_ascii=False)[:300]}"
        try:
            img_bytes = base64.b64decode(b64)
        except Exception as e:
            return f"[Error] gpt_image b64 解码失败: {e}"

        dest_png = stamped_path(self.working_dir / "figures", ".png")
        dest_png.write_bytes(img_bytes)

        elapsed = time.monotonic() - t0
        # 网关回填的实际尺寸/质量/tokens(价目未知期成本记 price snapshot,tokens 留对账)
        actual_size = self._png_size(img_bytes) or str(resp.get("size") or chosen_size)
        actual_quality = str(resp.get("quality") or chosen_quality)
        usage = resp.get("usage") or {}
        output_tokens = int(usage.get("output_tokens") or 0)

        meta = {
            "prompt": prompt,
            "model_id": model_id,
            "size": actual_size,
            "requested_size": chosen_size,
            "quality": actual_quality,
            "requested_quality": chosen_quality,
            "mode": "i2i" if is_i2i else "t2i",
            "reference_images": [ref_disp] if is_i2i else [],
            "cost_cny": price,
            "output_tokens": output_tokens,
            "elapsed_s": round(elapsed, 2),
            "ts": datetime.now().isoformat(timespec="seconds"),
        }
        write_meta(dest_png, meta)

        record_usage_safe(
            "gpt_image", record_image_usage,
            task_id=self.task_id,
            user_id=self.user_id,
            model_profile=f"unifyllm.{self.variant_key}",
            n_images=1,
            size=actual_size,
            price_cny_per_image=price,
            extra_units={"output_tokens": output_tokens, "quality": actual_quality},
        )

        disp = self._display(dest_png)
        # 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析);
        # 价格未知(price=0)时不放 cost 段,避免"¥0.00 = 免费"的误导。
        cost_seg = f" · cost=¥{price:.2f}" if price > 0 else ""
        mode_seg = " · mode=i2i" if is_i2i else ""
        ref_line = f"\nreference={ref_disp}" if is_i2i else ""
        result = (
            f"[gpt_image] model={model_id} · size={actual_size} · quality={actual_quality}"
            f"{cost_seg} · elapsed={elapsed:.1f}s{mode_seg}\n"
            f"saved: {disp}{ref_line}\n"
            f"prompt={prompt!r}"
        )
        if self.user_root is None:
            return result
        _, rel = resolve_artifact_path(
            str(dest_png), working_dir=self.working_dir, user_root=self.user_root,
        )
        return ToolExecutionResult(content=result, artifacts=(ArtifactRef(path=rel),))

    @staticmethod
    def _normalize_size(raw: object) -> tuple[str, str]:
        value = str(raw or "auto").strip().lower().replace("×", "x")
        if value == "auto":
            return "auto", ""
        match = re.fullmatch(r"([1-9]\d*)x([1-9]\d*)", value)
        if not match:
            return "", "[Error] size 必须是 auto 或 WIDTHxHEIGHT,例如 1536x864"
        width, height = (int(match.group(1)), int(match.group(2)))
        pixels = width * height
        if width % 16 or height % 16:
            return "", "[Error] size 的宽和高都必须是 16 的倍数"
        if max(width, height) > 3840:
            return "", "[Error] size 单边不能超过 3840"
        if max(width, height) > min(width, height) * 3:
            return "", "[Error] size 宽高比不能超过 3:1"
        if not 655_360 <= pixels <= 8_294_400:
            return "", "[Error] size 总像素必须在 655360 到 8294400 之间"
        return f"{width}x{height}", ""

    @staticmethod
    def _png_size(data: bytes) -> str:
        if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n" and data[12:16] == b"IHDR":
            width, height = struct.unpack(">II", data[16:24])
            return f"{width}x{height}"
        return ""