#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""单段生成（每段一个独立 ComfyUI 任务）—— 规避 32GB 内存下多段单任务的内存窒息。

用法:
  python3 h3_seg.py --seg 1
  python3 h3_seg.py --seg 2 --tail <上一段最后一帧.jpg>
  python3 h3_seg.py --seg 3 --tail <上一段最后一帧.jpg>
"""
import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

import requests

sys.path.insert(0, "/home/zyw/Downloads/dasiwa-h3-research")
import v18_pipeline as p  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
from h3_run import (BG_IMG, CHAR_IMG, COMFY_URL, FPS, W, H, SEED,  # noqa: E402
                    build_api, build_timeline, upload)

WORK = Path(__file__).parent
TASK = Path("/home/zyw/Downloads/dl-hub/10-视频生成流水线/20260916_2340_仙宫玉廊_汉服女子背影行走")
SEG_FRAMES = 260


def last_frame(video: Path, out: Path) -> Path:
    """抽最后一帧。"""
    subprocess.run(["ffmpeg", "-y", "-v", "error", "-sseof", "-0.2", "-i", str(video),
                    "-vsync", "0", "-q:v", "2", "-frames:v", "1", str(out)], check=True)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seg", type=int, required=True)
    ap.add_argument("--tail", default=None)
    ap.add_argument("--variant", default="b", choices=["b", "c"],
                    help="b: 尾帧作 <Picture 3>；c: 尾帧作 <Picture 1>（首帧锚定，接缝更稳）")
    ap.add_argument("--frames", type=int, default=SEG_FRAMES)
    ap.add_argument("--seed", type=int, default=SEED)
    ap.add_argument("--tag", default="", help="文件名/输出前缀标记，如 v2")
    args = ap.parse_args()

    seg = args.seg
    suffix = args.variant if args.tail else ""
    prompt_file = WORK / f"prompt_seg{seg}{suffix}.txt"
    prompt = prompt_file.read_text(encoding="utf-8")
    print(f"[seg{seg}] 提示词: {prompt_file.name}")

    bg_remote = upload(BG_IMG, "xianjia_bg_corridor.jpg")
    ch_remote = upload(CHAR_IMG, "xianjia_char_turnaround.jpg")
    if args.tail:
        tp = Path(args.tail)
        tail_remote = upload(tp, f"xianjia_seg{seg-1}_tail.jpg")
        # 上传顺序 = <Picture 1..N>：variant c 把尾帧放第一位当首帧锚
        refs = ([tail_remote, ch_remote, bg_remote] if args.variant == "c"
                else [bg_remote, ch_remote, tail_remote])
    else:
        refs = [bg_remote, ch_remote]
    print(f"[seg{seg}] 参考图: {refs}")

    tl = build_timeline([prompt], refs, args.frames, W, H, FPS,
                        continuity=False, overlap=0)
    tag = f"_{args.tag}" if args.tag else ""
    prefix = f"xianjia_seg{seg}{tag}_{time.strftime('%m%d_%H%M')}"
    api = build_api(tl, args.frames, args.seed, prefix, refine=True)
    (TASK / f"api_seg{seg}{'_' + args.tag if args.tag else ''}.json").write_text(json.dumps(api, ensure_ascii=False, indent=1),
                                             encoding="utf-8")

    print(f"[seg{seg}] 提交：{W}x{H} 一采 4 步 → 二采 1280x704｜{args.frames} 帧 "
          f"= {args.frames/FPS:.2f}s｜seed={args.seed}")
    t0 = time.time()
    entry = p.submit_and_wait(api, timeout_min=150)
    print(f"[seg{seg}] 完成，用时 {(time.time()-t0)/60:.1f} 分钟")
    info = p.extract_video_output(entry)
    local = p.download_video(info, TASK)
    out = TASK / f"段{seg}{'_' + args.tag if args.tag else ''}_1280x704.mp4"
    if local != out:
        local.replace(out)
    probe = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames",
         "-show_entries", "stream=nb_read_frames,width,height",
         "-show_entries", "format=duration", "-of", "json", str(out)],
        capture_output=True, text=True).stdout
    print(f"[seg{seg}] 已保存: {out}")
    print(probe)


if __name__ == "__main__":
    main()