#!/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 FastAPI, HTTPException, UploadFile, File, Request
from fastapi.responses import StreamingResponse, FileResponse
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，够单张参考图/首尾帧，防止占满磁盘
ENABLE_SUPERVISOR = os.environ.get("H3_ENABLE_SUPERVISOR") == "1"

# 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)


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
            )
        """)


init_db()


# ---------------- 资产上传 ----------------
# api.py 的 GenReq 只认 URL 字段（first_frame_url / reference_image_urls），内部用
# urllib.urlopen 拉取。所以上传接口的活是「收文件 -> 落盘 -> 吐一个后端能拉到的 URL」，
# 不改 api.py 一行代码。
@app.post("/v1/assets")
async def upload_asset(file: UploadFile = File(...)):
    ext = Path(file.filename or "").suffix.lower()
    if ext not in (".jpg", ".jpeg", ".png", ".webp"):
        raise HTTPException(400, f"不支持的文件类型 {ext}，只收图片(jpg/png/webp)")
    body = await file.read()
    if len(body) > MAX_UPLOAD_BYTES:
        raise HTTPException(413, f"文件超过 {MAX_UPLOAD_BYTES // 1024 // 1024}MB 上限")
    aid = uuid.uuid4().hex[:16]
    path = ASSETS / f"{aid}{ext}"
    path.write_bytes(body)
    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}"}


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("/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:
    # 和 gateway.py 里 `want = "fl2va" if (...) else "ref2va"` 保持完全一致的判断，
    # 不要各写一套、两边判断分岔的话 console 统计出来的 pending 计数就和 gateway 实际路由对不上。
    return "fl2va" if (req.get("first_frame_url") or req.get("last_frame_url")) else "ref2va"


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"]
    # gateway.py/api.py 原样带过来的 `poll`/`download` 字段用的是 gateway 内部 id（不是 console
    # 自己的 id），照抄会指向一个 console 自己的 sqlite 里根本查不到的路径，直接 404——这两个
    # 字段本来就该跟 `id` 一样，统一改成用 console 的 id 重新拼。
    if "poll" in view:
        view["poll"] = f"/v1/videos/{row['id']}"
    if "download" in view:
        view["download"] = f"/v1/videos/{row['id']}/file"
    return view


async def _try_submit(cid: str):
    """把一条本地排队的任务尝试提交给 gateway。给 create_video 首次调用和 reconciler_loop 复用。"""
    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  # 任务不存在，或已经提交过
    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))


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"):
            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),
        )


@app.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)


@app.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)


@app.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 一行代码。
    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")
        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")}
    return StreamingResponse(stream(), status_code=upstream.status_code, headers=passthrough)


@app.get("/v1/videos")
async def list_videos(limit: int = 50):
    # console 自己的 sqlite 是任务列表的唯一真相来源（本地排队中的任务，gateway 根本不知道它存在，
    # 没法从 gateway 反查出来）。实时进度由后台 reconciler_loop 持续刷新进 last_known_json，
    # 这里直接读，不在请求路径里再打一遍 gateway，避免列表接口被拖慢。
    with db() as con:
        rows = con.execute("SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,)).fetchall()
    return {"jobs": [_row_view(r) for r in rows]}


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


@app.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},
        ],
    }


# ---------------- 弹性 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 决定显示哪个视图），这两条路由只是让
# 浏览器地址栏里出现符合预期的路径，不是两个不同的页面文件。必须放在 StaticFiles 挂载之前，
# 不然 `/task.html/xxx` 会被当成一个不存在的静态文件直接 404。
@app.get("/task.html/{tid}")
async def task_page(tid: str):
    return FileResponse(str(BASE / "static" / "index.html"))


@app.get("/newtask.html")
async def newtask_page():
    return FileResponse(str(BASE / "static" / "index.html"))


app.mount("/", StaticFiles(directory=str(BASE / "static"), html=True), name="static")