# -*- coding: utf-8 -*-
"""
orchestrator.py · 核心节选
==========================

完整脚本约 1000 行，负责一整条漫剧生产管线：

    项目解析（md / json）→ 逐镜头构图 → 批量提交 ComfyUI → 轮询 → 产物落盘 → 失败重试

这里只保留最关键的「提交 / 轮询 / 落盘」三段，路径已改为相对占位，
配合同目录下的工作流 JSON，可以直接在自己的 ComfyUI 上复现。
"""

import json
import os
import re
import time
import urllib.parse
import urllib.request
from pathlib import Path

# ==================== 配置 ====================
COMFYUI_URL = "http://127.0.0.1:8188"
WORKFLOW_DIR = Path("./workflows")     # 放 ComfyUI 导出的 UI 格式工作流
OUTPUT_DIR = Path("./output")          # 产物落盘目录
LOG_FILE = Path("./orchestrator.log")

WF_T2I = "01-Krea2文生图.json"
WF_DUAL = "导演台双采_12G优化_完善版.json"


# ==================== 工具函数 ====================
def log(msg):
    line = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}"
    print(line, flush=True)
    try:
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            f.write(line + "\n")
    except Exception:
        pass


def http_get(url, timeout=30):
    req = urllib.request.Request(url)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


def http_post(url, data, timeout=30):
    req = urllib.request.Request(
        url,
        data=json.dumps(data).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def sanitize(name):
    """文件名消毒：去掉 Windows 非法字符，截断到 80 字符"""
    name = re.sub(r'[\\/:*?"<>|]', "_", name)
    name = name.replace("：", "_").replace("（", "(").replace("）", ")")
    return name[:80]


def server_ok():
    try:
        http_get(f"{COMFYUI_URL}/system_stats", timeout=5)
        return True
    except Exception:
        return False


def load_workflow(name):
    """加载工作流 JSON（ComfyUI UI 格式）"""
    p = os.path.join(WORKFLOW_DIR, name)
    return json.load(open(p, encoding="utf-8-sig"))


# ==================== 提交与轮询 ====================
def submit_and_wait(prompt, timeout=1800, label=""):
    """提交 prompt，轮询 history，返回 (成功, task)"""
    try:
        result = http_post(f"{COMFYUI_URL}/prompt", {"prompt": prompt})
        pid = result.get("prompt_id")
        log(f"  [{label}] 已提交 {pid}")
    except Exception as e:
        log(f"  [{label}] 提交失败: {e}")
        return False, None

    start = time.time()
    last_log = 0
    while time.time() - start < timeout:
        try:
            h = json.loads(http_get(f"{COMFYUI_URL}/history/{pid}", timeout=10))
            if pid in h:
                st = h[pid].get("status", {})
                if st.get("completed"):
                    log(f"  [{label}] 完成 ({int(time.time() - start)}s)")
                    return True, h[pid]
                msgs = st.get("messages", [])
                if msgs:
                    log(f"  [{label}] 异常: {json.dumps(msgs, ensure_ascii=False)[:300]}")
                    return False, h[pid]
        except Exception:
            pass
        el = int(time.time() - start)
        if el - last_log >= 60:
            log(f"  [{label}] 生成中... {el}s")
            last_log = el
        time.sleep(5)
    log(f"  [{label}] 超时({timeout}s)")
    return False, None


def download_outputs(task, out_dir, rename=None):
    """从 history 里把所有图片 / 视频下载到 out_dir"""
    os.makedirs(out_dir, exist_ok=True)
    saved = []
    outputs = task.get("outputs", {})
    for nid, node_out in outputs.items():
        for key in ("images", "gifs", "videos"):
            if key in node_out:
                for f_info in node_out[key]:
                    fn = f_info.get("filename", "")
                    sub = f_info.get("subfolder", "")
                    url = (
                        f"{COMFYUI_URL}/view?filename={urllib.parse.quote(fn)}"
                        f"&subfolder={urllib.parse.quote(sub)}&type=output"
                    )
                    try:
                        data = http_get(url, timeout=60)
                        if len(data) > 10240:
                            out_path = os.path.join(out_dir, rename or fn)
                            with open(out_path, "wb") as f:
                                f.write(data)
                            saved.append(out_path)
                            log(f"  保存: {os.path.basename(out_path)} ({len(data)//1024}KB)")
                    except Exception as e:
                        log(f"  下载失败 {fn}: {e}")
    return saved


# ==================== 调用示例 ====================
if __name__ == "__main__":
    if not server_ok():
        raise SystemExit("ComfyUI 未启动：http://127.0.0.1:8188")

    wf = load_workflow(WF_T2I)
    # 实际使用时，这里会把 wf 转成 API 格式并填入镜头参数，
    # 再交给 submit_and_wait；产物由 download_outputs 落盘。
    log(f"已加载工作流：{WF_T2I}（{len(wf.get('nodes', []))} 个节点）")