"""FastAPI 应用装配:OpenAI 路由 + 引擎调参 API + 极简 admin 页。"""
from __future__ import annotations

import asyncio
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse

from mlx_streaming.server.admin import ADMIN_HTML
from mlx_streaming.server.openai import make_openai_router
from mlx_streaming.server.state import EngineManager, rss_bytes
from mlx_streaming.tui.backend import FakeBackend


def create_app(args, backend=None) -> FastAPI:
    """装配应用。

    args: argparse Namespace 风格对象(至少含 model/k/max_tokens/expert_slots,
    真后端还需要 expert_dir/mtp_out/qn_config/spec_slots)。
    backend: 显式注入的后端(测试用);None 时按 SPARKLE_FAKE=1 选 FakeBackend,
    否则 MLXBackend。
    """
    if backend is None:
        if os.environ.get("SPARKLE_FAKE") == "1":
            backend = FakeBackend()
        else:
            from mlx_streaming.tui.backend import MLXBackend

            backend = MLXBackend(args)
    mgr = EngineManager(args, backend)

    @asynccontextmanager
    async def lifespan(app):
        # 加载放线程且不等待:真模型以分钟计,不能阻塞 uvicorn 启动;
        # 就绪前 loaded=False,chat 请求返回 503。
        loop = asyncio.get_running_loop()
        app.state.load_future = loop.run_in_executor(None, mgr.load_initial)
        yield

    app = FastAPI(title="sparkle server", lifespan=lifespan)
    # 本机工具服务:放开 CORS,让浏览器能直连 /v1/models、/api/* 做探测
    from fastapi.middleware.cors import CORSMiddleware

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )
    app.state.mgr = mgr
    app.include_router(make_openai_router(mgr))

    @app.get("/api/stats")
    async def api_stats():
        return {"tok_per_s": mgr.last_tok_per_s,
                "rss_bytes": rss_bytes(),
                "expert_slots": mgr.expert_slots,
                "loaded": mgr.loaded,
                # 专家池「分不到槽被迫落 0 号槽」的累计数。>0 = 有前向拿错专家权重算过,
                # 输出不可信(调小 PREFILL_CHUNK 或调大 expert_slots)。健康值恒为 0。
                "unplaced_experts": mgr.unplaced_experts(),
                # 专家池命中/读盘指标:hit_rate 低就是慢的直接原因(每 miss 一个要读 ~3.3MB)
                "pool": mgr.pool_metrics(),
                "model": str(mgr.args.model)}

    @app.get("/api/engine/config")
    async def get_engine_config():
        return {"expert_slots": mgr.expert_slots,
                "k": mgr.k,
                "max_tokens": mgr.max_tokens}

    @app.post("/api/engine/config")
    async def set_engine_config(req: Request):
        """有多少改多少:k/max_tokens 即时生效;expert_slots 变化触发后台重建。"""
        body = await req.json()
        if "k" in body:
            mgr.k = int(body["k"])
        if "max_tokens" in body:
            mgr.max_tokens = int(body["max_tokens"])
        reloading = False
        if "expert_slots" in body:
            slots = int(body["expert_slots"])
            if slots != mgr.expert_slots:
                mgr.expert_slots = slots
                mgr.reloading = True  # 即刻对 chat 可见(503),不等后台任务起跑
                reloading = True

                async def _bg_rebuild():
                    async with mgr.lock:  # 等当前生成结束再换引擎
                        await asyncio.get_running_loop().run_in_executor(
                            None, mgr.rebuild)

                asyncio.create_task(_bg_rebuild())
        return {"ok": True, "reloading": reloading,
                "expert_slots": mgr.expert_slots,
                "k": mgr.k, "max_tokens": mgr.max_tokens}

    @app.get("/admin", response_class=HTMLResponse)
    async def admin_page():
        return ADMIN_HTML

    return app