#!/usr/bin/env python3
"""关键帧生成 —— 千问AI平台（DashScope）多图融合

2026-10-04 实测确认可用（¥0.2/张）：
  · 端点  https://maas.qianwenaiapi.com/api/v1/services/aigc/multimodal-generation/generation
  · 模型  wan2.7-image-pro（最多 9 张输入图 / 顺序模式一次 12 张）
  · 格式  input.messages[0].content[] = [{text}, {image}, {image}, ...]
  · 图片  支持 base64 data URL（无需先上传 OSS）；实测 769×1362 竖屏

用法（CLI 自测，会花钱）：
    python3 scripts/keyframe_gen.py --role <角色图> --scene <场景图> \
        --prompt "把图1的人放进图2的房间里，中景，平视" --out /tmp/kf.png
    python3 scripts/keyframe_gen.py --text-only --prompt "空旷病房" --out /tmp/kf.png
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import urllib.error
import urllib.request
from pathlib import Path

SCRIPTS = Path(__file__).resolve().parent
if str(SCRIPTS) not in sys.path:
    sys.path.insert(0, str(SCRIPTS))
from lib.http import _NO_PROXY_OPENER  # noqa: E402
from lib.logger import get_logger       # noqa: E402

log = get_logger("keyframe_gen")

API = "https://maas.qianwenaiapi.com/api/v1/services/aigc/multimodal-generation/generation"
DEFAULT_MODEL = "wan2.7-image-pro"      # 9 图输入 / 顺序 12 张；备选 qwen-image-3.0-pro
UA = {"User-Agent": "Mozilla/5.0 (AgentOS-Keyframe)"}

# ── 画幅预设（官方文档实测：size 格式为 `宽*高`，不是只有 1K/2K/4K 档位）──
#   qwen-image 系列：宽高自由 512~2048
#   wan2.5+：总像素 1280² ~ 1440²（推荐 9:16 = 960*1696）
#   wan2.2 及更早：边长 512~1440 · FLUX：仅几个固定值
SIZE_PRESETS = {
    "竖屏 9:16 · 720P（测试推荐）": "720*1280",
    "竖屏 9:16 · 标准":            "960*1696",
    "竖屏 9:16 · 1080P":           "1080*1920",
    "方图 1:1":                    "1024*1024",
    "横屏 16:9 · 720P":            "1280*720",
    "横屏 16:9 · 标准":            "1696*960",
    "省算力 · 小图":               "768*1024",
}

DEFAULT_SIZE = "720*1280"               # 竖屏 720P —— 测试期默认，省算力
                                        # （正式出片可切 960*1696 或 1080*1920）


# ── 模型注册表（config/image_models.yaml）────────────────────
def _models_yaml() -> Path:
    return Path(os.environ.get("AVE_IMAGE_MODELS")
                or (SCRIPTS.parent / "config" / "image_models.yaml"))


def image_models() -> list[dict]:
    """读取模型注册表（含能力注释）。读不到时回退到内置最小集合。"""
    try:
        import yaml
        d = yaml.safe_load(_models_yaml().read_text(encoding="utf-8")) or {}
        ms = d.get("models") or []
        if ms:
            return ms
    except Exception as e:
        log.warning("读取 image_models.yaml 失败: %s", e)
    return [{"id": DEFAULT_MODEL, "name": DEFAULT_MODEL, "vendor": "百炼（千问AI平台）",
             "price": "¥0.2/张", "note": "内置回退项", "supports": ["multi", "single", "text2img"],
             "limits": {}, "key_field": "aliyun"}]


def models_for_mode(mode: str) -> list[dict]:
    """某生成方式可用的模型（按 supports 过滤）"""
    m = (mode or "").strip()
    return [x for x in image_models() if m in (x.get("supports") or [])]


def model_info(mid: str) -> dict:
    for x in image_models():
        if x.get("id") == mid:
            return x
    return {}

# ── 按「生成方式」选模型 ──────────────────────────────────
#   多图融合 / 文生图 / 图生图 → wan2.7-image-pro（9 图输入，融合强）
#   改造已有图（结构保持重绘）→ qwen-image-edit-plus（指令编辑，保结构）
MODEL_BY_MODE = {
    "multi":    "wan2.7-image-pro",
    "text2img": "wan2.7-image-pro",
    "single":   "wan2.7-image-pro",
    "edit":     "qwen-image-edit-plus",
}

# ⚠️ 官方示例里 qwen-image 系列的 content 顺序是「先 image 后 text」，
#    而 wan 系列是「先 text 后 image」。混用会导致指令被忽略。
IMAGE_FIRST_PREFIXES = ("qwen-image",)


def model_for_mode(mode: str) -> str:
    return MODEL_BY_MODE.get((mode or "").strip(), DEFAULT_MODEL)


def _api_key() -> str:
    """读取百炼/千问平台 key（agent-local 私有配置）"""
    import yaml
    p = os.environ.get("AVE_LOCAL_CONFIG") or os.path.expanduser(
        "~/workbuddy-agent-os/agent-local/tools/ave/config/local.yaml")
    try:
        return ((yaml.safe_load(open(p)) or {}).get("aliyun") or {}).get("api_key", "")
    except Exception as e:
        log.warning("读取 key 失败: %s", e)
        return ""


def _assets_dir() -> Path:
    """资产库根目录（与 Dashboard 的 AssetStore 保持一致）"""
    return Path(os.environ.get("AVE_ASSETS_DIR")
                or os.path.expanduser("~/workbuddy-agent-os/agent-local/tools/ave/assets"))


def _url_to_local(s: str) -> Path | None:
    """把「后端 URL 形式」还原成本地文件路径

    前端资产选择器给出的是可访问 URL：
        /api/ave/assets/{kind}/{aid}/file/{name}
    它与磁盘上的
        {assets}/{kind}/{aid}/{name}
    是同一个文件。这里做还原（含 URL 解码），并保证不越出资产目录。
    """
    from urllib.parse import unquote
    if not s.startswith("/api/ave/assets/"):
        return None
    parts = [x for x in s.split("/") if x]          # api ave assets kind aid file name
    if len(parts) < 7 or parts[5] != "file":
        return None
    kind, aid, name = (unquote(parts[3]), unquote(parts[4]), unquote(parts[6]))
    base = _assets_dir().resolve()
    cand = (base / kind / aid / name)
    try:
        cand = cand.resolve()
        if not str(cand).startswith(str(base)):     # 防路径穿越
            return None
    except Exception:
        return None
    return cand if cand.exists() else None


def _as_image_value(path_or_url: str) -> str:
    """参考图 → 可提交给模型的值

    支持四种输入：
      1. http(s) URL        → 原样返回
      2. data: base64       → 原样返回
      3. /api/ave/assets/…  → 还原成本地路径后转 base64（前端选择器给的格式）
      4. 本地文件路径        → 转 base64
    """
    s = str(path_or_url)
    if s.startswith(("http://", "https://", "data:")):
        return s
    p = _url_to_local(s) or Path(s)
    if not p.exists():
        raise FileNotFoundError(
            f"参考图不存在: {p}"
            + ("（未能从 /api/ave/assets/ 形式还原出本地文件）" if s.startswith("/api/") else ""))
    b64 = base64.b64encode(p.read_bytes()).decode()
    ext = p.suffix.lower().lstrip(".") or "png"
    mime = {"jpg": "jpeg", "jpeg": "jpeg", "webp": "webp"}.get(ext, "png")
    return f"data:image/{mime};base64,{b64}"


def build_keyframe(
    prompt: str,
    ref_images: list[str] | None = None,
    *,
    negative_prompt: str = "",
    size: str = DEFAULT_SIZE,
    n: int = 1,
    model: str = DEFAULT_MODEL,
    timeout: int = 180,
) -> list[str]:
    """生成关键帧，返回图片 URL 列表（1~n 张）

    prompt        构图/光线/动作瞬间的描述（视觉描述，不含"运动"—— 运动属于视频阶段）
    ref_images    参考图（本地路径或 URL），顺序即提示词里的"图1/图2…"
                  推荐：角色正面 + 角色 45° + 场景对应机位 [+ 道具]
    """
    key = _api_key()
    if not key:
        raise RuntimeError("未找到 aliyun.api_key（agent-local/tools/ave/config/local.yaml）")

    # ── 模型能力校验（依据 config/image_models.yaml）──────────
    info = model_info(model)
    if info.get("status"):
        raise RuntimeError(
            f"模型「{info.get('name', model)}」尚未接线（状态：{info['status']}）。"
            f"请先选百炼系模型，或告知我接入。")
    if info and info.get("key_field", "aliyun") != "aliyun":
        raise RuntimeError(
            f"模型「{info.get('name', model)}」用的是 {info.get('key_field')} 凭证，"
            f"当前 keyframe_gen 只实现了百炼通道。")
    mx = (info.get("limits") or {}).get("max_input_images")
    n_given = len(ref_images or [])
    if mx and n_given > int(mx):
        raise RuntimeError(
            f"模型「{info.get('name', model)}」最多接受 {mx} 张输入图，"
            f"当前给了 {n_given} 张 —— 请减少参考图，或换模型。")
    n_max = (info.get("limits") or {}).get("n_max")
    if n_max and int(n) > int(n_max):
        raise RuntimeError(f"模型「{info.get('name', model)}」一次最多出 {n_max} 张（当前 {n}）。")

    full_text = prompt + (f"\n\n【避免】{negative_prompt}" if negative_prompt else "")
    imgs = [{"image": _as_image_value(i)} for i in (ref_images or [])]
    # 顺序按模型家族区分（见上方 IMAGE_FIRST_PREFIXES 的注释）
    if str(model).startswith(IMAGE_FIRST_PREFIXES):
        content: list[dict] = imgs + [{"text": full_text}]
    else:
        content = [{"text": full_text}] + imgs

    payload = {
        "model": model,
        "input": {"messages": [{"role": "user", "content": content}]},
        "parameters": {"size": size, "n": int(n), "watermark": False},
    }
    log.info("🖼 关键帧生成：model=%s refs=%d size=%s n=%d",
             model, len(ref_images or []), size, n)
    req = urllib.request.Request(API, data=json.dumps(payload).encode(),
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        method="POST")
    with _NO_PROXY_OPENER.open(req, timeout=timeout) as r:
        d = json.loads(r.read().decode())
    out: list[str] = []
    for ch in ((d.get("output") or {}).get("choices") or []):
        for c in ((ch.get("message") or {}).get("content") or []):
            if c.get("image"):
                out.append(c["image"])
    log.info("✅ 完成，产出 %d 张（usage=%s）", len(out),
             json.dumps(d.get("usage") or {}, ensure_ascii=False))
    return out


def download(url: str, dest: str | Path, timeout: int = 90) -> str:
    dest = Path(dest)
    dest.parent.mkdir(parents=True, exist_ok=True)
    with _NO_PROXY_OPENER.open(urllib.request.Request(url, headers=UA), timeout=timeout) as r:
        dest.write_bytes(r.read())
    log.info("已保存: %s (%d KB)", dest, dest.stat().st_size // 1024)
    return str(dest)


def main() -> int:
    ap = argparse.ArgumentParser(description="关键帧生成（多图融合）")
    ap.add_argument("--prompt", required=True, help="画面描述（视觉，不含运动）")
    ap.add_argument("--role", action="append", default=[], help="角色参考图（可多次）")
    ap.add_argument("--scene", action="append", default=[], help="场景参考图（可多次）")
    ap.add_argument("--extra", action="append", default=[], help="其他参考图（道具等）")
    ap.add_argument("--negative", default="", help="负向提示词")
    ap.add_argument("--size", default=DEFAULT_SIZE,
                    help="画幅 宽*高（如 960*1696 竖屏 9:16 / 1280*720 横屏）")
    ap.add_argument("--n", type=int, default=1)
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--out", default="/tmp/keyframe.png")
    ap.add_argument("--text-only", action="store_true", help="不用参考图（纯文生图）")
    args = ap.parse_args()

    refs = [] if args.text_only else (args.role + args.scene + args.extra)
    urls = build_keyframe(args.prompt, refs, negative_prompt=args.negative,
                          size=args.size, n=args.n, model=args.model)
    if not urls:
        print("❌ 未返回图片")
        return 1
    for i, u in enumerate(urls):
        dest = args.out if len(urls) == 1 else args.out.replace(".png", f"_v{i+1}.png")
        download(u, dest)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())