"""MiniMax-H3 文生视频 HTTP API（异步任务式）。

为什么不是同步接口：模型加载 238s、单条视频 7 分钟（50 步 90 分钟），
同步 HTTP 必然超时。所以是「提交任务拿 id → 轮询状态 → 下载文件」。

为什么是串行单 worker：一条生成占 ~12.3G 显存 + ~32G CPU 内存，
一张 22G 卡同时只能跑一条。并发提交会排队，不会 OOM。
"""
import os, sys, time, uuid, threading, queue, traceback, io, sqlite3, urllib.request
from typing import Optional, Literal
from pathlib import Path

BASE = Path("/data/models/minimax-h3")
OUT = BASE / "outputs"; OUT.mkdir(exist_ok=True)
# 一个进程只跑一个分区。两套权重的 DiT、processor、Turbo LoRA 三样都不同，
# 同进程装两套要么重建 pipe（每次重读 16G 权重），要么常驻 77G 内存。
# 每卡一个进程本来就是这台机器唯一的并发方式（H3 不支持模型并行），顺手就把分区分开了。
PARTITION = os.environ.get("H3_PARTITION", "fl2va")
# 没上 GPU 的那部分权重存哪。cpu 是每进程一份 40G 匿名内存，并发被 125G 内存卡在 3 个；
# disk 让权重留在 safetensors 文件里（layers.py 里 offload() 就是 .to("meta")），
# 每次前向从文件读进显存，走 OS 页缓存——多个进程读同一个文件，物理内存里只有一份。
OFFLOAD = os.environ.get("H3_OFFLOAD", "cpu")
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "torch")
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(BASE / "models"))
os.environ.setdefault("MODELSCOPE_CACHE", str(BASE / ".mscache"))
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

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

def _rows(width: int, height: int, frames: int) -> float:
    """序列行数：潜帧数 x 宽 x 高 / 1024。真正决定显存/算力的量，帧数和分辨率是同一笔预算。"""
    latent_frames = (max(frames, 5) - 5) // 17 * 5 + 2
    return latent_frames * width * height / 1024

# 22G 卡的显存墙按 rows 算，不是单独卡 frames——实测边界（2026-08-20 long15s.sh）：
# 832x480x124f rows=14430 稳定基准；448x256x362f rows=11984、576x320x362f rows=19260 都能跑
# （后者峰值 14.5G，换入换出吃紧但没崩）；640x384x362f rows=25680 OOM（撞到 20.6G）。
# 20000 卡在 19260（能跑）和 25680（OOM）之间，留安全余量。
# 支持环境变量覆盖：只用于临时拉高某一个 worker 做边界探测实验，默认值（生产用的这道
# 安全线）不变——测完记得把起这个 worker 的进程重启回不带这个环境变量的状态。
ROWS_LIMIT = int(os.environ.get("H3_ROWS_LIMIT", 20000))

# ---------------- 任务模型 ----------------
class GenReq(BaseModel):
    prompt: str = Field(min_length=1, max_length=2000)
    # 上界取模型自己的默认档 1344x768。pipeline 本身不限分辨率，只把长宽向上取整到 32 的倍数，
    # 真正的墙是这张 22G 卡：序列行数按 h*w/1024 每潜帧算，832x480 是 390 行/帧 x 37 帧 = 14430 行，
    # 1344x768 是 1008 行/帧 = 37296 行，注意力是平方级，涨 2.6 倍行数就是涨约 6.7 倍算力。
    # 832x480 带三张参考图实测峰值 13.7G，往上哪一档撞墙没测过，超了就是这条任务 failed，不会拖死服务。
    width:  int = Field(832, ge=64,  le=1344)
    height: int = Field(480, ge=64,  le=768)
    # 上界 362=15s（协议规格上限），是否真能生成看下面 _check_rows——832x480 撞了还是会被拒。
    frames: int = Field(124, ge=8,   le=362)
    steps:  int = Field(4,   ge=1,   le=50)
    turbo:  bool = True
    seed:   int = Field(0, ge=0, le=2**31 - 1)
    flow_shift: Optional[float] = None
    # 首末帧条件（FL2VA）。给了首帧，这一段就从那张画面长出去——同一间屋子、同一个人。
    # 不给就是纯文生视频，屋子和人由模型自己编，跨条对不上。
    first_frame_url: Optional[str] = None
    last_frame_url:  Optional[str] = None
    # 参考图（Ref2VA）。首帧只管这一段从哪开始，参考图全程参与注意力，管的是「中间别乱改」。
    # 张数没有硬上限，代价是序列变长：一张按 ref_image_short_edge 缩放后折成 token，
    # 短边 480 时约 390 行，而目标视频本身才 14430 行。
    reference_image_urls: list[str] = []
    ref_image_short_edge: int = Field(480, ge=64, le=2048)

    @model_validator(mode="after")
    def _check_rows(self):
        rows = _rows(self.width, self.height, self.frames)
        if rows > ROWS_LIMIT:
            raise ValueError(f"分辨率×帧数超出显存安全上限（rows={rows:.0f} > {ROWS_LIMIT}）；"
                              f"降低 frames 或分辨率，例如 448x256/576x320 可到 362 帧，832x480 建议 ≤124 帧")
        return self

class Job(BaseModel):
    id: str
    status: Literal["queued", "loading", "running", "done", "failed"]
    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
    peak_vram_gb: Optional[float] = None

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

# ---------------- 任务记录落盘 ----------------
# 之前 JOBS 只在内存里，进程重启（比如 h3-console 的弹性调度器切分区）就把这张表清空——
# 视频文件在 outputs/ 没事，但 id→文件路径这层映射没了，前端/API 都查不到。改成每次状态变化
# 都往 sqlite 写一份，进程起来时先读回内存。四个 GPU 各一个进程，按 GPU id 分开的文件，
# 不共享一份——共享一份要处理多进程并发写 sqlite 的锁竞争，没必要：反正一个进程只会在同一张
# GPU 上起来，per-GPU 的文件天然不会有另一个进程来抢。
JOBS_DB = BASE / f"jobs_gpu{os.environ.get('CUDA_VISIBLE_DEVICES', 'x')}.sqlite3"
_JOBS_CON = sqlite3.connect(JOBS_DB, check_same_thread=False)
_JOBS_CON.execute(
    "CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, updated_at REAL NOT NULL, data TEXT NOT NULL)"
)
_JOBS_CON.commit()

def _persist(job: Job):
    with LOCK:
        _JOBS_CON.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()),
        )
        _JOBS_CON.commit()

def _load_jobs():
    """进程刚起来时把上次落盘的任务记录读回内存。in-flight 的状态没法恢复——Q 是空的，
    GPU 上下文也是全新的，没有任何东西能"接着跑"。落盘时正好是 queued/running 的一律
    改判为 failed 并写清原因，不然它会卡成一条永远显示"运行中"但没有 worker 会再碰的
    死记录，比直接查不到这个 id 更容易让人误判成"卡住了，再等等"。"""
    n = 0
    for (data,) in _JOBS_CON.execute("SELECT data FROM jobs"):
        try:
            j = Job.model_validate_json(data)
        except Exception:
            continue  # 单条记录损坏不该挡住整个服务启动
        if j.status in ("queued", "running"):
            j.status = "failed"
            j.error = "进程重启导致这条任务的执行状态丢失（GPU 上下文和排队状态不落盘，无法恢复）"
            j.finished_at = j.finished_at or time.time()
        JOBS[j.id] = j
        n += 1
    if n:
        print(f"[api] 从磁盘恢复了 {n} 条历史任务记录（{JOBS_DB.name}）", flush=True)

_load_jobs()

# ---------------- 推理 worker ----------------
# DiT、processor、Turbo LoRA 三样成套，跨分区混用会对不上层。
# ref2v 的 Turbo LoRA 只有 v0.1，fl2v 那个是 v1.0 768p——这一项两边成色不同。
PARTITIONS = {
    "fl2va": {"dit": "minimax-h3-fl2va-nf4.safetensors",
              "processor": "FL2VA/processor/",
              "lora": "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"},
    # ref2v 的 bf16 版实测出 124 帧噪点：同尺寸同 seed，挂上它是噪点、摘掉就出画面，
    # 换 flow_shift 无效（见 docs/H3 的二分表）。改用 comfyui 版试另一套键名。
    "ref2va": {"dit": "minimax-h3-ref2va-nf4.safetensors",
               "processor": "Ref2VA/processor/",
               "lora": "minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors"},
}

def _turbo_lora_config():
    from diffsynth.pipelines.minimax_h3_audio_video import ModelConfig
    return ModelConfig(model_id="lightx2v/Minimax-h3-Turbo",
                       origin_file_pattern=PARTITIONS[PARTITION]["lora"])

def _build_pipe():
    from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig
    part = PARTITIONS[PARTITION]
    DT = torch.float32          # fp16 在 sm_75 上出全黑视频，见 README
    # 常驻显存里留给权重多少 GB，剩下的全是激活。默认公式是按 832x480 调的（留 14G 给激活，实测峰值 13.7G）。
    # 分辨率往上走就得调低它：1280x720 在默认值下要 22.3G、卡只有 21.49G，差约 1G 撞 OOM；
    # 调低后权重多从内存流几趟，慢一点，换来激活空间。用 H3_VRAM_LIMIT 覆盖。
    free = torch.cuda.mem_get_info("cuda")[0] / 1024**3
    vram_limit = float(os.environ.get("H3_VRAM_LIMIT", max(4.0, free - 14.0)))
    offload_dt = "disk" if OFFLOAD == "disk" else DT
    vc = dict(offload_dtype=offload_dt, offload_device=OFFLOAD,
              onload_dtype=offload_dt, onload_device=OFFLOAD,
              preparing_dtype=DT, preparing_device="cuda",
              computation_dtype=DT, computation_device="cuda")
    NF4 = "DiffSynth-Studio/MiniMax-H3-NF4"
    pipe = MiniMaxH3Pipeline.from_pretrained(
        torch_dtype=DT, device="cuda",
        model_configs=[ModelConfig(model_id=NF4, origin_file_pattern=f, **vc) for f in
            [part["dit"], "minimax-h3-text-encoder-nf4.safetensors",
             "video_vae_nf4.safetensors", "audio_vae_nf4.safetensors"]],
        processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3",
                                     origin_file_pattern=part["processor"]),
        vram_limit=vram_limit)
    pipe.load_lora(pipe.dit, _turbo_lora_config())
    return pipe

def _fetch_image(url):
    from PIL import Image
    with urllib.request.urlopen(url, timeout=60) as resp:
        return Image.open(io.BytesIO(resp.read())).convert("RGB")

def _keyframes(req):
    """取首末帧图，缩到出片尺寸。都没给就返回 (None, None)，pipeline 退回纯文生视频。"""
    from PIL import Image
    images, indices = [], []
    for index, url in ((0, req.first_frame_url), (-1, req.last_frame_url)):
        if url is None:
            continue
        images.append(_fetch_image(url).resize((req.width, req.height), Image.LANCZOS))
        indices.append(index)
    return (images, indices) if images else (None, None)

def _references(req):
    """参考图按请求顺序取下来。尺寸交给 pipeline 按 ref_image_short_edge 归一，这里不动。"""
    if not req.reference_image_urls:
        return None
    return [{"type": "image", "image": _fetch_image(url)} for url in req.reference_image_urls]

def worker():
    from diffsynth.utils.data.audio_video import write_video_audio
    import numpy as np
    t0 = time.time()
    pipe = _build_pipe()
    STATE["model"] = "ready"; STATE["loaded_at"] = time.time()
    print(f"[api] 模型就绪，耗时 {time.time()-t0:.0f}s", flush=True)

    while True:
        jid = Q.get()
        j = JOBS[jid]
        try:
            j.status = "running"; j.started_at = time.time()
            j.total_steps = j.req.steps
            _persist(j)
            torch.cuda.reset_peak_memory_stats()
            fs = j.req.flow_shift if j.req.flow_shift is not None else (6.0 if j.req.turbo else 12.0)
            if not j.req.turbo:
                pipe.clear_lora()
            keyframes, keyframe_indices = _keyframes(j.req)
            video, audio = pipe(prompt=j.req.prompt, height=j.req.height, width=j.req.width,
                                num_frames=j.req.frames, num_inference_steps=j.req.steps,
                                seed=j.req.seed, flow_shift=fs,
                                keyframes=keyframes, keyframe_indices=keyframe_indices,
                                references=_references(j.req),
                                ref_image_short_edge=j.req.ref_image_short_edge)
            # 产物必须验过才算成功——全黑视频也能写出文件（踩过）
            a = np.stack([np.asarray(f.convert("RGB"), dtype=np.float32)
                          for f in video[:: max(1, len(video)//8)]])
            if a.std() < 1.0:
                raise RuntimeError(f"生成结果近乎纯色 (std={a.std():.3f})，判为失败")
            if torch.is_tensor(audio) and not torch.isfinite(audio).all():
                raise RuntimeError("音频含非有限值，判为失败")
            out = OUT / f"api_{jid}.mp4"
            write_video_audio(video=video, audio=audio, output_path=str(out),
                              fps=24, audio_sample_rate=32000)
            j.file = str(out); j.size_bytes = out.stat().st_size
            j.peak_vram_gb = round(torch.cuda.max_memory_allocated()/1024**3, 2)
            j.status = "done"
        except Exception as e:
            j.status = "failed"; j.error = f"{type(e).__name__}: {e}"
            traceback.print_exc()
        finally:
            j.finished_at = time.time()
            _persist(j)
            if not j.req.turbo:      # 下一个任务可能要 turbo，把 LoRA 挂回去
                try:
                    pipe.load_lora(pipe.dit, _turbo_lora_config())
                except Exception:
                    traceback.print_exc()
            Q.task_done()

# ---------------- HTTP ----------------
app = FastAPI(title="MiniMax-H3 Video API", version="1.0")

@app.on_event("startup")
def _start():
    threading.Thread(target=worker, daemon=True).start()

@app.get("/health")
def health():
    return {"status": "ok" if STATE["model"] == "ready" else "loading",
            "model": STATE["model"], "gpu": STATE["gpu"], "partition": PARTITION,
            "offload": OFFLOAD, "queued": Q.qsize(),
            "running": sum(1 for j in JOBS.values() if j.status == "running")}

@app.post("/v1/videos", status_code=202)
def create(req: GenReq):
    if req.turbo and req.steps > 8:
        raise HTTPException(400, "turbo 模式步数应 ≤8（LoRA 是 4 步蒸馏的）；要多步请传 turbo=false")
    jid = uuid.uuid4().hex[:12]
    with LOCK:
        JOBS[jid] = Job(id=jid, status="queued", req=req, created_at=time.time())
    _persist(JOBS[jid])  # LOCK 外面调用 —— _persist 自己会拿 LOCK，Lock 不可重入，嵌套会死锁
    Q.put(jid)
    # 耗时的自变量是序列行数（潜帧数 x 宽 x 高 / 1024），不是分辨率——分辨率和帧数是同一笔预算。
    # 基准点 832x480x124f = 14430 行，4 步实测 692 秒 = 4 x 132 + 163
    #（163 是文本编码 + 参考图编码 + VAE 解码，它也跟帧数走，所以一起按行数缩放）。
    # 回代：1280x736x56f（15640 行）估 749 / 实测 736；832x480x22f（2730 行）估 131 / 实测 96。
    # 注意力是平方级，行数越大这个线性估计偏得越小。
    rows = _rows(req.width, req.height, req.frames)
    est = round((req.steps * 132 + 163) * rows / 14430)
    return {"id": jid, "status": "queued", "queue_position": Q.qsize(),
            "eta_seconds": est, "poll": f"/v1/videos/{jid}"}

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

@app.get("/v1/videos/{jid}/file")
def download(jid: str):
    j = JOBS.get(jid)
    if not j: raise HTTPException(404, "无此任务")
    if j.status != "done": raise HTTPException(409, f"任务尚未完成（当前 {j.status}）")
    return FileResponse(j.file, media_type="video/mp4", filename=f"{jid}.mp4")

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