"""Probe: seedream 5.0 i2i 是否接受 base64 data: URI 作 image_urls 项。

跑法: .venv/Scripts/python.exe scripts/probe_seedream_i2i.py
**base64 接受 -> 真出图,产生 ~0.22 CNY 费用**;base64 被拒 -> ArkError,不计费。

目的:确认 i2i 走 base64 通路是否可行(zcbot 内网部署,无法提供公网 URL 给 ARK 反向 fetch)。
如果 base64 通 → 后续 seedream i2i tool 可以直接用 base64,工程简单;
如果 base64 拒 → 要查 ARK 是否有"上传文件拿 file_id"接口或外接对象存储。

输出:
  - HTTP 状态 / 响应 JSON(节选)
  - 结论标签 [RESULT] base64-supported / base64-rejected / unclear
"""
from __future__ import annotations

import base64
import io
import json
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

# 读 .env
env_file = ROOT / ".env"
if env_file.exists():
    for line in env_file.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, _, v = line.partition("=")
        os.environ.setdefault(k.strip(), v.strip())

from PIL import Image

from core.ark_client import ArkClient, ArkConfig, ArkError


def make_test_png() -> bytes:
    """256x256 蓝底 PNG,中间一个红圆 —— prompt 改色易验证 i2i 真生效。"""
    img = Image.new("RGB", (256, 256), (50, 100, 200))
    # 中心红圆
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    draw.ellipse((80, 80, 176, 176), fill=(220, 60, 60))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def main() -> int:
    cfg = ArkConfig.load()
    if cfg is None:
        print("[SKIP] ARK_API_KEY 未设 / config/media/doubao.yaml 缺,无法测真接口")
        return 0

    image_cfg = (cfg.raw.get("image") or {})
    if not image_cfg:
        print("[SKIP] doubao.yaml 无 image 段")
        return 0
    variant_key, variant_cfg = next(iter(image_cfg.items()))
    model_id = variant_cfg["model_id"]
    print(f"[setup] model={model_id}  endpoint=/images/generations")

    png_bytes = make_test_png()
    b64 = base64.b64encode(png_bytes).decode("ascii")
    data_uri = f"data:image/png;base64,{b64}"
    print(f"[setup] reference image: {len(png_bytes)} bytes -> data URI len={len(data_uri)}")

    body = {
        "model": model_id,
        "prompt": "Keep the layout but change the red circle to a yellow star, blue background unchanged.",
        "image_urls": [data_uri],
        "size": "2048x2048",  # ARK 最小要求 3686400 像素 (~1920²),1024² 会被拒
        "response_format": "url",
        "watermark": False,
    }

    print("[probe] POST /images/generations with image_urls=[data:image/png;base64,...]")
    try:
        with ArkClient(cfg, timeout_s=120.0) as client:
            resp = client.post_json("/images/generations", body, timeout_s=120.0)
    except ArkError as e:
        msg = str(e)
        print(f"[ArkError] {msg}")
        low = msg.lower()
        if "base64" in low or "image_url" in low or "invalid_url" in low or "url" in low:
            print("\n[RESULT] base64-rejected (error mentions image_url/base64/url)")
            print("         -> i2i 需要走公网 URL(对象存储)或 ARK file 上传 endpoint(待查)")
        else:
            print("\n[RESULT] unclear (ArkError but not obviously a base64 format issue)")
        return 2

    print(f"[response keys] {list(resp.keys())}")
    snippet = json.dumps(resp, ensure_ascii=False)[:600]
    print(f"[response (first 600ch)]\n{snippet}")

    # 解析出图 URL
    data = resp.get("data")
    new_url = None
    if isinstance(data, list) and data:
        first = data[0]
        if isinstance(first, dict):
            new_url = first.get("url") or first.get("image_url")
    elif isinstance(data, dict):
        imgs = data.get("images")
        if isinstance(imgs, list) and imgs and isinstance(imgs[0], dict):
            new_url = imgs[0].get("url")

    if new_url:
        print(f"\n[RESULT] base64-supported (got new image url: {new_url[:90]}...)")
        print(f"         estimated cost ~{variant_cfg.get('price_cny_per_image', 0.22)} CNY")
        return 0
    print("\n[RESULT] unclear (HTTP 200 but no image url in response — check JSON above)")
    return 1


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