"""四个 H3 实例前面的调度层。

api.py 是单 worker 串行，一个进程独占一张卡，所以并发的唯一办法是往不同端口投。
这一层把「投哪个端口」收掉：对外只有一个地址，接口和 api.py 逐字相同，
调用方原来怎么写现在还怎么写，只把端口换成 8000。

自己不存任何状态。任务 id 归哪个后端，问一遍就知道——四个后端一次扫描的代价可以忽略，
换来的是这个进程随时能重启，不丢任何东西（api.py 的 JOBS 本来也在内存里，
调度层再存一份 id→后端 的映射，就是第二个会丢的地方）。
"""
import os, json, threading, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse

BACKENDS = [f"http://127.0.0.1:{port.strip()}" for port in
            os.environ.get("H3_BACKENDS", "8010,8011,8012,8013").split(",")]
POOL = ThreadPoolExecutor(max_workers=16)
# 最闲的那个后端已经压了这么多条时，再来的提交直接 429 顶回去。
# 没有这个上限的话队列是无界的：投 1000 条全部返回 202，最后一条两天后才开始跑，
# 调用方却收到「已接受」。宁可当场说满了。默认 4 条 ≈ 45 分钟积压。
MAX_QUEUE = int(os.environ.get("H3_MAX_QUEUE", "4"))
# 选后端是「读一遍各家队列深度，再往最闲的那个投」。这中间不能插进另一个提交，
# 否则同时来的请求会读到同一份快照、一起挑中同一个后端，全堆在一张卡上。
# 锁只圈住选 + 投这一段（约 30 毫秒），任务本身要跑十几分钟，序列化的代价可以忽略。
PICK = threading.Lock()


def fetch(url, data=None, timeout=30):
    req = urllib.request.Request(url, data=data,
                                 headers={"Content-Type": "application/json"} if data else {})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.load(resp)


def health_all():
    def one(backend):
        try:
            return dict(fetch(backend + "/health", timeout=5), backend=backend)
        except Exception as e:
            return {"backend": backend, "status": "down", "error": f"{type(e).__name__}: {e}"}
    return list(POOL.map(one, BACKENDS))


def find_owner(jid):
    """任务归谁，挨个问。后端对陌生 id 返回 404，urlopen 抛异常，接住当没有。"""
    def one(backend):
        try:
            return backend, fetch(f"{backend}/v1/videos/{jid}", timeout=5)
        except Exception:
            return backend, None
    for backend, job in POOL.map(one, BACKENDS):
        if job is not None:
            return backend, job
    raise HTTPException(404, "无此任务")


app = FastAPI(title="MiniMax-H3 Gateway", version="1.0")


@app.get("/health")
def health():
    backends = health_all()
    ready = [b for b in backends if b.get("status") == "ok"]
    return {"status": "ok" if ready else "unavailable",
            "ready": len(ready), "total": len(backends),
            "queued": sum(b.get("queued", 0) for b in ready),
            "running": sum(b.get("running", 0) for b in ready),
            "backends": backends}


# 同步 def，让 FastAPI 把它丢进线程池。写成 async def 会让下面这些阻塞的 HTTP 调用
# 卡住事件循环——那样确实也能歪打正着地把提交排成队，但任何一个后端卡住，
# 5 秒超时就会连轮询一起堵死。要串行就用锁明说。
@app.post("/v1/videos", status_code=202)
def create(req: dict):
    # 给了首末帧就只能走 fl2va；其余都走 ref2va——ref2va 不给参考图就退回纯文生视频，
    # 所以「没有首末帧」这一类不需要再细分。
    want = "fl2va" if (req.get("first_frame_url") or req.get("last_frame_url")) else "ref2va"
    with PICK:
        candidates = [b for b in health_all()
                      if b.get("status") == "ok" and b.get("partition") == want]
        if not candidates:
            raise HTTPException(503, f"没有就绪的 {want} 实例")
        # 排队最短的那个。queued 是还在队列里的，running 是正在跑的那条，两个都要算。
        pick = min(candidates, key=lambda b: b["queued"] + b["running"])
        ahead = pick["queued"] + pick["running"]
        if ahead >= MAX_QUEUE:
            raise HTTPException(429, f"队列满：{len(candidates)} 个 {want} 实例每个都压着 "
                                     f"{ahead} 条（上限 {MAX_QUEUE}）。等一会儿再投，"
                                     f"或调大 H3_MAX_QUEUE")
        try:
            job = fetch(pick["backend"] + "/v1/videos", data=json.dumps(req).encode())
        except urllib.error.HTTPError as e:
            raise HTTPException(e.code, e.read().decode())
    # eta_seconds 是这一条自己的执行时间，不含排队。前面还压着 ahead 条，
    # 真正的等待还要加上它们——它们规格可能不同，这里不替调用方瞎乘。
    return dict(job, backend=pick["backend"], queue_ahead=ahead, poll=f"/v1/videos/{job['id']}")


@app.get("/v1/videos/{jid}")
def get(jid: str):
    backend, job = find_owner(jid)
    return dict(job, backend=backend)


@app.get("/v1/videos/{jid}/file")
def download(jid: str):
    backend, job = find_owner(jid)
    if job["status"] != "done":
        raise HTTPException(409, f"任务尚未完成（当前 {job['status']}）")
    resp = urllib.request.urlopen(f"{backend}/v1/videos/{jid}/file", timeout=300)
    return StreamingResponse(resp, media_type="video/mp4",
                             headers={"Content-Disposition": f'attachment; filename="{jid}.mp4"'})


@app.get("/v1/videos")
def listjobs(limit: int = 50):
    def one(backend):
        try:
            return [dict(j, backend=backend)
                    for j in fetch(f"{backend}/v1/videos?limit={limit}", timeout=5)["jobs"]]
        except Exception:
            return []
    jobs = [job for part in POOL.map(one, BACKENDS) for job in part]
    jobs.sort(key=lambda j: j["created_at"], reverse=True)
    return {"jobs": jobs[:limit]}