#!/usr/bin/env python3
"""
agent_short_drama.py — 苍耳 · 一键短剧 Agent
剧本 → 分镜(豆包) → 路由(豆包) → 逐镜出图(本地Z-Image)
      → 逐镜视频(LTX I2V, 本地) → 拼接(FFmpeg) → 成片

用法:
  python agent_short_drama.py <剧本> -e 1 --until render    # 只到出图
  python agent_short_drama.py <剧本> -e 1 --until video     # 到视频
  python agent_short_drama.py <剧本> -e 1 --until compose   # 到成片
  python agent_short_drama.py <分镜.json> --from render     # 从出图续跑
  python agent_short_drama.py <manifest> --from video       # 从视频续跑
"""
import os, sys, json, time, subprocess, urllib.request, argparse, glob, shutil

COMFY = "http://127.0.0.1:8188"
VIDEO_AI = os.path.dirname(os.path.abspath(__file__))
COMFY_OUT = os.path.expanduser("~/comfy/ComfyUI/output")
COMFY_IN = os.path.expanduser("~/comfy/ComfyUI/input")

LTX_NEG = ("static, frozen, no movement, flickering, jittery, choppy motion, morphing, "
           "deformed, twisted, contorted, distorted, blurry, low resolution, extra limbs, "
           "unnatural body, ugly, bad anatomy, disfigured, malformed, warped spine")

# ---------- 工具 ----------
def run_script(script, args):
    cmd = [sys.executable, os.path.join(VIDEO_AI, script)] + args
    print("  $", " ".join(cmd))
    r = subprocess.run(cmd)
    if r.returncode != 0:
        raise RuntimeError(f"{script} 退出码 {r.returncode}")

def comfy_submit(prompt):
    req = urllib.request.Request(f"{COMFY}/prompt",
        data=json.dumps({"prompt": prompt}).encode(),
        headers={"Content-Type": "application/json"})
    return json.load(urllib.request.urlopen(req, timeout=15))["prompt_id"]

def comfy_wait(pid, timeout=900):
    t0 = time.time()
    while time.time() - t0 < timeout:
        time.sleep(3)
        try:
            h = json.load(urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=8))
        except Exception:
            continue
        if pid in h and h[pid].get("outputs"):
            return h[pid]["outputs"]
    raise TimeoutError(f"ComfyUI 任务 {pid} 超时")

def wait_file(path, timeout=60):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if os.path.isfile(path) and os.path.getsize(path) > 0:
            return True
        time.sleep(2)
    return False

# ---------- Z-Image 出图 ----------
def zimage_workflow(prompt, seed, prefix, width=1024, height=1024):
    return {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
        "5": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}},
        "6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
        "7": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.0}},
        "8": {"class_type": "KSampler", "inputs": {"model": ["7", 0], "seed": seed, "steps": 8, "cfg": 1.0,
            "sampler_name": "res_multistep", "scheduler": "simple",
            "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "denoise": 1.0}},
        "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}},
        "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": prefix}},
    }

# ---------- LTX I2V 出视频 ----------
def ltx_i2v_workflow(image_name, prompt, seed, prefix, frames=49, width=768, height=512):
    return {
        "1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf"}},
        "2": {"class_type": "LTXAVTextEncoderLoader", "inputs": {"text_encoder": "gemma_3_12B_it.safetensors", "ckpt_name": "ltx-2.3-22b-embconn.safetensors", "device": "cpu"}},
        "3": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": LTX_NEG}},
        "5": {"class_type": "LTXVConditioning", "inputs": {"positive": ["3", 0], "negative": ["4", 0], "frame_rate": 24.0}},
        "6": {"class_type": "LoadImage", "inputs": {"image": image_name}},
        "7": {"class_type": "VAELoader", "inputs": {"vae_name": "LTX23_video_vae_bf16.safetensors"}},
        "8": {"class_type": "LTXVImgToVideo", "inputs": {"positive": ["5", 0], "negative": ["5", 1], "vae": ["7", 0], "image": ["6", 0], "width": width, "height": height, "length": frames, "batch_size": 1, "strength": 1.0}},
        "9": {"class_type": "RandomNoise", "inputs": {"noise_seed": seed}},
        "10": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler_ancestral_cfg_pp"}},
        "11": {"class_type": "CFGGuider", "inputs": {"model": ["1", 0], "positive": ["8", 0], "negative": ["8", 1], "cfg": 1.0}},
        "12": {"class_type": "ManualSigmas", "inputs": {"sigmas": "1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0"}},
        "13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["9", 0], "guider": ["11", 0], "sampler": ["10", 0], "sigmas": ["12", 0], "latent_image": ["8", 2]}},
        "14": {"class_type": "LTXVTiledVAEDecode", "inputs": {"vae": ["7", 0], "latents": ["13", 0], "horizontal_tiles": 2, "vertical_tiles": 2, "overlap": 6, "last_frame_fix": False, "working_device": "cpu", "working_dtype": "float16"}},
        "15": {"class_type": "CreateVideo", "inputs": {"images": ["14", 0], "fps": 24.0}},
        "16": {"class_type": "SaveVideo", "inputs": {"video": ["15", 0], "filename_prefix": prefix, "format": "mp4", "codec": "h264"}},
    }

def shot_to_prompt(shot, style=""):
    parts = []
    cam = shot.get('camera', '中景')
    scenes = shot.get('scenes', [])
    chars = shot.get('characters', [])
    desc = shot.get('description', '')
    if cam: parts.append(cam)
    if scenes: parts.append("in " + scenes[0])
    if chars: parts.append("with " + ", ".join(chars[:3]))
    if desc: parts.append(desc[:120])
    if style: parts.append(style)
    return ", ".join(parts)

def shot_to_ltx_prompt(desc, camera=""):
    """分镜 → LTX I2V 运动提示词（四段式简化：起势→动作→环境→收尾）"""
    base = (desc or "").strip()
    action = base if base else "the subject moves naturally"
    parts = [
        f"Opening on the first frame scene, {camera or 'medium shot'}. The subject {action}, "
        "with gentle flowing motion, grass and light drifting softly around. "
        "Camera holds steady, shallow depth of field, warm natural light. "
        "Slow, fluid, seamless motion, high detail. "
        "Ends with the subject settling into a calm final pose."
    ]
    return " ".join(parts)

# ---------- 各阶段 ----------
def stage_storyboard(script, episode, pro=False):
    args = [script, "-e", str(episode)]
    if pro: args.append("--pro")
    run_script("tools/run_storyboard.py", args)
    d = os.path.dirname(os.path.abspath(script))
    m = glob.glob(os.path.join(d, "STORYBOARD*.json"))
    if not m:
        raise RuntimeError("未找到分镜 JSON")
    return m[0]

def stage_route(sb_json):
    run_script("tools/route_shots.py", [sb_json])
    return sb_json

def stage_render(sb_json, out_dir=None, style="", max_shots=0):
    sb = json.load(open(sb_json, encoding="utf-8"))
    shots = sb.get("shots", [])
    if max_shots: shots = shots[:max_shots]
    if not shots: raise RuntimeError("分镜无 shots")
    if not out_dir:
        out_dir = os.path.join(os.path.dirname(os.path.abspath(sb_json)), "renders")
    os.makedirs(out_dir, exist_ok=True)
    manifest = os.path.join(out_dir, "shot_manifest.json")
    entries = []
    for i, sh in enumerate(shots):
        p = shot_to_prompt(sh, style)
        seed = 1000 + i * 17
        prefix = f"eed_shot_{i+1:04d}"
        pid = comfy_submit(zimage_workflow(p, seed, prefix))
        print(f"  🎬 [{i+1}/{len(shots)}] {sh.get('shot_number', f'S{i+1:02d}')}: {p[:50]}...")
        comfy_wait(pid)
        # 等文件落盘
        img = os.path.join(COMFY_OUT, f"{prefix}_00001_.png")
        wait_file(img)
        entries.append({"shot": sh.get("shot_number", f"S{i+1:02d}"), "prompt": p, "ltx_desc": sh.get("description",""),
                        "camera": sh.get("camera",""), "seed": seed, "prefix": prefix})
    with open(manifest, "w", encoding="utf-8") as f:
        json.dump({"out_dir": out_dir, "shots": entries}, f, ensure_ascii=False, indent=2)
    print(f"  📦 出图清单: {manifest}")
    return manifest

def stage_video(manifest, frames=49, max_shots=0):
    data = json.load(open(manifest, encoding="utf-8"))
    entries = data["shots"]
    if max_shots: entries = entries[:max_shots]
    vdir = os.path.join(data["out_dir"], "videos")
    os.makedirs(vdir, exist_ok=True)
    vlist = []
    for i, e in enumerate(entries):
        src = os.path.join(COMFY_OUT, f"{e['prefix']}_00001_.png")
        if not os.path.isfile(src):
            print(f"  ⚠️ 缺起始图 {src}，跳过 {e['shot']}"); continue
        iname = f"agent_shot_{i+1:04d}.png"
        shutil.copy(src, os.path.join(COMFY_IN, iname))
        prompt = shot_to_ltx_prompt(e.get("ltx_desc", ""), e.get("camera", ""))
        seed = 5000 + i * 29
        prefix = f"eed_vid_{i+1:04d}"
        pid = comfy_submit(ltx_i2v_workflow(iname, prompt, seed, prefix, frames=frames))
        print(f"  🎥 [{i+1}/{len(entries)}] {e['shot']} LTX I2V {frames}帧 提交，预计 {frames*4}s...")
        comfy_wait(pid, timeout=900)
        vout = os.path.join(COMFY_OUT, f"{prefix}_00001_.mp4")
        if not wait_file(vout, timeout=60):
            print(f"  ⚠️ 视频未落盘 {vout}")
            continue
        dest = os.path.join(vdir, f"shot_{i+1:04d}.mp4")
        shutil.copy(vout, dest)
        vlist.append({"shot": e["shot"], "file": dest, "prompt": prompt})
        print(f"  ✅ {e['shot']} → {dest}")
    vman = os.path.join(vdir, "video_manifest.json")
    with open(vman, "w", encoding="utf-8") as f:
        json.dump({"video_dir": vdir, "videos": vlist}, f, ensure_ascii=False, indent=2)
    print(f"  📦 视频清单: {vman}")
    return vman

def stage_compose(vman_or_dir, output="", title="", credits=""):
    if os.path.isfile(vman_or_dir) and vman_or_dir.endswith("video_manifest.json"):
        d = json.load(open(vman_or_dir, encoding="utf-8"))
        folder = d["video_dir"]
    else:
        folder = vman_or_dir
    vids = sorted(glob.glob(os.path.join(folder, "shot_*.mp4")))
    if not vids:
        raise RuntimeError(f"没有可拼接的视频片段: {folder}")
    if not output:
        output = os.path.join(os.path.dirname(folder), "EP01.mp4")
    if len(vids) == 1 and not title and not credits:
        # 单段且无片头片尾：直接拷贝（LTX 视频无声，video_composer 方案B 会因缺音频流失败）
        shutil.copy(vids[0], output)
        print(f"  🎬 单段成片(直接拷贝): {output}")
        return output
    args = [folder, "-o", output]
    if title: args += ["--title", title]
    if credits: args += ["--credits", credits]
    run_script("tools/video_composer.py", args)
    print(f"  🎬 成片: {output}")
    return output

def _video_duration(path):
    try:
        r = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
                            "-of", "csv=p=0", path], capture_output=True, text=True)
        return float(r.stdout.strip() or 0)
    except Exception:
        return 0

def stage_audio_pipeline(sb_json, composed, ep_dir):
    """⑥ 声画装配：配音 → 字幕 → BGM → 混音 → 有声终成片"""
    from tools import audio_pipeline as ap
    print("  🎙 配音 (Edge-TTS)...")
    voice = ap.generate_voice(sb_json, ep_dir)
    print("  📝 字幕...")
    ap.generate_srt(sb_json, ep_dir, voice)
    dur = _video_duration(composed) or 20
    print(f"  🎵 BGM ({dur:.0f}s)...")
    ap.generate_bgm(ep_dir, int(dur) + 2)
    final = os.path.join(ep_dir, "FINAL.mp4")
    print("  🎛 混音...")
    ap.mix_final(composed, os.path.join(ep_dir, "voice"),
                 os.path.join(ep_dir, "bgm.mp3"), final)
    print(f"  🎉 有声终成片: {final}")
    return final

def main():
    ap = argparse.ArgumentParser(description="苍耳 · 一键短剧 Agent")
    ap.add_argument("input", help="剧本 / 分镜JSON(--from render,route) / 清单(--from video,compose)")
    ap.add_argument("-e", "--episode", type=int, default=1)
    ap.add_argument("--pro", action="store_true", help="分镜用Pro模型")
    ap.add_argument("--style", default="", help="全局风格")
    ap.add_argument("--until", choices=["storyboard","route","render","video","compose","audio"], default="render")
    ap.add_argument("--from", dest="from_stage", choices=["route","render","video","compose","audio"])
    ap.add_argument("--max-shots", type=int, default=0, help="只跑前 N 镜")
    ap.add_argument("--frames", type=int, default=49, help="LTX 视频帧数(默认49)")
    ap.add_argument("-o", "--output", default="", help="成片路径")
    args = ap.parse_args()

    sb_json = None
    if args.from_stage in ("video", "compose"):
        manifest = args.input
        print(f"📖 从清单续跑: {manifest}")
        if args.from_stage == "video":
            vman = stage_video(manifest, frames=args.frames, max_shots=args.max_shots)
        if args.from_stage == "compose":
            out = stage_compose(manifest, args.output)
            print("🎯 成片完成:", out); return
        if args.from_stage == "audio":
            out = stage_compose(manifest, args.output)
            sb = glob.glob(os.path.join(os.path.dirname(os.path.dirname(out)), "STORYBOARD*.json"))
            if sb:
                ep_dir = os.path.dirname(os.path.abspath(out))
                final = stage_audio_pipeline(sb[0], out, ep_dir)
                print("🎯 有声终成片:", final)
            return
        print("🎯 阶段完成。续跑: --from compose")
        return

    if args.from_stage:
        sb_json = args.input
        if not os.path.isfile(sb_json):
            print(f"❌ 找不到分镜: {sb_json}"); sys.exit(1)
        print(f"📖 从分镜续跑: {sb_json}")
    else:
        if not os.path.isfile(args.input):
            print(f"❌ 找不到剧本: {args.input}"); sys.exit(1)
        print(f"🚀 苍耳 · 一键短剧 Agent | 剧本: {args.input} | EP{args.episode:02d} | 到: {args.until}")
        print("═" * 44); print("① 剧本→分镜 (豆包, ¥0.01/集 级)"); print("═" * 44)
        sb_json = stage_storyboard(args.input, args.episode, args.pro)
        print(f"✅ 分镜: {sb_json}")

    if not sb_json: sys.exit(1)

    if args.from_stage in (None, "storyboard") and args.until not in ("storyboard",):
        print("═" * 44); print("② 逐镜路由分配 (豆包)"); print("═" * 44)
        sb_json = stage_route(sb_json)
        print("✅ 路由完成")

    if args.until in ("render", "video", "compose"):
        print("═" * 44); print("③ 逐镜出图 (本地 Z-Image, 零成本)"); print("═" * 44)
        manifest = stage_render(sb_json, style=args.style, max_shots=args.max_shots)
    else:
        print("🎯 阶段完成。"); return

    if args.until in ("video", "compose"):
        print("═" * 44); print("④ 逐镜视频 (本地 LTX I2V, 零成本)"); print("═" * 44)
        vman = stage_video(manifest, frames=args.frames, max_shots=args.max_shots)
    else:
        print("🎯 出图完成。续跑: --from video <清单> 或 --until video")
        return

    if args.until in ("compose", "audio"):
        print("═" * 44); print("⑤ 剪辑拼接 (FFmpeg)"); print("═" * 44)
        out = stage_compose(vman, args.output)
        print(f"\n🎬 成片: {out}")
    if args.until == "audio":
        print("═" * 44); print("⑥ 声画装配 (配音+字幕+BGM+混音)"); print("═" * 44)
        sb = sb_json or (glob.glob(os.path.join(os.path.dirname(os.path.dirname(out)), "STORYBOARD*.json")) or [None])[0]
        if sb:
            ep_dir = os.path.dirname(os.path.abspath(out))
            final = stage_audio_pipeline(sb, out, ep_dir)
            print(f"\n🎉🎉 有声终成片: {final}")
        else:
            print("⚠️ 未找到分镜JSON，跳过声画装配")
    elif args.until != "compose":
        print("🎯 视频完成。续跑: --from compose <清单> 或 --until compose/audio")

if __name__ == "__main__":
    main()