#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""DSH 机 → GPU 机 V18 流水线核心模块。

职责：
1. 接收用户上传的图片 + 生成要求（中文）
2. 调用本地 Qwen2.5-VL-3B（llama.cpp）解析图片 → 生成 V18 REF2VA 英文提示词
3. 把图片上传到 GPU 机 ComfyUI input 目录
4. 基于 v18_api_template.json 构造 API prompt，POST /prompt
5. 轮询 /history 直到完成，返回输出视频信息
6. 从 GPU 机拉回视频到 dl-hub

依赖：vlm-env（llama-cpp-python）、requests
"""
import json
import os
import time
import uuid
import base64
import shutil
import subprocess
import sys
from pathlib import Path

import requests

# ---------------- 配置 ----------------
GPU_HOST = "192.168.31.25"
COMFY_PORT = 8189
COMFY_URL = f"http://{GPU_HOST}:{COMFY_PORT}"
DSH_DL_HUB = Path("/home/zyw/Downloads/dl-hub")
OUTPUT_SUBDIR = "10-视频生成流水线"
OUTPUT_DIR = DSH_DL_HUB / OUTPUT_SUBDIR

V18_TEMPLATE = Path("/home/zyw/Downloads/dasiwa-h3-research/v18_api_template.json")
VLM_MODEL = Path("/home/zyw/Downloads/vlm-models/Qwen2.5-VL-3B-Instruct-Q4_K_M.gguf")
VLM_MMPROJ = Path("/home/zyw/Downloads/vlm-models/mmproj-F16.gguf")
VLM_ENV_PY = "/home/zyw/Downloads/vlm-env/bin/python"
VLM_SERVER_SCRIPT = Path("/home/zyw/Downloads/vlm-env/vlm_server.py")

NO_PROXY = {"no_proxy": "*", "NO_PROXY": "*", "http": "", "https": ""}

# ---------------- V18 REF2VA 提示词生成 ----------------
# 加载用户的 H3 提示词写法规范（docx 提取的纯文本）
PROMPT_GUIDE = Path("/home/zyw/Downloads/dasiwa-h3-research/h3_prompt_guide.txt")
_H3_GUIDE = PROMPT_GUIDE.read_text(encoding="utf-8") if PROMPT_GUIDE.exists() else ""
# 截断到"可复制模板"之前（避免模板占位符诱导模型照抄）
_TEMPLATE_IDX = _H3_GUIDE.find("可复制模板")
if _TEMPLATE_IDX > 0:
    _H3_GUIDE = _H3_GUIDE[:_TEMPLATE_IDX]

SYSTEM_PROMPT = f"""You are an expert video prompt engineer for the MiniMax H3 video model (Ref2VA mode: reference images -> video).
The user provides reference image(s) and a generation requirement (in Chinese).
You MUST produce ONE English structured prompt with EXACTLY these six section labels, in this fixed order:
subject_definitions:
summary:
retention_analysis:
detailed_description:
overall_soundscape:
non_diegetic_music:

=== OFFICIAL RULES (follow strictly) ===
{_H3_GUIDE[:7000]}

=== CRITICAL OUTPUT RULES ===
- Upload order = picture number. The 1st uploaded image is <Picture 1>. Refer to it as <Picture 1>.
- summary starts with [reference generation] and states the clip duration in seconds.
- Each <Subject N> is defined in subject_definitions with its appearance from <Picture N>.
- retention_analysis: one line per subject/picture, format "<tag> (appears in [Shot 1]): fully_preserved - <what stays identical>". Use fully_preserved / partially_preserved / weak_reference.
- detailed_description: start with "[Shot 1] <style>. Vertical 9:16. One continuous take, no cut." then timeline segments with official camera words (static shot / pan / push in / truck / tracking shot). Dialogue verbatim inside <d>[Chinese] 原文</d> if the user gave dialogue. List forbidden items explicitly (no cut, no tracking shot if camera is static, etc).
- overall_soundscape: 1-4 English sentences of environment/action sounds only.
- non_diegetic_music: instruments + tempo + dynamics, or write N/A if no music.
- NEVER output template text, placeholder brackets, or the words "[Who does what, when]", "[What happens after the line]". Fill all sections with concrete content derived from the image and requirement.
- Output ONLY the six sections. No commentary, no markdown fences."""

def build_vlm_payload(image_path: str, requirement: str, duration: int, vision_description: str = "") -> dict:
    """构造 llama.cpp OpenAI 兼容 API 的请求体（自动缩放图片加速 clip 编码）。

    vision_description 传入后可让提示词生成阶段聚焦于"要求+识图结果"，减少重复识图负担。
    """
    image_url = _image_url(image_path) if image_path else None
    if vision_description:
        user_msg = (
            f"Reference image already analyzed:\n{vision_description}\n\n"
            f"Generation requirement (Chinese): {requirement}\n"
            f"Clip duration: {duration} seconds.\n"
            f"Using the image analysis above as ground truth for <Picture 1>, "
            f"now produce the six-section REF2VA prompt in English."
        )
    else:
        user_msg = (
            f"Reference image(s) attached. Generation requirement (Chinese): {requirement}\n"
            f"Clip duration: {duration} seconds.\n"
            f"Now produce the six-section REF2VA prompt in English."
        )
    content = []
    if image_url:
        content.append({"type": "image_url", "image_url": {"url": image_url}})
    content.append({"type": "text", "text": user_msg})
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": content},
        ],
        "temperature": 0.6,
        "max_tokens": 2048,
    }


def start_vlm_server():
    """确保 llama.cpp 视觉模型 HTTP 服务在跑，返回 base URL。"""
    base = "http://127.0.0.1:8089/v1"
    try:
        r = requests.get(base + "/models", timeout=3)
        if r.status_code == 200:
            return base
    except Exception:
        pass
    # 启动服务（后台）
    cmd = [
        VLM_ENV_PY, str(VLM_SERVER_SCRIPT),
        "--model", str(VLM_MODEL),
        "--mmproj", str(VLM_MMPROJ),
        "--host", "127.0.0.1", "--port", "8089",
    ]
    log = open("/home/zyw/Downloads/vlm-server.log", "a")
    subprocess.Popen(cmd, stdout=log, stderr=log, start_new_session=True)
    # 等待就绪
    for _ in range(60):
        time.sleep(2)
        try:
            r = requests.get(base + "/models", timeout=3)
            if r.status_code == 200:
                return base
        except Exception:
            continue
    raise RuntimeError("VLM server 启动超时，查看 /home/zyw/Downloads/vlm-server.log")


def analyze_image_and_build_prompt(image_path: str, requirement: str, duration: int):
    """两阶段解析：先识图（中文详细描述），再基于识图结果+要求生成六段式提示词。

    返回 (vision_description, prompt_text)。
    """
    base = start_vlm_server()

def describe_image(image_path: str) -> str:
    """阶段 1：识图——用中文详细描述图片内容，返回结构化中文描述。"""
    base = start_vlm_server()
    describe_payload = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": _image_url(image_path)}},
                    {"type": "text", "text":
                        "请用中文详细描述这张参考图片的内容，用于视频生成的身份/场景锚定。"
                        "按以下结构输出：\n"
                        "1) 主体人物：性别、年龄、发型、面部特征、服装款式与颜色、配饰\n"
                        "2) 场景环境：地点、建筑、背景元素、光线氛围\n"
                        "3) 关键道具：物品名称、外观、持握方式\n"
                        "4) 画面构图：主体位置、景别、视角\n"
                        "5) 整体风格：写实/插画/古风等，一句话概括"},
                ],
            }
        ],
        "temperature": 0.3,
        "max_tokens": 800,
    }
    r = requests.post(base + "/chat/completions", json=describe_payload, timeout=900)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"].strip()


def build_prompt_from_vision(image_path: str, vision: str, requirement: str, duration: int) -> str:
    """阶段 2：基于识图结果 + 用户要求，生成六段式 Ref2VA 提示词。"""
    base = start_vlm_server()
    prompt_payload = build_vlm_payload(image_path, requirement, duration, vision_description=vision)
    r = requests.post(base + "/chat/completions", json=prompt_payload, timeout=900)
    r.raise_for_status()
    raw = r.json()["choices"][0]["message"]["content"].strip()
    return ensure_six_sections(raw)


def analyze_image_and_build_prompt(image_path: str, requirement: str, duration: int,
                                   on_vision=None):
    """两阶段解析：先识图（中文详细描述），再基于识图结果+要求生成六段式提示词。

    返回 (vision_description, prompt_text)。
    on_vision 可选回调：识图完成时立即调用 on_vision(vision)，用于前端实时展示。
    """
    vision = describe_image(image_path)
    if on_vision:
        on_vision(vision)
    prompt_text = build_prompt_from_vision(image_path, vision, requirement, duration)
    return vision, prompt_text


def analyze_requirement(image_path: str, requirement: str):
    """独立拆解功能：识图 + 把生成要求拆解为结构化要素（不提交 GPU）。

    返回 dict: {"vision": 识图结果, "breakdown": 要求拆解, "raw": 原始响应}
    """
    base = start_vlm_server()
    result = {}

    # 1. 识图（有图时）
    if image_path and Path(image_path).exists():
        result["vision"] = describe_image(image_path)
    else:
        result["vision"] = "（未提供参考图）"

    # 2. 拆解要求
    breakdown_prompt = (
        "你是视频生成需求拆解助手。请把用户的生成要求拆解成以下结构化要素（中文输出，逐条给出）：\n"
        "1. 动作/情节：人物做什么、发生什么事件\n"
        "2. 运镜：镜头运动方式（静止/推/拉/摇/跟等）\n"
        "3. 场景/环境：地点、背景、氛围\n"
        "4. 人物/身份：主角是谁、外观特征（可结合参考图识别结果）\n"
        "5. 台词：如有明确台词，原文列出；没有则写'无'\n"
        "6. 风格：写实/插画/古风/电影感等\n"
        "7. 时长与节奏：起承转合的时间分配\n"
        "8. 禁止项/注意事项：明确不要出现的内容\n"
        "如果要求里有不清楚或矛盾的地方，明确指出。\n\n"
        f"参考图识别结果：\n{result['vision']}\n\n"
        f"用户要求：\n{requirement}"
    )
    payload = {
        "messages": [
            {"role": "system", "content": "你是专业的视频生成需求拆解助手，输出简洁准确的中文结构化拆解。"},
            {"role": "user", "content": breakdown_prompt},
        ],
        "temperature": 0.3,
        "max_tokens": 1000,
    }
    r = requests.post(base + "/chat/completions", json=payload, timeout=900)
    r.raise_for_status()
    result["breakdown"] = r.json()["choices"][0]["message"]["content"].strip()
    result["raw"] = json.dumps(result, ensure_ascii=False)
    return result


def _image_url(image_path: str) -> str:
    """构造图片 data URL（自动缩放）。"""
    read_path = _resize_image(image_path)
    with open(read_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    ext = Path(read_path).suffix.lstrip(".").lower() or "jpeg"
    if ext == "jpg":
        ext = "jpeg"
    return f"data:image/{ext};base64,{b64}"


def _resize_image(image_path: str) -> str:
    """缩放图片到 ≤768px，返回缩放后路径（无法缩放则返回原路径）。"""
    try:
        from PIL import Image
        img = Image.open(image_path)
        max_side = max(img.size)
        if max_side > 768:
            scale = 768.0 / max_side
            img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)
        tmp_path = image_path + ".vlm_resized.jpg"
        img.convert("RGB").save(tmp_path, quality=92)
        return tmp_path
    except Exception:
        return image_path


def ensure_six_sections(text: str) -> str:
    """确保六段式结构完整：缺失的段补默认占位，避免 V18 Director 解析失败。"""
    sections = [
        "subject_definitions",
        "summary",
        "retention_analysis",
        "detailed_description",
        "overall_soundscape",
        "non_diegetic_music",
    ]
    # 按顺序切分已有内容
    idx = {}
    for s in sections:
        i = text.find(s + ":")
        if i >= 0:
            idx[s] = i
    if len(idx) == len(sections):
        return text  # 六段齐全
    # 重建：把已识别的段按顺序保留，缺失段补空
    ordered = sorted(idx.items(), key=lambda kv: kv[1])
    parts = []
    for i, (name, pos) in enumerate(ordered):
        end = ordered[i + 1][1] if i + 1 < len(ordered) else len(text)
        parts.append(text[pos:end].strip())
    result_parts = []
    pi = 0
    for s in sections:
        if s in idx:
            result_parts.append(parts[pi])
            pi += 1
        else:
            if s == "non_diegetic_music":
                result_parts.append("non_diegetic_music:\nN/A")
            elif s == "overall_soundscape":
                result_parts.append("overall_soundscape:\nAmbient environmental sound.")
            else:
                result_parts.append(f"{s}:\n")
    return "\n\n".join(result_parts)


# ---------------- GPU 机交互 ----------------
def upload_image_to_comfy(local_path: str, remote_name: str) -> str:
    """上传图片到 GPU 机 ComfyUI input 目录（用 /upload/image API）。"""
    with open(local_path, "rb") as f:
        files = {"image": (remote_name, f, "image/jpeg")}
        data = {"type": "input", "overwrite": "true"}
        r = requests.post(f"{COMFY_URL}/upload/image", files=files, data=data, timeout=120)
    r.raise_for_status()
    return remote_name


def _upload_file_to_comfy(local_path, remote_name: str) -> str:
    """上传任意文件（含视频）到 GPU 机 ComfyUI input 目录。"""
    ext = Path(remote_name).suffix.lstrip(".").lower() or "bin"
    mime = {"mp4": "video/mp4", "webm": "video/webm", "mov": "video/quicktime"}.get(ext, "application/octet-stream")
    with open(local_path, "rb") as f:
        files = {"image": (remote_name, f, mime)}
        data = {"type": "input", "overwrite": "true"}
        r = requests.post(f"{COMFY_URL}/upload/image", files=files, data=data, timeout=600)
    r.raise_for_status()
    return remote_name


def build_v18_prompt(remote_images, prompt_text: str,
                     width: int, height: int, duration: int,
                     seed: int, steps: int, fps: float = 24.0,
                     upscale_rtx: bool = False, rtx_scale: float = 2.0,
                     force_h264: bool = False) -> dict:
    """基于模板构造 V18 API prompt。

    remote_images: str 或 list[str] —— GPU 机 input 目录里的参考图文件名，
    顺序即 <Picture 1>、<Picture 2>...（对应 timeline slot 0、1...）。
    upscale_rtx: 在 VAE 解码后插入 DaSiWa_RTX_UpscalerRefiner（NVIDIA VSR）放大。
    rtx_scale: 放大倍数（1.0-4.0）。
    force_h264: 主输出强制 H.264（默认 Auto→AV1）。OpenCV/FaceRefine 读不了 AV1，
                后续要接 FaceRefine 时建议开启，避免二次转码。
    """
    if isinstance(remote_images, str):
        remote_images = [remote_images]
    template = json.loads(V18_TEMPLATE.read_text())
    api = template["prompt"]

    # Director 节点
    director = api["2730"]["inputs"]
    director["mode"] = "REF2VA"
    director["prompt"] = prompt_text
    director["width"] = width
    director["height"] = height
    director["duration"] = duration
    director["frame_rate"] = fps
    # timeline_data：参考图 items（多图，slot 按上传顺序）
    now = int(time.time() * 1000)
    items = []
    for i, fn in enumerate(remote_images):
        items.append({
            "id": f"image-{now}-{i}",
            "enabled": True, "order": i, "slot": i, "start": 0, "duration": 1,
            "type": "image", "value": fn,
            "thumbnail": None, "source_width": 0, "source_height": 0,
        })
    timeline = {
        "version": 1,
        "items": items,
        "prompt_blocks": [],
    }
    director["timeline_data"] = json.dumps(timeline, ensure_ascii=False)
    # builder_state
    builder = {
        "version": 2, "mode": "REF2VA", "duration": duration,
        "ref": {"subject_definitions": "", "summary": "", "retention_analysis": "",
                "detailed_description": "", "soundscape": "", "music": ""},
        "prompt_mode": "simple",
        "simple_prompt": prompt_text,
        "resolution": {"aspect": "auto", "resolution": "0.52 MP - SD", "input_scaling": "Auto",
                       "custom_aspect_w": 16, "custom_aspect_h": 9, "custom_mp": 1,
                       "custom_width": width, "custom_height": height},
    }
    director["builder_state"] = json.dumps(builder, ensure_ascii=False)

    # 采样参数（两个 BasicScheduler）
    for nid in ("1512:2590", "1512:2679"):
        if nid in api:
            api[nid]["inputs"]["steps"] = steps
    # seed
    if "2739" in api:
        api["2739"]["inputs"]["seed_value"] = seed
        api["2739"]["inputs"]["seed_control_state"] = json.dumps(
            {"mode": "fixed", "last_seed": str(seed), "recent": [str(seed)]})

    # RTX VSR 放大：在 ComfySwitchNode(1512:2677) 之后、EnhancedVideoCombine(2568) 之前插入
    if upscale_rtx:
        rtx_id = "rtx_upscale"
        api[rtx_id] = {
            "class_type": "DaSiWa_RTX_UpscalerRefiner",
            "inputs": {
                "images": api["2568"]["inputs"]["images"],  # 接原 Combine 的输入（Switch 输出）
                # required 全量参数
                "denoise": False,
                "denoise_quality": "Low",
                "deblur": False,
                "deblur_quality": "Low",
                "upscale": "VSR",
                "upscale_quality": "High",
                "resize_type": "Scale",
                "scale": rtx_scale,
                "megapixels": 2.0,
                "width": int(width * rtx_scale),
                "height": int(height * rtx_scale),
                "divisible_by": "8",
                "ratio_preset": "16:9",
                "resize_method": "Center Crop (Fill)",
                "device_id": 0,
                # optional
                "empty_cache": False,
                "use_mmap": False,
                "auto_unload_models": True,
            },
        }
        # 把 Combine 的 images 改接 RTX 输出
        api["2568"]["inputs"]["images"] = [rtx_id, 0]

    # 主输出强制 H.264（默认 Auto → AV1；OpenCV/FaceRefine 读不了 AV1）
    if force_h264:
        api["2568"]["inputs"]["codec"] = "H.264"
        if api["2568"]["inputs"].get("container", "Auto") == "Auto":
            api["2568"]["inputs"]["container"] = "MP4"

    return api


def build_face_refine_prompt(video_filename: str, ref_images: list,
                             prompt_text: str, duration: int = 10,
                             seed: int = 666666, steps: int = 8,
                             video_abs_path: str = None,
                             upscale_rtx: bool = False, rtx_scale: float = 2.0) -> dict:
    """构造 FaceRefine 二次生成 API 图（ComfyUI-H3-FaceRefine 节点）。

    流程：加载 H3 主视频 → H3FaceTrackCrop 逐帧裁脸 → H3 重生成人脸区域
          → H3PerFrameDenoise → H3FaceStitch 缝合 →（可选 RTX VSR 放大）→ SaveVideo
    video_filename: 视频文件名（用于展示）。
    video_abs_path: 视频在 GPU 机上的绝对路径（VHS_LoadVideoPath 需要绝对路径）。
    upscale_rtx: 缝合后经 NVIDIA RTX VSR 放大。
    rtx_scale: 放大倍数。
    参考示例工作流 H3_Face_Refine_Auto_Select.json 的连接。
    """
    fps = 24.0
    length = int(duration * fps)
    ref0 = ref_images[0] if ref_images else "example.png"
    vpath = video_abs_path or video_filename

    api = {
        # --- 加载主视频（绝对路径）---
        "fr_load": {
            "class_type": "VHS_LoadVideoPath",
            "inputs": {"video": vpath, "force_rate": 0,
                       "custom_width": 0, "custom_height": 0, "frame_load_cap": 0,
                       "skip_first_frames": 0, "select_every_nth": 1, "start_frame": 0,
                       "end_frame": 0, "maximum_frames": 0, "force_size": "Disabled",
                       "custom_audio": 0, "audio_start_index": 0,
                       "pingpong": False},
        },
        # --- 人脸跟踪 + 裁剪 ---
        "fr_track": {
            "class_type": "H3FaceTrackCrop",
            "inputs": {
                "images": ["fr_load", 0],
                "identity_reference": None,
                "detector": "bbox/face_yolov8m.pt",
                "confidence": 0.35,
                "crop_factor": 2.5,
                "canvas_width": 768, "canvas_height": 768,
                "canvas_mode": "auto_capped_768",
                "smooth_window": 21, "size_smooth_window": 51,
                "smooth_method": "gaussian",
                "size_mode": "per_frame",
                "identity_track": True,
                "identity_threshold": 0.28,
                "select": "largest_face",
                "fallback_detector": "none",
                "fallback_head_frac": 0.5,
                "select_index": 0,
                "identity_model": "insightface",
                "cut_detection": "none",
                "cut_threshold": 3.0,
                "absent_shots": "off",
                "X": 0, "Y": 0, "frame_index": 0,
            },
        },
        # --- 模型链 ---
        "fr_unet": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": "minimax_h3_fl2va_pruned_int8_convrot.safetensors", "weight_dtype": "default"},
        },
        "fr_lora": {
            "class_type": "LoraLoaderModelOnly",
            "inputs": {"model": ["fr_unet", 0], "lora_name": "minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", "strength_model": 1.0},
        },
        "fr_clip": {
            "class_type": "CLIPLoader",
            "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax", "device": "default"},
        },
        "fr_vae": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": "minimax_h3_video_vae_int8_convrot.safetensors"},
        },
        "fr_avae": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"},
        },
        # --- 参考图（身份锚定）---
        "fr_refimg": {
            "class_type": "LoadImage",
            "inputs": {"image": ref0, "upload": "image"},
        },
        # --- H3 重生成人脸：ReferenceToVideo + latent 注入 ---
        "fr_reftv": {
            "class_type": "MiniMaxH3ReferenceToVideo",
            "inputs": {
                "clip": ["fr_clip", 0], "vae": ["fr_vae", 0], "audio_vae": ["fr_avae", 0],
                "prompt": prompt_text, "width": 768, "height": 768, "length": length,
                "ref_image_size": "match",
                "ref_images.ref_image_0": ["fr_refimg", 0],
            },
        },
        "fr_inject": {
            "class_type": "H3InjectVideoLatent",
            "inputs": {
                "av_latent": ["fr_reftv", 1], "images": ["fr_track", 0], "vae": ["fr_vae", 0],
            },
        },
        # 逐帧降噪（模型路径；不用 NativeAudioLock，音频从原视频直接走）
        "fr_denoise": {
            "class_type": "H3PerFrameDenoise",
            "inputs": {
                "model": ["fr_lora", 0], "av_latent": ["fr_inject", 0],
                "transform": ["fr_track", 1],
                "denoise_multiplier_small_face": 0.4,
                "denoise_multiplier_large_face": 0.2,
                "scale_mode": "relative_to_clip",
                "face_px_small": 30.0,
                "face_px_large": 120.0,
                "gamma": 1.0,
                "smooth_frames": 9,
            },
        },
        # --- 采样 ---
        "fr_noise": {"class_type": "RandomNoise", "inputs": {"noise_seed": seed}},
        "fr_guider": {"class_type": "BasicGuider", "inputs": {"model": ["fr_denoise", 2], "conditioning": ["fr_reftv", 0]}},
        "fr_sampler": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}},
        "fr_sched": {"class_type": "BasicScheduler", "inputs": {"model": ["fr_denoise", 2], "scheduler": "simple", "steps": steps, "denoise": 1.0}},
        "fr_samp": {
            "class_type": "SamplerCustomAdvanced",
            "inputs": {
                "noise": ["fr_noise", 0], "guider": ["fr_guider", 0],
                "sampler": ["fr_sampler", 0], "sigmas": ["fr_sched", 0],
                "latent_image": ["fr_denoise", 0],
            },
        },
        # --- VAE 解码 ---
        "fr_decode": {"class_type": "VAEDecode", "inputs": {"samples": ["fr_samp", 0], "vae": ["fr_vae", 0]}},
        # --- 缝合（不用 SAM 蒙版，用 transform 检测框蒙版；SAM 模型缺失时可跑）---
        "fr_stitch": {
            "class_type": "H3FaceStitch",
            "inputs": {
                "base_images": ["fr_load", 0], "refined_crops": ["fr_decode", 0],
                "transform": ["fr_track", 1],
                "paste_region": "face_ellipse",
                "mask_dilation": 16, "feather": 6,
                "colour_match": 1.0, "blend": 1.0,
                "undetected_frames": "fade_out",
            },
        },
        # --- 可选：RTX VSR 放大（缝合后）---
        "fr_stitch_out": ["fr_stitch", 0],
    }
    if upscale_rtx:
        api["fr_rtx"] = {
            "class_type": "DaSiWa_RTX_UpscalerRefiner",
            "inputs": {
                "images": ["fr_stitch", 0],
                "denoise": False, "denoise_quality": "Low",
                "deblur": False, "deblur_quality": "Low",
                "upscale": "VSR", "upscale_quality": "High",
                "resize_type": "Scale", "scale": rtx_scale,
                "megapixels": 2.0,
                "width": 2560, "height": 1440,
                "divisible_by": "8", "ratio_preset": "16:9",
                "resize_method": "Center Crop (Fill)",
                "device_id": 0,
                "empty_cache": False, "use_mmap": False, "auto_unload_models": True,
            },
        }
        api["fr_stitch_out"] = ["fr_rtx", 0]

    api["fr_save"] = {
        "class_type": "VHS_VideoCombine",
        "inputs": {
            "images": api["fr_stitch_out"], "audio": ["fr_load", 2],
            "frame_rate": fps, "loop_count": 0, "filename_prefix": "video/FaceRefine",
            "format": "video/h264-mp4", "pix_fmt": "yuv420p", "crf": 18, "save_output": True,
            "pingpong": False,
        },
    }
    del api["fr_stitch_out"]
    return api


def submit_and_wait(api_prompt: dict, timeout_min: int = 60) -> dict:
    """提交 prompt 并轮询直到完成，返回 history 条目。"""
    client_id = f"dasiwa-v18-{uuid.uuid4().hex[:8]}"
    payload = {"prompt": api_prompt, "client_id": client_id}
    r = requests.post(f"{COMFY_URL}/prompt", json=payload, timeout=30)
    if r.status_code != 200:
        raise RuntimeError(f"提交失败 {r.status_code}: {r.text[:500]}")
    prompt_id = r.json()["prompt_id"]

    deadline = time.time() + timeout_min * 60
    while time.time() < deadline:
        time.sleep(10)
        r = requests.get(f"{COMFY_URL}/history/{prompt_id}", timeout=30)
        if r.status_code != 200:
            continue
        data = r.json()
        if prompt_id in data:
            entry = data[prompt_id]
            status = entry.get("status", {})
            if status.get("completed"):
                return entry
            if status.get("status_str") == "error":
                raise RuntimeError(f"生成失败: {json.dumps(status, ensure_ascii=False)[:500]}")
    raise TimeoutError(f"生成超时（>{timeout_min} 分钟）")


def extract_video_output(history_entry: dict):
    """从 history 提取输出视频信息。"""
    outputs = history_entry.get("outputs", {})
    for nid, out in outputs.items():
        for key in ("gifs", "videos"):
            for v in out.get(key, []):
                fn = v.get("filename", "")
                if fn.endswith(".mp4") or fn.endswith(".webm") or fn.endswith(".mov"):
                    return {"filename": fn, "subfolder": v.get("subfolder", ""), "type": v.get("type", "output")}
    # 兜底：查所有 output 字段
    for nid, out in outputs.items():
        for key, val in out.items():
            if isinstance(val, list):
                for v in val:
                    if isinstance(v, dict) and v.get("filename"):
                        fn = v["filename"]
                        if fn.endswith((".mp4", ".webm", ".mov", ".gif")):
                            return {"filename": fn, "subfolder": v.get("subfolder", ""), "type": v.get("type", "output")}
    raise RuntimeError("未在输出中找到视频文件")


def download_video(video_info: dict, local_dir: Path) -> Path:
    """从 GPU 机下载视频到本地目录。"""
    local_dir.mkdir(parents=True, exist_ok=True)
    params = {"filename": video_info["filename"], "subfolder": video_info.get("subfolder", ""), "type": video_info.get("type", "output")}
    r = requests.get(f"{COMFY_URL}/view", params=params, timeout=600, stream=True)
    r.raise_for_status()
    local_path = local_dir / video_info["filename"]
    with open(local_path, "wb") as f:
        shutil.copyfileobj(r.raw, f)
    return local_path


# ---------------- 命名辅助 ----------------
def _safe_name(s: str, max_len: int = 24) -> str:
    """清洗为文件系统安全的短名称（保留中英文数字，去特殊字符）。"""
    import re
    s = re.sub(r'[\\/:*?"<>|\s]+', '_', s.strip())
    s = re.sub(r'_+', '_', s).strip('_')
    return s[:max_len] or "video"


def derive_topic(requirement: str, topic: str = None, max_len: int = 24) -> str:
    """确定主题：优先显式 topic，否则从 requirement 提取关键片段。"""
    if topic and topic.strip():
        return _safe_name(topic, max_len)
    # 从要求提取：取前若干字符（中文按语义截断到完整词）
    req = requirement.strip()
    if not req:
        return "video"
    # 去掉常见引导词，取第一个句子片段
    for sep in ("，", ",", "。", ".", "！", "？", "："):
        idx = req.find(sep)
        if 0 < idx <= max_len:
            req = req[:idx]
            break
    return _safe_name(req, max_len)


def make_task_dir(requirement: str, topic: str = None) -> Path:
    """创建语义化任务目录：YYYYMMDD_HHMM_主题。"""
    ts = time.strftime("%Y%m%d_%H%M")
    name = f"{ts}_{derive_topic(requirement, topic)}"
    task_dir = OUTPUT_DIR / name
    # 同分钟同名冲突时追加序号
    n = 2
    while task_dir.exists():
        task_dir = OUTPUT_DIR / f"{name}_{n}"
        n += 1
    task_dir.mkdir(parents=True, exist_ok=True)
    return task_dir


def _rename_video(src: Path, new_base: str, task_dir: Path) -> Path:
    """把下载的视频重命名为 语义化基础名.扩展名，返回新路径。"""
    ext = src.suffix or ".mp4"
    dest = task_dir / f"{_safe_name(new_base, 60)}{ext}"
    if dest != src:
        if dest.exists():
            dest.unlink()
        shutil.move(str(src), str(dest))
    return dest


# ---------------- 流水线主流程 ----------------
def run_pipeline(image_path, requirement: str,
                 width: int = 736, height: int = 1280, duration: int = 10,
                 seed: int = 666666, steps: int = 8, fps: float = 24.0,
                 on_vision=None, prompt_override: str = None,
                 vision_override: str = None,
                 extra_images: list = None,
                 upscale_rtx: bool = False, rtx_scale: float = 2.0,
                 face_refine: bool = False,
                 topic: str = None, variant: str = None) -> dict:
    """完整流水线：解析 → 上传 → 提交 → 下载 → 归档。

    on_vision 可选回调：识图完成时立即调用 on_vision(vision_desc)，用于前端实时展示。
    prompt_override: 若提供，则跳过 VLM 解析，直接用该提示词提交 GPU（步骤3独立按钮用）。
    vision_override: 若提供，作为已完成的识图结果（不重新识图）。
    extra_images: 附加参考图列表（本地路径），按顺序作为 <Picture 2>、<Picture 3>...
                 主图 image_path 始终是 <Picture 1>。
    upscale_rtx: 生成后经 NVIDIA RTX VSR 放大（需 GPU 机已装 nvidia-vfx）。
    rtx_scale: 放大倍数。
    face_refine: 生成后经 ComfyUI-H3-FaceRefine 二次生成修复小脸（远景/中景崩脸）。
    topic: 任务主题（用于目录/文件命名）；未传则从 requirement 提取。
    variant: 版本标记（如 "v2_来了就好"），追加到视频文件名，同主题多版本区分。
    """
    # 语义化任务目录 + 基础文件名
    task_dir = make_task_dir(requirement, topic)
    task_id = task_dir.name  # 目录名即 task_id（YYYYMMDD_HHMM_主题）
    base_name = task_id
    if variant and variant.strip():
        base_name = f"{task_id}_{_safe_name(variant, 30)}"
    print(f"[{task_id}] 任务目录: {task_dir}")
    print(f"[{task_id}] 基础文件名: {base_name}")

    # seed 处理：0 = 随机
    seed_was_random = (seed == 0)
    if seed == 0:
        import random as _rnd
        seed = _rnd.randint(1, 2**31 - 1)
        print(f"[{task_id}] seed=0 → 随机 seed: {seed}")

    # 记录 seed 到 TXT
    seed_info = (
        f"SEED: {seed}\n"
        f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n"
        f"随机: {'是（seed=0）' if seed_was_random else '否（固定值）'}\n"
        f"时长: {duration}s | 分辨率: {width}x{height} | 步数: {steps} | 帧率: {fps}\n"
    )
    (task_dir / "seed.txt").write_text(seed_info, encoding="utf-8")
    # seed 也追加到 prompt.txt 末尾（方便直接看）
    _prompt_path = task_dir / "prompt.txt"
    if _prompt_path.exists():
        _prompt_path.write_text(
            _prompt_path.read_text(encoding="utf-8").rstrip() + "\n\n--- 生成参数 ---\n" + seed_info,
            encoding="utf-8")

    # 1. 保存用户图片（主图 + 附加图）
    user_img = task_dir / f"ref_{Path(image_path).name}"
    shutil.copy(image_path, user_img)
    extra_local = []
    for i, ep in enumerate(extra_images or [], start=2):
        dest = task_dir / f"ref_pic{i}_{Path(ep).name}"
        shutil.copy(ep, dest)
        extra_local.append(dest)

    # 2. 提示词来源：外部传入 或 两阶段解析
    if prompt_override and prompt_override.strip():
        prompt_text = prompt_override.strip()
        vision_desc = vision_override or "（外部传入提示词，未识图）"
        (task_dir / "prompt.txt").write_text(prompt_text, encoding="utf-8")
        if vision_override:
            (task_dir / "vision_description.txt").write_text(vision_override, encoding="utf-8")
        print(f"[{task_id}] 使用外部提示词（{len(prompt_text)} 字符）")
    else:
        print(f"[{task_id}] 解析图片（识图 + 生成提示词）...")
        vision_desc, prompt_text = analyze_image_and_build_prompt(
            str(user_img), requirement, duration, on_vision=on_vision)
        (task_dir / "vision_description.txt").write_text(vision_desc, encoding="utf-8")
        (task_dir / "prompt.txt").write_text(prompt_text, encoding="utf-8")
        print(f"[{task_id}] 识图完成（{len(vision_desc)} 字符）")
        print(f"[{task_id}] 提示词生成完成（{len(prompt_text)} 字符）")

    # 3. 上传参考图到 GPU 机（主图 + 附加图）
    remote_names = []
    rn = f"v18pipe_{task_id}_{Path(image_path).name}"
    upload_image_to_comfy(str(user_img), rn)
    remote_names.append(rn)
    for i, ep in enumerate(extra_local):
        rn = f"v18pipe_{task_id}_pic{i+2}_{Path(ep).name}"
        upload_image_to_comfy(str(ep), rn)
        remote_names.append(rn)
    print(f"[{task_id}] 图片已上传: {remote_names}")

    # 4. 构造并提交 V18 prompt
    # 后续接 FaceRefine 时强制 H.264 输出（OpenCV 读不了 AV1，避免二次转码）
    api = build_v18_prompt(remote_names, prompt_text, width, height, duration, seed, steps, fps,
                           upscale_rtx=upscale_rtx, rtx_scale=rtx_scale,
                           force_h264=face_refine)
    (task_dir / "api_prompt.json").write_text(json.dumps(api, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"[{task_id}] 提交 V18 工作流 (WxH={width}x{height}, dur={duration}s, steps={steps}, seed={seed}, refs={len(remote_names)}, rtx={upscale_rtx and f'x{rtx_scale}' or 'off'})...")
    entry = submit_and_wait(api)
    print(f"[{task_id}] 生成完成")

    # 5. 提取并下载视频（重命名为语义化文件名）
    video_info = extract_video_output(entry)
    local_video = download_video(video_info, task_dir)
    final_video = _rename_video(local_video, base_name, task_dir)
    print(f"[{task_id}] 视频已下载: {final_video}")

    # 5.5 人脸修复阶段（二次生成）
    if face_refine:
        print(f"[{task_id}] 启动 FaceRefine 人脸修复（二次生成，约 5-15 分钟）...")
        # 转码为 H.264（OpenCV 不支持 AV1/H265 视频）
        import subprocess
        h264_path = str(task_dir / "main_h264.mp4")
        cmd = ["ffmpeg", "-y", "-i", str(final_video), "-c:v", "libx264",
               "-pix_fmt", "yuv420p", "-c:a", "aac", h264_path]
        subprocess.run(cmd, capture_output=True, timeout=600)
        # 上传转码后视频到 GPU 机 input 目录
        remote_video = f"v18pipe_{task_id}_main.mp4"
        _upload_file_to_comfy(h264_path, remote_video)
        remote_ref = remote_names[0]
        fr_api = build_face_refine_prompt(
            remote_video, [remote_ref], prompt_text,
            duration=duration, seed=seed, steps=steps,
            video_abs_path=f"/home/zyw/ComfyUI/input/{remote_video}",
            # 主阶段已 RTX 放大到 2K；FaceRefine 阶段不再重复放大（避免 5K 双重放大）
            upscale_rtx=False, rtx_scale=rtx_scale)
        (task_dir / "api_prompt_facerefine.json").write_text(
            json.dumps(fr_api, ensure_ascii=False, indent=1), encoding="utf-8")
        fr_entry = submit_and_wait(fr_api, timeout_min=60)
        fr_info = extract_video_output(fr_entry)
        fr_video = download_video(fr_info, task_dir / "facerefine")
        final_video = _rename_video(fr_video, f"{base_name}_facerefine", task_dir)
        print(f"[{task_id}] 人脸修复完成: {final_video}")

    result = {
        "task_id": task_id,
        "vision_description": vision_desc,
        "prompt": prompt_text,
        "video_local": str(final_video),
        "video_url": f"http://192.168.31.76:8899/{OUTPUT_SUBDIR}/{task_id}/{final_video.name}",
        "task_dir": str(task_dir),
    }
    # 记录任务元信息
    meta = {
        "task_id": task_id,
        "requirement": requirement,
        "params": {"width": width, "height": height, "duration": duration, "seed": seed, "steps": steps, "fps": fps},
        "vision_description": vision_desc,
        "prompt": prompt_text,
        "video_url": result["video_url"],
        "time": time.strftime("%Y-%m-%d %H:%M:%S"),
    }
    (task_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8")
    return result


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("用法: python3 pipeline.py <图片路径> <要求文本> [时长秒] [宽] [高] [步数] [seed]")
        sys.exit(1)
    img = sys.argv[1]
    req = sys.argv[2]
    dur = int(sys.argv[3]) if len(sys.argv) > 3 else 10
    w = int(sys.argv[4]) if len(sys.argv) > 4 else 736
    h = int(sys.argv[5]) if len(sys.argv) > 5 else 1280
    steps = int(sys.argv[6]) if len(sys.argv) > 6 else 8
    seed = int(sys.argv[7]) if len(sys.argv) > 7 else 666666
    res = run_pipeline(img, req, w, h, dur, seed, steps)
    print(json.dumps(res, ensure_ascii=False, indent=2))