#!/usr/bin/env python3
"""H3 Console：给 gateway.py 加一层 Web UI + 上传 + 任务持久化。

不重做调度——gateway.py 已经做了选后端/分区路由/限流，这里只反代它的
/v1/videos* 接口，加三样它没有的东西：文件上传（返回可喂给 api.py 的 URL）、
SQLite 任务影子记录（gateway.py 的 JOBS 在后端进程内存里，进程重启就丢）、
静态前端。

弹性 partition 调度器（kill+restart 某个 GPU 后端切 fl2va/ref2va）只在
H3_ENABLE_SUPERVISOR=1 时启用——这一层直接操作服务器上的真实进程，本地
Mac 开发环境没有意义，见 DESIGN.md 「四点五」。
"""
import os, re, sqlite3, time, uuid, json, asyncio, subprocess
from pathlib import Path
from typing import Optional

import httpx
from fastapi import APIRouter, FastAPI, HTTPException, UploadFile, File, Request
from fastapi.responses import StreamingResponse, FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image

BASE = Path(__file__).resolve().parent
ASSETS = BASE / "assets"
ASSETS.mkdir(exist_ok=True)
DB_PATH = BASE / "console.sqlite3"

GATEWAY_URL = os.environ.get("H3_GATEWAY_URL", "http://192.168.110.67:8000")
CONSOLE_PUBLIC_BASE = os.environ.get("H3_CONSOLE_PUBLIC_BASE")  # 供后端 worker 回读上传图片用；见 README
MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20MB，够单张参考图/首尾帧，防止占满磁盘
MAX_VIDEO_UPLOAD_BYTES = 100 * 1024 * 1024  # 100MB，参考视频（2-15秒）比图片大得多
MAX_AUDIO_UPLOAD_BYTES = 30 * 1024 * 1024  # 30MB，独立音频参考
ENABLE_SUPERVISOR = os.environ.get("H3_ENABLE_SUPERVISOR") == "1"
# 部署在共享域名的子路径下时设这个（比如挂在 qwen.wudaima.com/console 而不是独立域名），
# 见 README「子路径部署」。留空就是原来的行为，挂在根路径，本地/局域网直连不受影响。
BASE_PATH = os.environ.get("H3_CONSOLE_BASE_PATH", "").rstrip("/")

# GPU/端口映射，对应服务器上 start-all.sh 写死的分配。仅 supervisor 用得到。
BACKEND_TOPOLOGY = [
    {"gpu": 1, "port": 8010},
    {"gpu": 2, "port": 8011},
    {"gpu": 0, "port": 8012},
    {"gpu": 3, "port": 8013},
]
H3_ROOT = Path("/data/models/minimax-h3")  # 仅服务器上存在，本地开发不会用到

app = FastAPI(title="H3 Console")
client = httpx.AsyncClient(timeout=30.0)
# 所有 /v1/* 接口挂在这个 router 上而不是直接挂在 app 上，是为了能整体套一层 BASE_PATH
# 前缀（app.include_router(router, prefix=BASE_PATH)，见文件末尾）——BASE_PATH 留空时
# 前缀是空字符串，行为和直接挂在 app 上完全一样。
router = APIRouter()


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


def init_db():
    with db() as con:
        con.execute("""
            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                gateway_id TEXT,
                partition TEXT NOT NULL,
                backend TEXT,
                status TEXT NOT NULL,
                req_json TEXT NOT NULL,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL,
                last_known_json TEXT
            )
        """)
        # 老库升级：CREATE TABLE IF NOT EXISTS 不会给已存在的表加列，只能单独 ALTER。
        # 第二次启动起这行必然报"列已存在"，吞掉就行，不影响其它列。
        try:
            con.execute("ALTER TABLE jobs ADD COLUMN deleted_at REAL")
        except sqlite3.OperationalError:
            pass


init_db()


# ---------------- 资产上传 ----------------
# api.py 的 GenReq 只认 URL 字段（first_frame_url / reference_image_urls），内部用
# urllib.urlopen 拉取。所以上传接口的活是「收文件 -> 落盘 -> 吐一个后端能拉到的 URL」，
# 不改 api.py 一行代码。
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv")
AUDIO_EXTS = (".mp3", ".wav", ".m4a", ".ogg", ".flac")


@router.post("/v1/assets")
async def upload_asset(file: UploadFile = File(...)):
    ext = Path(file.filename or "").suffix.lower()
    # 2026-08-23：原来只收图片，现在参考视频/独立音频参考也走这个接口——按扩展名分流到
    # 各自的大小上限，视频/音频没有 PIL 那样现成的解码校验，先不做，坏文件会在 ComfyUI
    # LoadVideo/LoadAudio 那一步报错，属于可接受的失败模式。
    if ext in IMAGE_EXTS:
        kind, max_bytes = "image", MAX_UPLOAD_BYTES
    elif ext in VIDEO_EXTS:
        kind, max_bytes = "video", MAX_VIDEO_UPLOAD_BYTES
    elif ext in AUDIO_EXTS:
        kind, max_bytes = "audio", MAX_AUDIO_UPLOAD_BYTES
    else:
        raise HTTPException(400, f"不支持的文件类型 {ext}，只收图片(jpg/png/webp)/"
                                  f"视频(mp4/mov/webm/mkv)/音频(mp3/wav/m4a/ogg/flac)")
    body = await file.read()
    if len(body) > max_bytes:
        raise HTTPException(413, f"文件超过 {max_bytes // 1024 // 1024}MB 上限")
    aid = uuid.uuid4().hex[:16]
    path = ASSETS / f"{aid}{ext}"
    path.write_bytes(body)
    if kind == "image":
        try:
            Image.open(path).verify()  # 校验确实是能解码的图片，不是随便什么二进制
        except Exception:
            path.unlink(missing_ok=True)
            raise HTTPException(400, "文件不是有效图片")
    return {"id": aid, "url": f"{_public_base()}/assets/{aid}{ext}", "kind": kind}


def _public_base() -> str:
    if CONSOLE_PUBLIC_BASE:
        return CONSOLE_PUBLIC_BASE.rstrip("/")
    # 本地开发时没配这个环境变量会直接报错提示——见 README「已知限制」：
    # 上传图片的 URL 必须是 api.py 所在服务器能反向拉取到的地址，本地 Mac 默认猜不出这个地址。
    raise HTTPException(
        500, "未设置 H3_CONSOLE_PUBLIC_BASE：上传接口需要知道自己对外可访问的地址，"
             "这个地址会被服务器上的 api.py worker 反向拉取，见 README。")


app.mount(f"{BASE_PATH}/assets", StaticFiles(directory=str(ASSETS)), name="assets")


# ---------------- 生成任务：console 自己是排队的主人，gateway.py 只是"当前有空的执行者" ----------------
# 关键点：gateway.py 对没有就绪后端的 partition 直接 503（`没有就绪的 fl2va 实例`），不排队。
# 如果 console 只是把这个 503 原样透传给用户，"有几个任务排队就用几张卡跑 fl2va"就无从谈起——
# 请求在 503 那一步就死了，根本没机会被弹性调度器看到、更别说触发 flip。
# 所以 console 自己维护身份（console 生成的 id）和排队状态：提交时先落 sqlite 记一条 status=queued，
# 尝试提交给 gateway；提交不上（503/429）就留在本地队列里，交给下面的 reconciler_loop 每隔几秒重试，
# 弹性调度器看到的 pending 计数正是这些「还没提交上去」的本地排队任务。
def _partition_of(req: dict) -> str:
    # 2026-08-23 修复：原来是二选一（没有first/last frame就当ref2va），会把纯文生视频
    # （reference_image_urls也是空的）误标成"ref2va"——跟前端 rerunJob() 自己的 kind 判断
    # （index.html 里早就正确区分 text/fl2va/ref2va 三种）对不上，在 task.html 上看起来
    # "标着ref2va却没有参考图"，容易让人怀疑参考图是不是丢了/没显示出来，其实只是这个分类
    # 少了纯文生视频这一档。改成跟前端一致的三分类。
    if req.get("first_frame_url") or req.get("last_frame_url"):
        return "fl2va"
    # 2026-08-23：新增参考视频/独立音频参考字段时漏了这里，跟 raylight/server.py 自己的
    # _partition_of 对不上——那边路由图用的是它自己那份判断（功能不受影响，只是这里的
    # 展示标签会显示错），保持两份判断逻辑一致。
    if req.get("reference_image_urls") or req.get("reference_video_urls") or req.get("reference_audio_urls"):
        return "ref2va"
    return "text"


def _row_view(row) -> dict:
    """把一行 sqlite 记录还原成对外的任务视图，用 console 自己的 id，不暴露 gateway 内部 id。"""
    if row["last_known_json"]:
        view = json.loads(row["last_known_json"])
    else:
        view = {}
    view["id"] = row["id"]
    view["status"] = row["status"]
    view["partition"] = row["partition"]
    view["backend"] = row["backend"]
    view.setdefault("req", json.loads(row["req_json"]))
    view["created_at"] = row["created_at"]
    view["deleted_at"] = row["deleted_at"]
    # gateway.py/api.py 原样带过来的 `poll`/`download` 字段用的是 gateway 内部 id（不是 console
    # 自己的 id），照抄会指向一个 console 自己的 sqlite 里根本查不到的路径，直接 404——这两个
    # 字段本来就该跟 `id` 一样，统一改成用 console 的 id 重新拼。
    if "poll" in view:
        view["poll"] = f"{BASE_PATH}/v1/videos/{row['id']}"
    if "download" in view:
        view["download"] = f"{BASE_PATH}/v1/videos/{row['id']}/file"
    return view


# create_video 首次提交、get_video 轮询时顺手提交、reconciler_loop 每 5 秒扫描，三条路都会在
# gateway_id 还是 NULL 的窗口里调 _try_submit(cid)。三个协程共享同一个事件循环，`await
# client.post(...)` 让出控制权的那段时间里，另外两条路完全可能也读到「还没提交过」，
# 于是同一个 cid 被 POST 给 gateway 两次——两张卡各烧一遍 GPU，console 只记得住最后一次
# UPDATE 写回的 gateway_id，另一条产物是真实存在但没人知道、也没人能停的孤儿任务。
# 2026-08-20 实测踩到过：一个 15s/832×480 的探测任务真的在两张卡上各 OOM 了一次。
# 修法：进程内占坑（不是数据库锁——server.py 单进程单事件循环，内存里一个 set 就够，
# 两行判断中间没有 await，天然原子，不会被其它协程插进来）。
_SUBMITTING: set[str] = set()


async def _try_submit(cid: str):
    """把一条本地排队的任务尝试提交给 gateway。给 create_video 首次调用和 reconciler_loop 复用。"""
    if cid in _SUBMITTING:
        return  # 已经有一次提交在飞行中，等它写完 gateway_id 再说，不要并发再发一次
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    if row is None or row["gateway_id"] is not None:
        return  # 任务不存在，或已经提交过
    _SUBMITTING.add(cid)
    try:
        req = json.loads(row["req_json"])
        try:
            resp = await client.post(f"{GATEWAY_URL}/v1/videos", json=req)
        except httpx.HTTPError:
            return  # gateway 暂时连不上，下一轮再试
        now = time.time()
        if resp.status_code < 400:
            job = resp.json()
            gateway_id = job["id"]
            job["id"] = cid  # 对外只暴露 console 自己的 id
            with db() as con:
                con.execute(
                    "UPDATE jobs SET gateway_id=?, backend=?, status=?, updated_at=?, last_known_json=? WHERE id=?",
                    (gateway_id, job.get("backend"), job.get("status", "queued"),
                     now, json.dumps(job, ensure_ascii=False), cid),
                )
        elif resp.status_code in (503, 429):
            return  # 没有就绪后端 / 队列满，留在本地队列，等 flip 或别的任务跑完再试
        else:
            # 真正的客户端错误（比如 422 参数越界、400 turbo+steps>8），重试也不会成功，直接判失败
            err_view = {"id": cid, "status": "failed", "error": resp.text,
                        "req": req, "partition": row["partition"], "created_at": row["created_at"]}
            with db() as con:
                con.execute("UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                            (now, json.dumps(err_view, ensure_ascii=False), cid))
    finally:
        _SUBMITTING.discard(cid)


async def _refresh_from_gateway(cid: str, gateway_id: str):
    """已经提交给 gateway 的任务，去问一下最新状态，刷新本地记录。"""
    try:
        resp = await client.get(f"{GATEWAY_URL}/v1/videos/{gateway_id}")
    except httpx.HTTPError:
        return
    now = time.time()
    if resp.status_code == 404:
        # gateway 完全不认识这个 gateway_id——通常是所在后端进程被重启（比如手动 flip 分区
        # 或崩溃），内存态清空，这条任务已经不会再有任何进展。不能像别的非 200 一样悄悄跳过，
        # 那样 status 会永远卡在 queued/running，界面看着像在跑，实际已经死了、永远不会再更新。
        with db() as con:
            row = con.execute("SELECT last_known_json FROM jobs WHERE id=?", (cid,)).fetchone()
        if row is None:
            return
        view = json.loads(row["last_known_json"]) if row["last_known_json"] else {"id": cid}
        if view.get("status") in ("done", "failed", "cancelled"):
            return  # 已经是终态，不要覆盖
        view["status"] = "failed"
        view["error"] = "任务在执行途中丢失：所在后端进程被重启，内存态被清空，不会再有进展"
        with db() as con:
            con.execute(
                "UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                (now, json.dumps(view, ensure_ascii=False), cid),
            )
        return
    if resp.status_code != 200:
        return  # 其它非 200（后端一时不可达等）多半是暂时的，保留原状态，下一轮再试
    job = resp.json()
    job["id"] = cid
    with db() as con:
        con.execute(
            "UPDATE jobs SET status=?, backend=?, updated_at=?, last_known_json=? WHERE id=?",
            (job.get("status"), job.get("backend"), now, json.dumps(job, ensure_ascii=False), cid),
        )


@router.post("/v1/videos", status_code=202)
async def create_video(req: dict):
    partition = _partition_of(req)
    cid = uuid.uuid4().hex[:12]
    now = time.time()
    synth = {"id": cid, "status": "queued", "req": req, "partition": partition, "created_at": now}
    with db() as con:
        con.execute(
            "INSERT INTO jobs (id, gateway_id, partition, backend, status, req_json, created_at, updated_at, last_known_json) "
            "VALUES (?,NULL,?,NULL,'queued',?,?,?,?)",
            (cid, partition, json.dumps(req, ensure_ascii=False), now, now, json.dumps(synth, ensure_ascii=False)),
        )
    await _try_submit(cid)  # 尽力立即提交一次；提交不上就留给 reconciler_loop
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    return _row_view(row)


@router.post("/v1/videos/{jid}/cancel")
async def cancel_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["status"] in ("done", "failed", "cancelled"):
        raise HTTPException(409, f"任务当前状态是 {row['status']}，已经是终态，取消不了")
    now = time.time()
    if row["gateway_id"]:
        # 已经提交给后端了，得让后端也知道要取消（后端只允许取消它自己还在排队的任务，
        # 如果已经开始跑了会返回409——这里直接把后端的错误透传出去，不假装取消成功了）
        try:
            resp = await client.post(f"{GATEWAY_URL}/v1/videos/{row['gateway_id']}/cancel")
        except httpx.HTTPError:
            raise HTTPException(502, "取消请求发不到后端，稍后再试")
        if resp.status_code >= 400:
            raise HTTPException(resp.status_code, resp.text)
        job = resp.json()
        job["id"] = jid
        with db() as con:
            con.execute(
                "UPDATE jobs SET status='cancelled', updated_at=?, last_known_json=? WHERE id=?",
                (now, json.dumps(job, ensure_ascii=False), jid),
            )
    else:
        # 还没提交给后端（本地排队中），直接在 console 自己这边标掉，reconciler_loop 只挑
        # status='queued' 的行重试提交，改成 cancelled 之后它就不会再去试着提交这条了
        with db() as con:
            con.execute("UPDATE jobs SET status='cancelled', updated_at=? WHERE id=?", (now, jid))
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_view(row)


@router.get("/v1/videos/{jid}")
async def get_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["gateway_id"]:
        await _refresh_from_gateway(jid, row["gateway_id"])
    else:
        await _try_submit(jid)  # 还在本地排队时，顺手趁轮询这一下也试着提交
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_view(row)


@router.get("/v1/videos/{jid}/file")
async def download_video(jid: str, request: Request):
    with db() as con:
        row = con.execute(
            "SELECT gateway_id, backend, status, last_known_json FROM jobs WHERE id=?", (jid,)
        ).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if not row["gateway_id"] or not row["backend"]:
        raise HTTPException(409, "任务还在本地排队，尚未分配到可用后端")

    # 不经过 gateway.py 转发：它自己那层 `StreamingResponse(urlopen(...))` 不认 Range 请求，
    # 探测 Range 一律回 200 而不是 206，还不带 Content-Length/Accept-Ranges——`<video>`
    # 标签播放前先探一次 Range 拿时长/可否拖动，探不到就一直卡在 readyState=0 播不出来
    # （下载能成功是因为普通 GET 不需要这些，只有网页里的播放器会被这个坑绊住）。
    # api.py 自己那层是 Starlette `FileResponse`，原生支持 Range；console 既然已经从
    # gateway 的响应里存了 `backend`（真实拥有这条任务的那个 GPU 后端），直接绕过 gateway
    # 这层转发、把 Range 头透传给 api.py，把它的状态码/头原样带回去，问题在这一层就解决，
    # 不用碰 gateway.py/api.py 一行代码。
    # 下载/播放到本地的文件名——不带这个的话浏览器只能从 URL 最后一段猜，我们的 URL 都是
    # 以 "/file" 结尾（没有扩展名），猜出来清一色叫 "file.mp4"，多个任务的视频存到同一个
    # 下载目录会互相覆盖。用 console 自己的任务 id 拼一个每条任务都不同的文件名。
    download_filename = f"h3-{jid}.mp4"
    content_disposition = f'attachment; filename="{download_filename}"'

    upstream_url = f"{row['backend']}/v1/videos/{row['gateway_id']}/file"
    headers = {}
    if "range" in request.headers:
        headers["range"] = request.headers["range"]
    try:
        req = client.build_request("GET", upstream_url, headers=headers)
        upstream = await client.send(req, stream=True)
    except httpx.HTTPError:
        upstream = None

    if upstream is None or upstream.status_code >= 400:
        if upstream is not None:
            await upstream.aclose()
        # 后端进程重启会把 api.py 内存里的 JOBS 字典清空（这是它本来的设计，任务状态
        # 不落盘），导致它忘了这条任务——但视频文件本身在磁盘上没被动过。console 和
        # api.py 跑在同一台机器上，任务完成时已经把 `file` 路径存进了自己的 sqlite，
        # 后端忘了就直接读文件兜底，不用等它想起来（它也不会想起来，重启就是清空）。
        cached = json.loads(row["last_known_json"]) if row["last_known_json"] else {}
        file_path = cached.get("file")
        if file_path and os.path.exists(file_path):
            return FileResponse(file_path, media_type="video/mp4",
                                 headers={"Content-Disposition": content_disposition})
        raise HTTPException(404, "后端进程重启后忘了这条任务的记录，磁盘上也没找到对应文件")

    async def stream():
        async for chunk in upstream.aiter_bytes():
            yield chunk
        await upstream.aclose()

    passthrough = {k: v for k, v in upstream.headers.items()
                   if k.lower() in ("content-length", "content-range", "accept-ranges", "content-type")}
    passthrough["content-disposition"] = content_disposition
    return StreamingResponse(stream(), status_code=upstream.status_code, headers=passthrough)


@router.get("/v1/videos")
async def list_videos(limit: int = 50, trash: bool = False):
    # console 自己的 sqlite 是任务列表的唯一真相来源（本地排队中的任务，gateway 根本不知道它存在，
    # 没法从 gateway 反查出来）。实时进度由后台 reconciler_loop 持续刷新进 last_known_json，
    # 这里直接读，不在请求路径里再打一遍 gateway，避免列表接口被拖慢。
    # 软删除：deleted_at 非空的任务默认不出现在列表里，?trash=true 反过来只看这些——
    # 数据还在同一张表里，删除/恢复就是 UPDATE 一个时间戳，不用挪表、不用真删文件。
    where = "deleted_at IS NOT NULL" if trash else "deleted_at IS NULL"
    order = "deleted_at DESC" if trash else "created_at DESC"
    with db() as con:
        rows = con.execute(f"SELECT * FROM jobs WHERE {where} ORDER BY {order} LIMIT ?", (limit,)).fetchall()
    return {"jobs": [_row_view(r) for r in rows]}


@router.delete("/v1/videos/{jid}")
async def delete_video(jid: str):
    """软删除：只挪出主列表，不碰 gateway/api.py 上的任务本身（可能还在跑），不删磁盘文件。"""
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=? WHERE id=? AND deleted_at IS NULL", (time.time(), jid))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务，或已经在回收站里")
    return {"id": jid, "deleted": True}


@router.post("/v1/videos/{jid}/restore")
async def restore_video(jid: str):
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=NULL WHERE id=? AND deleted_at IS NOT NULL", (jid,))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务，或不在回收站里")
    return {"id": jid, "deleted": False}


@router.get("/v1/backends")
async def backends():
    resp = await client.get(f"{GATEWAY_URL}/health")
    resp.raise_for_status()
    return resp.json()


@router.get("/v1/models")
def models():
    # 静态信息，对应服务器上 api.py 里的 PARTITIONS 常量。只读展示，不做运行时切换。
    return {
        "partitions": {
            "fl2va": {"dit": "minimax-h3-fl2va-nf4.safetensors",
                      "lora": "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors",
                      "desc": "首末帧条件生成"},
            "ref2va": {"dit": "minimax-h3-ref2va-nf4.safetensors",
                       "lora": "minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors",
                       "desc": "文生视频 / 多参考图生成"},
        },
        "presets": [
            {"label": "标准档 832×480×124f(5.17s)", "width": 832, "height": 480, "frames": 124},
            {"label": "长时长档 448×256×362f(15s)", "width": 448, "height": 256, "frames": 362},
            {"label": "高分辨率档 1344×768×64f(2.7s)", "width": 1344, "height": 768, "frames": 64},
        ],
    }


# ---------------- 弹性 partition 调度器（仅服务器部署时启用） ----------------
def _pending_counts() -> dict:
    with db() as con:
        rows = con.execute(
            "SELECT partition, COUNT(*) c FROM jobs WHERE status IN ('queued','running') GROUP BY partition"
        ).fetchall()
    return {r["partition"]: r["c"] for r in rows}


async def _flip_backend(gpu: int, port: int, new_partition: str):
    # 和今天手动测试时用的手法一致：按端口查 PID、杀、用现成的 api-run.sh 重启到新分区。
    # 不新写进程管理代码，复用已经在跑的脚本。
    # 注意：ss 的输出里 pid 不是独立的空格分隔 token，是嵌在 `users:(("python",pid=NNNN,fd=14))`
    # 里面的——按空格 split 后拿 `tok.startswith("pid=")` 永远不会命中，得用正则从整段里抠出来
    # （手动测试时用的是 `grep -o "pid=[0-9]*"`，这里必须做等价的事，不能简化成 split）。
    pid_out = subprocess.run(["ss", "-lptn", f"sport = :{port}"], capture_output=True, text=True).stdout
    for pid in re.findall(r"pid=(\d+)", pid_out):
        subprocess.run(["kill", pid])
    subprocess.Popen(
        ["setsid", "nohup", "./api-run.sh", str(gpu), str(port), new_partition],
        cwd=str(H3_ROOT), stdout=open(H3_ROOT / "logs" / f"api-{port}.log", "a"),
        stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, start_new_session=True,
    )


async def supervisor_loop():
    idle_since: dict[int, float] = {}
    # 一个 flip 从 kill 到新进程报 ready 要 13-20 秒（api.py 冷启动实测），远长于 10 秒的 tick 间隔。
    # 如果不记「已经在切换中」，接下来两三轮 tick 会看到 fl2va_ready 还是 0（新进程没起来），
    # 误判成需求仍未满足，再多切几张卡——今天调试时真实撞到过：一个 pending=1 的请求把三张卡都切了
    # 过去（而且当时 kill 那步还有个 bug 没真的杀掉旧进程，新进程绑端口失败悄悄退出，见下）。
    flipping: dict[int, tuple[str, float]] = {}  # port -> (目标 partition, 发起时间)
    while True:
        await asyncio.sleep(10)
        try:
            health = (await client.get(f"{GATEWAY_URL}/health")).json()["backends"]
            pending = _pending_counts()
            now = time.time()

            for port in list(flipping):
                target, started = flipping[port]
                b = next((x for x in health if x["backend"].endswith(f":{port}")), None)
                if b and b.get("partition") == target and b.get("status") == "ok":
                    flipping.pop(port)  # 切换完成，新分区已经在 health 里报 ready 了
                elif now - started > 60:
                    flipping.pop(port)  # 60 秒还没起来，当作这次尝试失败，让下一轮按实际情况重新决定

            fl2va_ready = [b for b in health if b.get("partition") == "fl2va" and b.get("status") == "ok"]
            ref2va_ready = [b for b in health if b.get("partition") == "ref2va" and b.get("status") == "ok"]
            fl2va_inflight = sum(1 for t, _ in flipping.values() if t == "fl2va")
            ref2va_inflight = sum(1 for t, _ in flipping.values() if t == "ref2va")

            def idle_backends(group):
                out = []
                for b in group:
                    port = int(b["backend"].rsplit(":", 1)[1])
                    if port in flipping:
                        continue  # 已经在切换中，不能再选它
                    if b.get("queued", 0) == 0 and b.get("running", 0) == 0:
                        idle_since.setdefault(port, now)
                        if now - idle_since[port] >= 30:
                            out.append(port)
                    else:
                        idle_since.pop(port, None)
                return out

            fl2va_pending = pending.get("fl2va", 0)
            ref2va_pending = pending.get("ref2va", 0)
            print(f"[supervisor] tick: pending={pending} "
                  f"fl2va={len(fl2va_ready)}(+{fl2va_inflight} 切换中) "
                  f"ref2va={len(ref2va_ready)}(+{ref2va_inflight} 切换中)", flush=True)

            if fl2va_pending > len(fl2va_ready) + fl2va_inflight:
                candidates = idle_backends(ref2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} ref2va -> fl2va", flush=True)
                    flipping[port] = ("fl2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "fl2va")
            elif ref2va_pending > len(ref2va_ready) + ref2va_inflight:
                candidates = idle_backends(fl2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} fl2va -> ref2va", flush=True)
                    flipping[port] = ("ref2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "ref2va")
        except Exception as e:
            print(f"[supervisor] tick failed: {e}", flush=True)


# ---------------- 常驻对账循环（本地/服务器都跑，跟 GPU 无关） ----------------
# 两件事：本地排队中的任务（gateway_id 还是空）隔几秒重试提交一次；已提交的任务刷新最新状态。
# 这是任务列表能实时更新、以及「没后端时先排队而不是直接报错」这两件事真正生效的地方。
async def reconciler_loop():
    while True:
        await asyncio.sleep(5)
        try:
            with db() as con:
                pending = con.execute(
                    "SELECT id FROM jobs WHERE gateway_id IS NULL AND status='queued'"
                ).fetchall()
                in_flight = con.execute(
                    "SELECT id, gateway_id FROM jobs WHERE gateway_id IS NOT NULL AND status IN ('queued','running')"
                ).fetchall()
            for r in pending:
                await _try_submit(r["id"])
            for r in in_flight:
                await _refresh_from_gateway(r["id"], r["gateway_id"])
        except Exception as e:
            print(f"[reconciler] tick failed: {e}", flush=True)


@app.on_event("startup")
async def _startup():
    asyncio.create_task(reconciler_loop())
    if ENABLE_SUPERVISOR:
        asyncio.create_task(supervisor_loop())
        print("[console] elastic partition supervisor enabled", flush=True)
    else:
        print("[console] supervisor disabled (H3_ENABLE_SUPERVISOR!=1) — fl2va/ref2va allocation is whatever the backends currently report", flush=True)


# 任务详情/新建各自有一个真实的 URL，方便复制粘贴分享、刷新不丢状态——都是同一个单页应用
# （static/index.html 自己那份 JS 读 location.pathname 决定显示哪个视图），这三条路由只是让
# 浏览器地址栏里出现符合预期的路径，不是三个不同的页面文件。
#
# 页面不再直接 FileResponse 原文件，而是读进来做一次字符串替换：把 __H3_BASE_PATH__ 占位符
# 换成 BASE_PATH。部署在共享域名子路径下时（比如 qwen.wudaima.com/console），前端 JS 里所有
# fetch/history.pushState 用的绝对路径都得带上这个前缀，否则请求会打到域名根路径而不是我们
# 这个子路径——BASE_PATH 留空时替换成空字符串，跟原来的行为完全一样。
def render_index() -> HTMLResponse:
    html = (BASE / "static" / "index.html").read_text(encoding="utf-8")
    return HTMLResponse(html.replace("__H3_BASE_PATH__", BASE_PATH))


@router.get("/")
async def index_page():
    return render_index()


@router.get("/task.html/{tid}")
async def task_page(tid: str):
    return render_index()


@router.get("/newtask.html")
async def newtask_page():
    return render_index()


# BASE_PATH 留空时这个前缀是空字符串，router 上的路径原样挂到 app 根上，行为跟以前
# 直接用 @app.xxx 一模一样；非空时（比如 "/console"）整体套一层前缀，一次性生效，
# 不用每个路由自己拼。
app.include_router(router, prefix=BASE_PATH)