#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Generate four separate scene-03 coin-on-pouch play frames via MXAI.
No base64 payloads are written to disk; only clean prompts, metadata, and downloaded outputs are saved.
"""
import base64
import datetime as dt
import json
import mimetypes
import os
import pathlib
import re
import sys
import time
import urllib.request
import urllib.error

BASE = "https://mcp.mxai.cn"
PROJECT = pathlib.Path("/Users/bot1/Volumes/root_for_ai/AI工作区/通用_产品宣传视频_古钱币杜邦纸钱袋包_20260530_1702")
RUN_ID = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
MODEL = os.environ.get("MXAI_IMAGE_MODEL", "gpt-image-2")
OUTDIR = PROJECT / "outputs" / "scene-03-soft-trampoline" / f"storyboard-v11-simple-play-single-frames-{MODEL}-{RUN_ID}"
PROMPT_DIR = PROJECT / "prompts"
REPORT_DIR = PROJECT / "reports"
OUTDIR.mkdir(parents=True, exist_ok=True)
PROMPT_DIR.mkdir(parents=True, exist_ok=True)
REPORT_DIR.mkdir(parents=True, exist_ok=True)

PRODUCT_REF = PROJECT / "outputs/scene-03-soft-trampoline/storyboard-v9-single-frames-gold-particle-transition-20260606_192204/input_refs/image1_product_lock.jpg"
CHAR_DIR = PROJECT / "outputs/coin-character-clean-set/v2_cartoon_simplified_20260605"

FRAMES = [
    {
        "name": "scene03_v11_frame_03_wobbly_walk",
        "title": "3｜小心走两步，身体轻轻晃一下",
        "characters": [
            ("04_green_disc_cartoon_clean_character_v2_1.png", "绿色圆形方孔钱小人：做主角，往前小心走两步"),
            ("02_ancient_script_oval_cartoon_clean_character_v2_1.png", "黑色椭圆铭文钱小人：在侧后方伸手扶一下"),
        ],
        "prompt": """Create one single cinematic vertical 3:4 storyboard frame, not a grid, not a contact sheet, no split panels.\n\nUse the product reference as the exact yellow pouch product reference. Use the character reference cards to keep the two coin characters recognizable. Do not redesign the product or the coin characters.\n\nFrame 3: a simple close view on the yellow pouch body surface. The green round square-hole coin character carefully takes two small steps across the pouch surface; its body tilts slightly because the surface is softly uneven. The black oval inscription coin character stands just behind and to the side, gently reaching out as if to steady it. The pouch surface shows only a very slight soft depression and contact shadows under the tiny feet. Keep the action small, cute, and physically plausible. Show mostly the plain yellow pouch body surface, with at most one subtle fold or stitched edge near the background.\n\nAvoid: big bouncing, high jumping, flying, falling hard, elastic trampoline deformation, complex bag hardware, motif lighting, pattern activation, gold particles, smoke, explosion, suction, text, labels, watermark, UI.""",
    },
    {
        "name": "scene03_v11_frame_04_touch_surface",
        "title": "4｜趴下来摸包面",
        "characters": [
            ("02_ancient_script_oval_cartoon_clean_character_v2_1.png", "黑色椭圆铭文钱小人：半跪/趴下触摸包身"),
            ("01_broken_disc_cartoon_clean_character_v2_1.png", "裂纹金色圆片钱小人：旁边蹲着观察"),
        ],
        "prompt": """Create one single cinematic vertical 3:4 storyboard frame, not a grid, not a contact sheet, no split panels.\n\nUse the product reference as the exact yellow pouch product reference. Use the character reference cards to keep the two coin characters recognizable. Do not redesign the product or the coin characters.\n\nFrame 4: a quiet close-up on the yellow pouch body surface. The black oval inscription coin character kneels low and leans forward, using one small hand to gently touch and feel the pouch surface. Beside it, the cracked golden round coin character crouches and looks down curiously. Their feet and hands have clear contact shadows; the pouch surface only has a tiny natural indentation where touched. The composition is simple and calm, focused on exploration and tactile curiosity. Show mostly the plain yellow pouch body surface, with only a soft fold or edge cue in the background.\n\nAvoid: digging, tearing, strong deformation, crawling like an insect, high jump, flying, complex bag hardware, motif lighting, pattern activation, gold particles, smoke, explosion, suction, text, labels, watermark, UI.""",
    },
    {
        "name": "scene03_v11_frame_05_soft_fall_sit_up",
        "title": "5｜轻轻倒下又坐起来",
        "characters": [
            ("06_bronze_script_hollow_cartoon_clean_character_v2_1.png", "青铜圆形方孔钱小人：主角，刚从软包面上坐起"),
            ("05_central_script_disc_cartoon_clean_character_v2_1.png", "青绿色方孔钱小人：旁边弯腰看它"),
        ],
        "prompt": """Create one single cinematic vertical 3:4 storyboard frame, not a grid, not a contact sheet, no split panels.\n\nUse the product reference as the exact yellow pouch product reference. Use the character reference cards to keep the two coin characters recognizable. Do not redesign the product or the coin characters.\n\nFrame 5: a simple close view on the yellow pouch body surface. The bronze round square-hole coin character has gently fallen back onto the soft pouch surface and is now sitting up, supported by its small arms. The teal square-hole coin character bends slightly nearby, looking at it with concern and curiosity. The pouch surface softly cushions the seated character, with only a small shallow crease and clear contact shadow. The mood is playful and harmless, like a soft cushion moment.\n\nAvoid: painful fall, strong impact, exaggerated squash, bouncing high, flying, chaos, complex bag hardware, motif lighting, pattern activation, gold particles, smoke, explosion, suction, text, labels, watermark, UI.""",
    },
    {
        "name": "scene03_v11_frame_06_sitting_swing_feet",
        "title": "6｜两个小人并排坐着晃脚",
        "characters": [
            ("01_broken_disc_cartoon_clean_character_v2_1.png", "裂纹金色圆片钱小人：坐在包身小隆起上晃脚"),
            ("09_textured_bar_cartoon_clean_character_v2_1.png", "绿色长柄铲形钱小人：并排坐着晃脚"),
        ],
        "prompt": """Create one single cinematic vertical 3:4 storyboard frame, not a grid, not a contact sheet, no split panels.\n\nUse the product reference as the exact yellow pouch product reference. Use the character reference cards to keep the two coin characters recognizable. Do not redesign the product or the coin characters.\n\nFrame 6: a warm simple close view on a gently raised area of the yellow pouch body. The cracked golden round coin character and the green long spade-shaped coin character sit side by side on the soft pouch surface, relaxed and happy, lightly swinging their short feet. Their bodies remain upright and stable; the pouch surface has only a tiny soft dip under them. Keep the shot minimal, cute, and calm, with a shallow depth of field and mostly plain yellow pouch body surface.\n\nAvoid: large jump, sliding off, dangling from handles, complex bag hardware, motif lighting, pattern activation, gold particles, smoke, explosion, suction, text, labels, watermark, UI.""",
    },
]


def redact(s, key=""):
    s = str(s)
    s = re.sub(r"nb_[A-Za-z0-9_\-]+", "nb_********", s)
    if key:
        s = s.replace(key, "********")
    return s


def image_to_data_url(path: pathlib.Path) -> str:
    mime = mimetypes.guess_type(str(path))[0] or "image/png"
    b64 = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{mime};base64,{b64}"


def request_json(method, path, key, body=None, timeout=120):
    headers = {
        "Content-Type": "application/json; charset=utf-8",
        "Authorization": "Bearer " + key,
        "User-Agent": "HermesVideo/1.0",
    }
    data = None if body is None else json.dumps(body, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(BASE + path, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8", errors="replace"))
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code} {path}: {redact(raw, key)}")


def nested(resp):
    return resp.get("data", resp) if isinstance(resp, dict) and isinstance(resp.get("data", resp), dict) else {}


def serial_from(resp):
    obj = nested(resp)
    for k in ("serial_no", "serialNo", "serial"):
        if obj.get(k):
            return obj[k]
    if isinstance(resp, dict):
        for k in ("serial_no", "serialNo", "serial"):
            if resp.get(k):
                return resp[k]
    return None


def download(url, dest):
    req = urllib.request.Request(url, headers={"User-Agent": "HermesVideo/1.0"})
    with urllib.request.urlopen(req, timeout=180) as resp:
        dest.write_bytes(resp.read())


def main():
    key = os.environ.get("MX_AI_API_KEY")
    if not key:
        raise SystemExit("MX_AI_API_KEY missing")
    if not PRODUCT_REF.exists():
        raise SystemExit(f"missing product ref: {PRODUCT_REF}")

    clean_prompt_md = ["# Scene 03 v11 simple play single frames", "", f"Run: {RUN_ID}", "", f"Model: {MODEL}, ratio 3:4, resolution 1K, quality high", "", "Reference policy: each frame uses only product lock + the listed coin-character cards; no old/rejected storyboard images.", ""]
    metadata = {"run_id": RUN_ID, "outdir": str(OUTDIR), "model": MODEL, "aspect_ratio": "3:4", "frames": []}

    results = []
    for ix, frame in enumerate(FRAMES, 1):
        refs = [(PRODUCT_REF, "产品锁定参考：黄色杜邦纸钱袋包")]
        for fname, role in frame["characters"]:
            p = CHAR_DIR / fname
            if not p.exists():
                raise SystemExit(f"missing character ref: {p}")
            refs.append((p, role))

        clean_prompt_md += [f"## {frame['title']}", "", "### References", ""]
        for p, role in refs:
            clean_prompt_md.append(f"- `{p}` — {role}")
        clean_prompt_md += ["", "### Prompt", "", frame["prompt"], ""]

        body = {
            "model": MODEL,
            "aspect_ratio": "3:4",
            "resolution": "1K",
            "quality": "high",
            "count": 1,
            "prompt": frame["prompt"],
            "input_images": [image_to_data_url(p) for p, _ in refs],
        }
        resp = request_json("POST", "/mcp/api/generate/image", key, body, timeout=180)
        serial = serial_from(resp)
        rec = {"index": ix, "name": frame["name"], "title": frame["title"], "serial_no": serial, "references": [{"path": str(p), "role": role} for p, role in refs], "prompt": frame["prompt"], "status": "submitted"}
        results.append(rec)
        metadata["frames"].append({k: rec[k] for k in ("index", "name", "title", "serial_no", "references", "prompt")})
        print(json.dumps({"event": "submitted", "index": ix, "title": frame["title"], "serial_no": serial}, ensure_ascii=False), flush=True)
        # small spacing so the service is not hit with all requests in the exact same instant
        time.sleep(3)

    (PROMPT_DIR / f"storyboard_scene03_simple_play_v11_single_frames_{RUN_ID}.md").write_text("\n".join(clean_prompt_md), encoding="utf-8")

    pending = {r["serial_no"]: r for r in results if r.get("serial_no")}
    deadline = time.time() + 900
    while pending and time.time() < deadline:
        for serial in list(pending.keys()):
            rec = pending[serial]
            obj = nested(request_json("GET", f"/mcp/api/task/{serial}", key, timeout=90))
            status = str(obj.get("status"))
            rec["last_status"] = status
            print(json.dumps({"event": "poll", "serial_no": serial, "status": status, "title": rec["title"]}, ensure_ascii=False), flush=True)
            if status == "2":
                urls = obj.get("image_urls") or obj.get("images") or obj.get("urls") or []
                if isinstance(urls, str):
                    urls = [urls]
                rec["files"] = []
                for i, url in enumerate(urls, 1):
                    ext = ".png"
                    m = re.search(r"\.(png|jpg|jpeg|webp)(?:\?|$)", url, re.I)
                    if m:
                        ext = "." + m.group(1).lower().replace("jpeg", "jpg")
                    dest = OUTDIR / f"{rec['name']}_{i}{ext}"
                    download(url, dest)
                    rec["files"].append(str(dest))
                rec["status"] = "completed"
                pending.pop(serial, None)
            elif status in {"3", "4"}:
                rec["status"] = "failed"
                rec["fail_msg"] = obj.get("fail_msg") or obj.get("message") or obj
                pending.pop(serial, None)
        if pending:
            time.sleep(5)

    for rec in pending.values():
        rec["status"] = "timeout"

    metadata["results"] = []
    for rec in results:
        safe = {k: v for k, v in rec.items() if k not in {"prompt"}}
        metadata["results"].append(safe)
    meta_path = OUTDIR / "generation_metadata.json"
    meta_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")

    report_lines = ["# Scene 03 v11 simple play single frames QC", "", f"Run: {RUN_ID}", "", "Generated files:"]
    for rec in results:
        report_lines.append(f"- {rec['title']} — {rec.get('status')} — {rec.get('files', [])}")
    report_path = REPORT_DIR / f"scene03_simple_play_v11_single_frames_qc_{RUN_ID}.md"
    report_path.write_text("\n".join(report_lines), encoding="utf-8")

    print("FINAL_JSON_START")
    print(json.dumps({"ok": True, "outdir": str(OUTDIR), "prompt_file": str(PROMPT_DIR / f"storyboard_scene03_simple_play_v11_single_frames_{RUN_ID}.md"), "metadata": str(meta_path), "report": str(report_path), "results": metadata["results"]}, ensure_ascii=False, indent=2))
    print("FINAL_JSON_END")

if __name__ == "__main__":
    main()
