"""Smoke: 豆包 Seedream 图像生成 tool 端到端走通。

跑法: .venv/Scripts/python.exe scripts/smoke_seedream.py
依赖 .env 里 ARK_API_KEY / ZCBOT_DB_URL。**会真的调豆包 API,产生 ~¥0.22 费用**。

校验:
  1. ArkConfig.load() 拿到 cfg
  2. SeedreamTool.execute(prompt=...) 返回 [image saved: ...] 文案
  3. figures/<ts>-<rand>.png 落盘且大于 0 字节
  4. 同名 .meta.json 存在 + 含 prompt/model_id/cost_cny 字段
  5. usage_events 多出一行 kind="image",单价 snapshot 在 units jsonb
"""
from __future__ import annotations

import json
import os
import sys
import uuid
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 sqlalchemy import text

from core.ark_client import ArkConfig
from core.storage import session_scope
from core.storage.models import Task, User
from tools.seedream import SeedreamTool


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 {})
    variant_key, variant_cfg = next(iter(image_cfg.items()))
    print(f"[setup] variant={variant_key} model={variant_cfg.get('model_id')} price={variant_cfg.get('price_cny_per_image')}")

    # 准备一次性 user + task 行(usage_events FK 校验)
    uid = uuid.uuid4()
    tid = uuid.uuid4()
    ws_user = ROOT / "workspace" / "users" / str(uid)
    wd = ws_user / "smoke_seedream"
    wd.mkdir(parents=True, exist_ok=True)
    with session_scope() as s:
        s.add(User(user_id=uid))
        s.add(Task(task_id=tid, user_id=uid, name="smoke_seedream", working_dir=str(wd)))

    tool = SeedreamTool(
        ark_cfg=cfg,
        image_variant_cfg=variant_cfg,
        variant_key=variant_key,
        working_dir=wd,
        task_id=tid,
        user_id=uid,
        base_dir=wd,
        user_root=ws_user,
    )

    print(f"[call] prompt='一只橙色的小猫坐在窗台上望向远方,水彩风格'")
    result = tool.execute(prompt="一只橙色的小猫坐在窗台上望向远方,水彩风格")
    print(f"[tool result]\n{result}\n")

    if result.startswith("[Error]"):
        print(f"[FAIL] tool 返回错误")
        return 2

    # 校验 figures 目录与 meta
    figs = list((wd / "figures").glob("*.png"))
    assert len(figs) == 1, f"figures/*.png 应当 1 个,实际 {len(figs)}"
    png = figs[0]
    assert png.stat().st_size > 0, f"{png} 大小为 0"
    print(f"[OK] png 落盘 {png.name} ({png.stat().st_size} bytes)")

    meta_path = png.with_suffix(".meta.json")
    assert meta_path.exists(), f"meta 文件不存在 {meta_path}"
    meta = json.loads(meta_path.read_text(encoding="utf-8"))
    for k in ("prompt", "model_id", "size", "cost_cny", "ts"):
        assert k in meta, f"meta 缺字段 {k}"
    print(f"[OK] meta 字段齐全: {list(meta.keys())}")

    # 校验 usage_events
    with session_scope() as s:
        rows = s.execute(text(
            "SELECT kind, model_profile, units, cost_cny FROM usage_events "
            "WHERE task_id = :tid"
        ), {"tid": str(tid)}).all()
    assert len(rows) == 1, f"usage_events 行数应 1,实际 {len(rows)}"
    row = rows[0]
    assert row.kind == "image", f"kind 应 image,实际 {row.kind}"
    assert row.model_profile == f"doubao.{variant_key}", f"model_profile = {row.model_profile}"
    assert "price_cny_per_image" in row.units, "units 缺 price_cny_per_image snapshot"
    print(f"[OK] usage_events: kind={row.kind} model={row.model_profile} cost_cny={row.cost_cny}")
    print(f"     units snapshot: {row.units}")

    print("\n[PASS] smoke_seedream 全部通过")
    return 0


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