#!/usr/bin/env python3
"""生产网关：把 h3-console 现有的 gateway.py/api.py 契约（/health, /v1/videos...）搬到
ComfyUI + Raylight 之上。跟旧架构最大的不同：旧架构是4卡4个独立worker、gateway负责按partition
选卡；新架构是4卡一个 FSDP 集群，任何时刻只能跑一个请求，所以这里把 gateway 的"选后端"和 api.py
的"单卡worker"合并成一个服务：一条队列，一个后台 worker，串行往 ComfyUI 提交/等待/取结果。

RayInitializer/RayUNETLoader/CLIPLoader/VAELoader 这几个"重"节点的输入值每次提交必须逐字节
相同——这是 ComfyUI 节点级缓存跳过重新建集群/重新加载模型的唯一机制（Raylight 自己在这条量化+
FSDP 的配置下没有安全网，见 docs/todo/raylight-comfyui-multi-gpu-eval.md）。所以下面的
RAY_INIT_CFG 是全局常量，任何请求都不应该改它。
"""
import asyncio
import contextlib
import json
import os
import queue
import re
import sqlite3
import subprocess
import threading
import time
import urllib.parse
import urllib.request
import uuid
from pathlib import Path
from typing import Literal, Optional

import httpx
import websockets
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field, model_validator

COMFY_HOST = os.environ.get("COMFY_HOST", "127.0.0.1")
COMFY_PORT = os.environ.get("COMFY_PORT", "8188")
COMFY_BASE = f"http://{COMFY_HOST}:{COMFY_PORT}"
COMFY_WS = f"ws://{COMFY_HOST}:{COMFY_PORT}/ws"
# console 的 _try_submit/_row_view 依赖响应里的 backend 字段来判断"任务分配到哪个后端了"，
# 没有这个字段就会一直认为任务"还没分配"，播放/下载会被 409 拒绝——踩过这个坑，2026-08-22。
SERVER_PORT = os.environ.get("PORT", "8100")
BACKEND_URL = f"http://127.0.0.1:{SERVER_PORT}"

ROOT = Path("/data/models/raylight/ComfyUI")
INPUT_DIR = ROOT / "input"
OUTPUT_DIR = ROOT / "output" / "video"
DB_PATH = Path("/data/models/raylight/jobs.sqlite3")
COMFYUI_LOG_PATH = Path("/data/models/raylight/comfyui-server.log")
# ponytail: Raylight 的 FSDP 采样在 Ray 子进程里跑，从不调 comfy.utils.ProgressBar 上报进度，
# 所以下面 _run_job_async 里监听的 progress_state websocket 永远拿不到真实 step——只能退而
# 求其次，轮询子进程 tqdm 直接打到 stdout、被日志捕获的这行文字。上游改了这行的措辞/格式，
# 这个正则就会失效（退化为进度条不动，不影响任务本身跑成功）。只信任 total 等于本次请求
# steps 的那次匹配，因为同一句"Model Initializing"文案在 FSDP 初始化阶段也会打印、但那里的
# 总数是 GPU 数（4），跟真正的采样步数（这里的 steps 字段）容易混淆。
_PROGRESS_RE = re.compile(r"Model Initializing.*?(\d+)\.\d+/(\d+)\.\d+ \[")
# 2026-08-24: 单纯截尾([-2000:])在参考图/参考视频的报错里会失效——Ray traceback里经常在真正
# 报错行后面接着打印一大坨tensor内容（比如OOM报错后把出错的输入张量也dump出来），尾部2000字符
# 全是"0., 0., 0., ..."这种数字，看不到真正的报错类型/原因。改成直接从整段文本里找
# "XxxError: ..."这种模式，取最后一个匹配（Python异常链里最后抛出的通常是最具体的结论），
# 找不到才退回截尾兜底。
_ERROR_LINE_RE = re.compile(r'\b([A-Za-z_][\w.]*(?:Error|Exception)):\s*([^\n"\\]*)')


def _extract_error_summary(messages: list) -> str:
    raw = json.dumps(messages, ensure_ascii=False)
    text = raw.replace("\\n", "\n")  # 还原JSON转义的换行，不然"n"会跟前一个词粘连（比如"ntorch..."）
    matches = _ERROR_LINE_RE.findall(text)
    if matches:
        return f"{matches[-1][0]}: {matches[-1][1].strip()}"[:2000]
    return raw[-2000:]

# 每次提交必须逐字节相同才能让 ComfyUI 跳过重建集群——不要给任何一个请求字段改这几个值。
RAY_INIT_CFG = {
    "ray_cluster_address": "local",
    "ray_cluster_namespace": "default",
    "GPU": 4,
    "ulysses_degree": 4,
    "ring_degree": 1,
    "cfg_degree": 1,
    "dp_degree": 1,
    "sync_ulysses": False,
    "clear_vram_after_sampling": False,
    "FSDP": True,
    "FSDP_CPU_OFFLOAD": False,
    # 2026-08-22：TORCH_MATH 物化完整attention矩阵，是之前 832x480x124帧/448x256x362帧/
    # 1024x576x60帧 反复 OOM 的根因，不是权重装不下。yunchang 其实原生支持 SAGE_FP16，只是环境里
    # 从没装过给Turing(sm_75)编译的sageattention——装上以后三个档位全部通过，GPU0峰值~13.3GB，
    # 还快很多（832x480x124帧从OOM变成100秒）。sageattention来自社区分支
    # github.com/1506086927/SageAttention2_Optimized_Test（经 github.com/Eanya-Tonic/2080ti-minimax-h3
    # 的补丁接入），详见 notes.md。
    "XFuser_attention": "SAGE_FP16",
    "RAYLIGHT_ULYSSES_KV_INT8": "off",  # 通信层INT8量化，跟这次的显存修复无关，保持关闭
    "skip_comm_test": False,
    "use_mmap": True,
}
# 2026-08-23：4个分辨率×帧数排列组合实测出的真实CUDA OOM边界（832×480/1024×576/1344×768）
# 或画质边界（448×256，见notes.md——技术上测到3000帧都没OOM，是用户看过实际输出判定50秒往后
# 画质不能看，不是显存拦的）。跟 h3-console 前端 RESOLUTIONS 里的 maxFrames 保持一致——
# 两边算法对不上就会出现"前端放行、后端拒绝"或反过来，数字改动两边一起改。
RESOLUTION_CEILINGS = {
    (832, 480): 1125,
    (448, 256): 1200,
    (1024, 576): 743,
    (1344, 768): 412,
}
CLIP_NAME = "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"
VAE_VIDEO = "minimax_h3_video_vae_fp16.safetensors"
VAE_AUDIO = "minimax_h3_audio_vae_fp32.safetensors"
DIT_FL2VA = "minimax_h3_fl2va_pruned_int8_convrot.safetensors"
DIT_REF2VA = "minimax_h3_ref2va_pruned_int8_convrot.safetensors"


class GenReq(BaseModel):
    prompt: str
    width: int = Field(832, ge=64, le=2048)
    height: int = Field(480, ge=64, le=2048)
    frames: int = Field(124, ge=5, le=3000)  # 2026-08-23: 448x256在1200帧(旧上限)时GPU远没跑满
    # （轻松成功，646秒），真实显存上限没探到——旧的1200是这个字段自己的硬编码，不是显存拦的。
    # 抬到3000留探测空间，其他分辨率的真实上限已经在1200帧以内测到了，不受这次调整影响。见 notes.md
    steps: int = Field(20, ge=1, le=50)
    turbo: bool = False  # Raylight 图里还没接turbo LoRA，目前是no-op，见 notes.md
    seed: int = Field(0, ge=0, le=2**31 - 1)
    cfg_scale: Optional[float] = None       # 目前是 no-op：MiniMax H3 在 Raylight 里不支持 CFG
    negative_prompt: Optional[str] = None   # 同上，暂不生效，见 notes.md
    first_frame_url: Optional[str] = None
    last_frame_url: Optional[str] = None
    reference_image_urls: list[str] = []
    # 2026-08-23：MiniMaxH3ReferenceToVideo 节点本身支持参考视频/独立音频参考（ComfyUI
    # /object_info 里的 ref_videos/ref_audios），之前只接了 ref_images 这一种，见 notes.md。
    reference_video_urls: list[str] = Field(default_factory=list, max_length=3)  # 最多3段，
    # 每段的音轨会跟着自动带上（对应 ref_video_audios，不用单独传）
    reference_audio_urls: list[str] = Field(default_factory=list, max_length=3)  # 独立音频参考，
    # 不跟任何参考视频绑定，最多3个（对应节点的 ref_audios）

    @model_validator(mode="after")
    def _check_known_resolution_ceiling(self):
        # 只挡已经实测过的4个分辨率——没测过的组合不拦，跟以前一样"能提交但结果未验证"，
        # 不假装知道我们没测过的显存边界。
        ceiling = RESOLUTION_CEILINGS.get((self.width, self.height))
        if ceiling is not None and self.frames > ceiling:
            raise ValueError(
                f"{self.width}x{self.height} 实测上限是 {ceiling} 帧，{self.frames} 帧会真实CUDA OOM"
                f"（或448x256画质会退化，见notes.md），不接受这个请求"
            )
        return self


class Job(BaseModel):
    id: str
    status: Literal["queued", "running", "done", "failed", "cancelled"]
    req: GenReq
    created_at: float
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    step: int = 0
    total_steps: int = 0
    file: Optional[str] = None
    size_bytes: Optional[int] = None
    error: Optional[str] = None
    partition: str = "ref2va"
    comfy_prompt_id: Optional[str] = None
    backend: str = BACKEND_URL


JOBS: dict[str, Job] = {}
Q: "queue.Queue[str]" = queue.Queue()
LOCK = threading.Lock()
STATE = {"model": "loading", "loaded_at": None}


def db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = db()
    conn.execute(
        "CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, updated_at REAL NOT NULL, data TEXT NOT NULL)"
    )
    conn.commit()
    conn.close()


def _persist(job: Job):
    with LOCK:
        conn = db()
        conn.execute(
            "INSERT INTO jobs(id, updated_at, data) VALUES (?,?,?) "
            "ON CONFLICT(id) DO UPDATE SET updated_at=excluded.updated_at, data=excluded.data",
            (job.id, time.time(), job.model_dump_json()),
        )
        conn.commit()
        conn.close()


def _load_jobs():
    conn = db()
    rows = list(conn.execute("SELECT data FROM jobs"))
    conn.close()  # 2026-08-23: 之前在这个连接的SELECT游标还开着的时候就调 _persist() 开新连接写，
    # SQLite 默认rollback-journal模式下第二个连接的写commit会被第一个连接的读事务锁住，报
    # 'database is locked'，导致服务启动直接崩溃——先把结果取完关掉连接，再在下面循环里写。
    for row in rows:
        try:
            j = Job.model_validate_json(row["data"])
        except Exception:
            continue
        # 进程重启=ComfyUI侧的执行状态和排队状态都丢了，跟老系统的策略一致：不可恢复的直接判定
        # 终态。2026-08-23区分开：只是排队中、根本还没开始跑的，判"cancelled"更准确——没有任何
        # 执行过程可言，"failed"意味着跑起来又出错了，这两者不是一回事。真的在跑到一半被打断的
        # 才算failed。
        if j.status == "queued":
            j.status = "cancelled"
            j.error = "进程重启导致这条排队中的任务被取消，需要重新提交"
            j.finished_at = time.time()
            _persist(j)
        elif j.status == "running":
            j.status = "failed"
            j.error = "进程重启导致这条任务的执行状态丢失，无法恢复"
            j.finished_at = time.time()
            _persist(j)  # 2026-08-23: 之前只改内存不落盘，sqlite里的status永远停在running/queued，
            # 下次重启会被这段逻辑重新"发现"一遍、把人工写的详细错因覆盖回这句通用文案——必须写回去
        JOBS[j.id] = j
    conn.close()


def _partition_of(req: GenReq) -> str:
    if req.reference_image_urls or req.reference_video_urls or req.reference_audio_urls:
        return "ref2va"
    return "fl2va"


_INPUT_EXTS = {
    "image": (".png", ".jpg", ".jpeg", ".webp"),
    "video": (".mp4", ".mov", ".webm", ".mkv"),
    "audio": (".mp3", ".wav", ".m4a", ".ogg", ".flac"),
}


def _fetch_to_input(url: str, kind: str = "image") -> str:
    """下载一个文件到 ComfyUI 的 input/ 目录，LoadImage/LoadVideo/LoadAudio 节点按文件名读取。"""
    allowed = _INPUT_EXTS[kind]
    ext = os.path.splitext(urllib.parse.urlparse(url).path)[1].lower()
    if ext not in allowed:
        ext = allowed[0]
    fname = f"{uuid.uuid4().hex}{ext}"
    dest = INPUT_DIR / fname
    with urllib.request.urlopen(url, timeout=120) as resp, open(dest, "wb") as f:
        f.write(resp.read())
    return fname


def _static_nodes() -> dict:
    """三套 workflow 共用的静态节点：CLIP/VAE loader + RayInitializer。节点 id 用字符串前缀
    "s_" 避免跟每次请求动态生成的节点 id 冲突。

    CLIP/VAE 不走 FSDP 分片，默认会全部堆在 GPU0（跟 Ray rank0 的 DiT 分片挤同一张卡，是
    GPU0 在所有 OOM 测试里最先报错的原因之一，见 notes.md）。这个 ComfyUI fork 自带
    SelectCLIPDevice/SelectVAEDevice 节点（comfy_extras/nodes_multigpu.py），用它们把这三样
    分别钉到 GPU1/2/3，把 GPU0 净空出来只留 DiT 分片——2026-08-22 加，效果见 notes.md。"""
    return {
        "s_clip": {"class_type": "CLIPLoader",
                   "inputs": {"clip_name": CLIP_NAME, "type": "minimax", "device": "default"}},
        "s_clip_gpu": {"class_type": "SelectCLIPDevice", "inputs": {"clip": ["s_clip", 0], "device": "gpu:1"}},
        "s_vae_v": {"class_type": "VAELoader", "inputs": {"vae_name": VAE_VIDEO}},
        "s_vae_v_gpu": {"class_type": "SelectVAEDevice", "inputs": {"vae": ["s_vae_v", 0], "device": "gpu:2"}},
        "s_vae_a": {"class_type": "VAELoader", "inputs": {"vae_name": VAE_AUDIO}},
        "s_vae_a_gpu": {"class_type": "SelectVAEDevice", "inputs": {"vae": ["s_vae_a", 0], "device": "gpu:3"}},
        "s_rayinit": {"class_type": "RayInitializer", "inputs": dict(RAY_INIT_CFG)},
    }


def build_fl2va_graph(req: GenReq) -> dict:
    """文生视频 / 首尾帧生视频 —— MiniMaxH3ImageToVideo，不给 first/last frame 就是纯 T2V。"""
    g = _static_nodes()
    g["s_unet"] = {"class_type": "RayUNETLoader",
                   "inputs": {"unet_name": DIT_FL2VA, "weight_dtype": "default", "lora": None,
                              "ray_actors_init": ["s_rayinit", 0]}}

    conditioning_inputs = {
        "prompt": req.prompt, "clip": ["s_clip_gpu", 0], "vae": ["s_vae_v_gpu", 0],
        "width": req.width, "height": req.height, "length": req.frames,
    }
    if req.first_frame_url:
        fname = _fetch_to_input(req.first_frame_url)
        g["n_first"] = {"class_type": "LoadImage", "inputs": {"image": fname}}
        conditioning_inputs["first_frame"] = ["n_first", 0]
    if req.last_frame_url:
        fname = _fetch_to_input(req.last_frame_url)
        g["n_last"] = {"class_type": "LoadImage", "inputs": {"image": fname}}
        conditioning_inputs["last_frame"] = ["n_last", 0]
    g["n_cond"] = {"class_type": "MiniMaxH3ImageToVideo", "inputs": conditioning_inputs}

    _append_sampling_tail(g, req, positive_node="n_cond")
    return g


def build_ref2va_graph(req: GenReq) -> dict:
    """多参考图生视频 —— MiniMaxH3ReferenceToVideo，最多9张参考图（协议规格上限）。"""
    g = _static_nodes()
    g["s_unet"] = {"class_type": "RayUNETLoader",
                   "inputs": {"unet_name": DIT_REF2VA, "weight_dtype": "default", "lora": None,
                              "ray_actors_init": ["s_rayinit", 0]}}

    conditioning_inputs = {
        "prompt": req.prompt, "clip": ["s_clip_gpu", 0], "vae": ["s_vae_v_gpu", 0],
        "audio_vae": ["s_vae_a_gpu", 0], "width": req.width, "height": req.height,
        "length": req.frames, "ref_image_size": "match",
    }
    for i, url in enumerate(req.reference_image_urls[:9]):
        fname = _fetch_to_input(url)
        nid = f"n_ref{i}"
        g[nid] = {"class_type": "LoadImage", "inputs": {"image": fname}}
        conditioning_inputs[f"ref_images.ref_image_{i}"] = [nid, 0]
    for i, url in enumerate(req.reference_video_urls[:3]):
        fname = _fetch_to_input(url, kind="video")
        load_nid, split_nid = f"n_refvid{i}_load", f"n_refvid{i}_split"
        g[load_nid] = {"class_type": "LoadVideo", "inputs": {"file": fname}}
        # GetVideoComponents 把 VIDEO 拆成 images/audio/fps/bit_depth/color_space——
        # ref_video 要 IMAGE 不是 VIDEO，ref_video_audio 用的就是同一段视频自带的音轨，
        # 不需要用户再单独传一份音频。
        g[split_nid] = {"class_type": "GetVideoComponents", "inputs": {"video": [load_nid, 0]}}
        conditioning_inputs[f"ref_videos.ref_video_{i}"] = [split_nid, 0]
        conditioning_inputs[f"ref_video_audios.ref_video_audio_{i}"] = [split_nid, 1]
    for i, url in enumerate(req.reference_audio_urls[:3]):
        fname = _fetch_to_input(url, kind="audio")
        nid = f"n_refaudio{i}"
        g[nid] = {"class_type": "LoadAudio", "inputs": {"audio": fname}}
        conditioning_inputs[f"ref_audios.ref_audio_{i}"] = [nid, 0]
    g["n_cond"] = {"class_type": "MiniMaxH3ReferenceToVideo", "inputs": conditioning_inputs}

    _append_sampling_tail(g, req, positive_node="n_cond")
    return g


def _append_sampling_tail(g: dict, req: GenReq, positive_node: str):
    """采样+解码+存视频这一段两套 workflow 完全一样，抽出来一份。
    cfg_scale/negative_prompt 目前没有入口——MiniMax H3 在 Raylight 这条 pipeline 里不支持 CFG
    parallel（cfg_degree 只能是1），RayBasicGuider 只吃一路 positive conditioning，见 notes.md。"""
    g["n_guider"] = {"class_type": "RayBasicGuider",
                      "inputs": {"ray_actors": ["s_unet", 0], "conditioning": [positive_node, 0]}}
    g["n_sched"] = {"class_type": "RayBasicScheduler",
                     "inputs": {"scheduler": "simple", "steps": req.steps, "denoise": 1.0,
                                "ray_actors": ["s_unet", 0]}}
    g["n_sampler_sel"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}}
    g["n_sample"] = {"class_type": "XFuserSamplerCustomAdvanced",
                      "inputs": {"add_noise": True, "noise_seed": req.seed,
                                 "guider": ["n_guider", 0], "sampler": ["n_sampler_sel", 0],
                                 "sigmas": ["n_sched", 0], "latent_image": [positive_node, 1]}}
    g["n_vaedec"] = {"class_type": "VAEDecode", "inputs": {"samples": ["n_sample", 0], "vae": ["s_vae_v_gpu", 0]}}
    g["n_vaedec_a"] = {"class_type": "VAEDecodeAudio", "inputs": {"samples": ["n_sample", 0], "vae": ["s_vae_a_gpu", 0]}}
    g["n_video"] = {"class_type": "CreateVideo",
                     "inputs": {"fps": 24, "bit_depth": 8, "images": ["n_vaedec", 0], "audio": ["n_vaedec_a", 0]}}
    g["n_save"] = {"class_type": "SaveVideo",
                    "inputs": {"filename_prefix": "video/MiniMax_H3", "format": "auto",
                               "video": ["n_video", 0], "format.codec": "auto"}}


def build_graph(req: GenReq) -> tuple[dict, str]:
    partition = _partition_of(req)
    if partition == "ref2va":
        return build_ref2va_graph(req), partition
    return build_fl2va_graph(req), partition


# ---------------- ComfyUI 交互（单队列 worker，严格串行——4卡是一个集群，同时只能跑一个请求） ----------------

async def _tail_log_progress(j: Job, expected_total: int, stop: asyncio.Event):
    """见 _PROGRESS_RE 定义处的 ponytail 注释：progress_state websocket 拿不到真实 step 时的退路。"""
    while not stop.is_set():
        await asyncio.sleep(2)
        try:
            with open(COMFYUI_LOG_PATH, "rb") as f:
                f.seek(0, 2)
                f.seek(max(0, f.tell() - 8192))
                tail = f.read().decode("utf-8", "ignore")
        except OSError:
            continue
        for m in reversed(list(_PROGRESS_RE.finditer(tail))):
            cur, total = int(m.group(1)), int(m.group(2))
            if total == expected_total and (cur, total) != (j.step, j.total_steps):
                j.step, j.total_steps = cur, total
                _persist(j)
            break


async def _track_ws_progress(j: Job, client_id: str, prompt_id: str):
    """尽力而为的进度上报——完成判定不靠这个（见 _run_job_async 里的注释），这个连接断了/
    超时/漏消息都无所谓，静默退出就行，不能让它的失败影响任务本身的成败判定。"""
    try:
        ws_url = f"{COMFY_WS}?clientId={client_id}"
        async with websockets.connect(ws_url, max_size=None) as ws:
            while True:
                raw = await ws.recv()
                if not isinstance(raw, str):
                    continue
                msg = json.loads(raw)
                data = msg.get("data", {})
                if msg.get("type") == "progress_state" and data.get("prompt_id") == prompt_id:
                    for node_info in data.get("nodes", {}).values():
                        if node_info.get("state") == "running":
                            j.step = int(node_info.get("value", 0))
                            j.total_steps = int(node_info.get("max", 0))
                            _persist(j)
    except Exception:
        pass


async def _poll_history(prompt_id: str, interval: float = 5.0) -> dict:
    async with httpx.AsyncClient(timeout=30) as http:
        while True:
            resp = await http.get(f"{COMFY_BASE}/history/{prompt_id}")
            resp.raise_for_status()
            data = resp.json()
            if prompt_id in data:
                return data[prompt_id]
            await asyncio.sleep(interval)


async def _run_job_async(j: Job):
    graph, _ = build_graph(j.req)
    client_id = uuid.uuid4().hex
    async with httpx.AsyncClient(timeout=60) as http:
        resp = await http.post(f"{COMFY_BASE}/prompt", json={"prompt": graph, "client_id": client_id})
        resp.raise_for_status()
        body = resp.json()
        if body.get("node_errors"):
            raise RuntimeError(f"workflow 校验失败: {json.dumps(body['node_errors'], ensure_ascii=False)}")
        prompt_id = body["prompt_id"]
        j.comfy_prompt_id = prompt_id
        _persist(j)

    stop_tail = asyncio.Event()
    tail_task = asyncio.create_task(_tail_log_progress(j, j.req.steps, stop_tail))
    ws_task = asyncio.create_task(_track_ws_progress(j, client_id, prompt_id))
    try:
        # 2026-08-23: 之前完成判定只靠 websocket 的 execution_success 消息——这条连接在大分辨率/
        # 长帧数任务上复现过至少2次"消息收不到"（ComfyUI日志明明有"Prompt executed"，这边永远
        # 卡在running），把单队列堵死。改成轮询 /history 作为完成判定的唯一权威来源，
        # 跟websocket解耦：websocket只用来拿实时进度，它断线/丢消息不影响完成判定。
        hist = await _poll_history(prompt_id)
    finally:
        stop_tail.set()
        tail_task.cancel()
        ws_task.cancel()
        # 2026-08-23: 真根因找到了——asyncio.CancelledError 在py3.8+是BaseException的子类，
        # 不是Exception的子类，contextlib.suppress(Exception)根本压不住它。之前这里一旦
        # cancel()完再await，CancelledError会从这个finally里直接飞出去，绕过worker()的
        # except Exception，导致status永远停在running但finished_at已经被worker()的finally
        # 写上了——外面看起来就是"denoise明明跑完了、任务却卡死"，其实是这个而不是websocket
        # 断线（早前怀疑的方向）。这就是这两天反复出现的那个"卡死"bug的真身。
        with contextlib.suppress(asyncio.CancelledError):
            await tail_task
        with contextlib.suppress(asyncio.CancelledError):
            await ws_task

    status = hist.get("status") or {}
    if status.get("status_str") == "error":
        raise RuntimeError(f"ComfyUI execution_error: {_extract_error_summary(status.get('messages', []))}")

    # 视频文件在 outputs.<save节点id>.images[] 里，不是 "videos" key——ComfyUI 内部历史遗留命名
    save_out = hist["outputs"]["n_save"]["images"][0]
    fpath = OUTPUT_DIR / save_out["filename"]
    if not fpath.exists():
        raise RuntimeError(f"ComfyUI 报告完成但文件不存在: {fpath}")
    j.file = str(fpath)
    j.size_bytes = fpath.stat().st_size


def worker():
    STATE["model"] = "ready"  # ComfyUI 自己管加载，这里没有独立的"模型加载"阶段
    STATE["loaded_at"] = time.time()
    loop = asyncio.new_event_loop()
    while True:
        jid = Q.get()
        j = JOBS.get(jid)
        if j is None:
            continue
        if j.status == "cancelled":
            # 排在队里的时候被取消了——直接跳过，不要覆盖回running，也不要真的去跑
            Q.task_done()
            continue
        j.status = "running"
        j.started_at = time.time()
        _persist(j)
        try:
            loop.run_until_complete(_run_job_async(j))
            j.status = "done"
        except Exception as e:
            j.status = "failed"
            j.error = f"{type(e).__name__}: {e}"
            print(f"[job {jid}] failed: {j.error}")
        finally:
            j.finished_at = time.time()
            _persist(j)
            Q.task_done()


# ---------------- HTTP API —— 跟 gateway.py 现有契约保持一致，console 不用改一行 UI 代码 ----------------

app = FastAPI()


@app.on_event("startup")
def _start():
    INPUT_DIR.mkdir(parents=True, exist_ok=True)
    init_db()
    _load_jobs()
    threading.Thread(target=worker, daemon=True).start()


def _gpu_status() -> list:
    # 2026-08-24：给 console 首页显示实时显存/利用率——之前撞过好几次OOM，后来发现是别的
    # 工作负载（比如旧 minimax-h3 四卡worker）在偷偷抢显存，提交任务前完全看不出来。
    # 直接读 nvidia-smi，不缓存，每次 /health 请求都是当下的真实数字。
    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,memory.used,memory.total,utilization.gpu",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=3, check=True,
        ).stdout
    except Exception:
        return []
    gpus = []
    for line in out.strip().splitlines():
        parts = [x.strip() for x in line.split(",")]
        if len(parts) != 4:
            continue
        idx, used, total, util = parts
        gpus.append({"index": int(idx), "memory_used_mb": int(used),
                     "memory_total_mb": int(total), "utilization_pct": int(util)})
    return gpus


def _comfyui_alive() -> bool:
    # 2026-08-24：这之前是唯一的健康检查盲区——/health 只看 STATE["model"]（worker()线程
    # 启动时设一次"ready"，之后再也不更新），ComfyUI(8188)本体真的崩了，这层壳(server.py)
    # 没死的话 /health 会一直报 ready。实测过一次：8188 进程在一次OOM后整个退出，/health
    # 骗了将近20分钟才被人工发现，期间提交的任务全部秒失败(ConnectError)。改成每次 /health
    # 都真的探一下 8188 活不活，跟 start.sh 自己的就绪检测用同一个端点(/system_stats)。
    try:
        r = httpx.get(f"{COMFY_BASE}/system_stats", timeout=2.0)
        return r.status_code == 200
    except Exception:
        return False


@app.get("/health")
def health():
    running = sum(1 for j in JOBS.values() if j.status == "running")
    comfy_alive = _comfyui_alive()
    ready = 1 if (STATE["model"] == "ready" and comfy_alive) else 0
    model_state = STATE["model"] if comfy_alive else "comfyui_unreachable"
    backend = {
        "status": "ok" if ready else "loading",
        "model": model_state,
        "gpu": "0,1,2,3",
        "partition": "raylight-fsdp",  # 不再是路由依据，纯展示——4卡是一个集群，没有单卡分区概念了
        "offload": "fsdp",
        "queued": Q.qsize(),
        "running": running,
        "backend": BACKEND_URL,
    }
    return {
        "status": "ok" if ready else "unavailable",
        "ready": ready, "total": 1,
        "queued": Q.qsize(), "running": running,
        "backends": [backend],
        "gpus": _gpu_status(),
    }


@app.post("/v1/videos")
def create_video(req: GenReq):
    if req.turbo and req.steps > 8:
        raise HTTPException(400, "turbo 模式建议 steps<=8（当前实现turbo LoRA还没接，steps按原样生效）")
    jid = uuid.uuid4().hex[:12]
    partition = _partition_of(req)
    with LOCK:
        JOBS[jid] = Job(id=jid, status="queued", req=req, created_at=time.time(), partition=partition)
    _persist(JOBS[jid])
    Q.put(jid)
    return {"id": jid, "status": "queued", "queue_position": Q.qsize(),
            "eta_seconds": None, "poll": f"/v1/videos/{jid}", "backend": BACKEND_URL}


def _trigger_restart_async():
    # sleep 2 秒让当前这个HTTP响应先发出去，再让 start.sh 把这个进程自己也杀掉重启，
    # 不能同步等 start.sh 跑完（它会把当前进程杀掉，等不到返回）。
    subprocess.Popen(
        ["bash", "-c", "sleep 2 && bash /data/models/raylight/start.sh "
                       ">> /data/models/raylight/stop-restart.log 2>&1"],
        start_new_session=True,
    )


@app.post("/v1/videos/{jid}/cancel")
def cancel_video(jid: str):
    j = JOBS.get(jid)
    if j is None:
        raise HTTPException(404, "无此任务")
    if j.status == "queued":
        with LOCK:
            j.status = "cancelled"
            j.finished_at = time.time()
            j.error = "用户手动取消"
        _persist(j)
        return j.model_dump()
    if j.status == "running":
        # 2026-08-24实测确认：ComfyUI的/interrupt打断不了Raylight——它的实际计算跑在独立的
        # Ray子进程里（FSDP分布式），根本不检查ComfyUI主进程的打断标志位（打断前后GPU一直
        # 100%利用率，采样步数照样往前走，完全没停）。唯一能真正把GPU立刻收回来的办法是硬
        # 重启整个服务，跟"actor崩溃/卡死"的恢复手段一样——不假装能优雅停止。
        with LOCK:
            j.status = "cancelled"
            j.finished_at = time.time()
            j.error = "用户手动停止（硬重启GPU——Raylight的Ray子进程收不到ComfyUI的/interrupt信号，见notes.md）"
        _persist(j)
        _trigger_restart_async()
        return j.model_dump()
    raise HTTPException(409, f"任务当前状态是 {j.status}，无法取消/停止")


@app.post("/admin/reset-gpu")
def reset_gpu():
    # 通用重置——不依赖有没有具体任务在跑。已知上游bug：Raylight actor崩溃/OOM后不会自愈
    # （ensure_fresh_actors里needs_restart被无条件覆盖成False，见README），显存可能卡在
    # 接近满载但JOBS里根本没有running任务，这时候"停止"按钮不会出现（只在有running任务时才
    # 显示），需要一个不挂在具体任务上的兜底重置入口。
    with LOCK:
        for j in JOBS.values():
            if j.status == "running":
                j.status = "cancelled"
                j.finished_at = time.time()
                j.error = "用户手动重置GPU（硬重启，见 /v1/videos/{id}/cancel 同款说明）"
                _persist(j)
    _trigger_restart_async()
    return {"status": "restarting", "message": "已触发服务重启，约15-20秒后GPU恢复空闲"}


@app.get("/v1/videos/{jid}")
def get_video(jid: str):
    j = JOBS.get(jid)
    if j is None:
        raise HTTPException(404, "无此任务")
    out = j.model_dump()
    if j.status == "running" and j.started_at:
        out["elapsed_seconds"] = time.time() - j.started_at
    if j.status == "done":
        out["download"] = f"/v1/videos/{jid}/file"
    return out


@app.get("/v1/videos/{jid}/file")
def get_video_file(jid: str):
    j = JOBS.get(jid)
    if j is None:
        raise HTTPException(404, "无此任务")
    if j.status != "done" or not j.file:
        raise HTTPException(409, "任务还没完成")
    return FileResponse(j.file, media_type="video/mp4", filename=f"{jid}.mp4")


@app.get("/v1/videos")
def list_videos(limit: int = 50):
    jobs = sorted(JOBS.values(), key=lambda j: j.created_at, reverse=True)[:limit]
    return [{"id": j.id, "status": j.status, "created_at": j.created_at,
             "size_bytes": j.size_bytes, "error": j.error} for j in jobs]