"""sparkle OpenAI 兼容 server 入口:python -m mlx_streaming.server

模型路径参数与 sparkle chat 一致(默认值同样来自 mlx_streaming/config.py 的环境变量)。
SPARKLE_FAKE=1 时用 FakeBackend,免模型开发/测试。
"""
import argparse
import os

from mlx_streaming import config


def _build_parser():
    p = argparse.ArgumentParser(
        prog="sparkle-server",
        description="sparkle:OpenAI 兼容的流式 MoE 推理 server")
    p.add_argument("--host", default="127.0.0.1", help="监听地址(默认 127.0.0.1)")
    p.add_argument("--port", type=int, default=8317, help="监听端口(默认 8317)")
    p.add_argument("--model", default=config.model_path(), help="主模型路径(MLX 量化)")
    p.add_argument("--expert-dir", default=config.expert_dir(),
                   help="拆分后的 per-expert 目录")
    p.add_argument("--mtp-out", default=config.mtp_out(), help="MTP 权重文件")
    p.add_argument("--qn-config", default=config.qn_config(),
                   help="Qwen3-Next 配置 JSON")
    p.add_argument("-k", "--k", type=int, default=3, help="MTP 投机宽度(默认 3)")
    p.add_argument("-n", "--max-tokens", type=int, default=4096,
                   help="每轮最多生成的新 token 数(默认 4096)")
    p.add_argument("--expert-slots", type=int, default=32,
                   help="常驻专家池容量(默认 32,同时作为侧区行数默认)")
    p.add_argument("--spec-slots", type=int, default=None,
                   help="侧区行数 POOL_SPEC_SLOTS(默认跟随 --expert-slots)")
    return p


def _safe_prefill_chunk(args) -> int:
    """按专家池实际可安放容量算 prefill 步长上限，取 min(启动方给的值, 上限)。

    一次前向路由的唯一专家最多 chunk × top_k 个，真实区能安放的只有
    expert_slots − AUTOPIN 钉死数（钉死的永不驱逐）。超出会被 block.py 分流到
    host/fetch 路径：结果仍正确，但每层每 chunk 重读盘，prefill 慢几倍。
    读不到 top_k 时按 Qwen3-Next 的 10 兜底（偏保守，宁可 chunk 小）。
    """
    import json

    top_k = 10
    try:
        with open(args.qn_config, encoding="utf-8") as f:
            top_k = max(1, int(json.load(f).get("num_experts_per_tok", top_k)))
    except Exception:      # noqa: BLE001  配置读不到/字段缺失都退到保守默认
        pass
    slots = max(1, int(args.expert_slots))
    pinned = int(slots * config.autopin_budget_frac()) if config.autopin() else 0
    limit = max(1, (slots - pinned) // top_k)
    want = int(os.environ.get("PREFILL_CHUNK") or limit)
    if want > limit:
        print(f"[PREFILL_CHUNK] {want} → {limit}：{slots} 槽扣掉 AUTOPIN 钉死 {pinned} 后只能安放 "
              f"{slots - pinned} 个专家，chunk×top_k({top_k}) 不得超过它，否则 prefill 退到读盘慢路径。",
              flush=True)
    return min(want, limit)


def main(argv=None):
    args = _build_parser().parse_args(argv)

    # KV_QUANT 维持 Sparkle 原始配方(开):此前"不按规则回答"已证实是 tools 未渲染所致,
    # 与 KV 量化无关(且其通过质量门:token 一致率≥95%、logits cosine≥0.99)。
    # KV_QUANT=0 会使 prefill 慢 ~3 倍(fp16 KV 带宽),业务长提示下代价不可接受。
    # 显式 export KV_QUANT=0 可关闭做质量对照。

    # 前缀快照头部的**上限**(实际长度由引擎按 system+tools 的真实边界自动取,见
    # MLXBackend._fixed_prefix_len)。写死数字曾是主要的慢因:业务 system+4 工具实测
    # 3972 token,固定切 2048 只覆盖一半,每个新会话都要白 prefill 1900+ token
    # (prefill 约 30 tok/s,即多等一分钟)。这里只兜住极端情况下快照过大。
    os.environ.setdefault("PREFIX_SNAPSHOT_HEAD", "8192")

    # AUTOPIN 用一半槽位钉历史热专家(冷启动不慢热),另一半留给 demand 动态换入。
    os.environ.setdefault("AUTOPIN_BUDGET_FRAC", "0.5")

    # PREFILL_CHUNK 必须由槽位推导,不能让启动方拍数字:每个 chunk 的一次前向最多路由
    # chunk × top_k 个唯一专家,而真实区能安放的只有 expert_slots − AUTOPIN 钉死数。超了
    # 会被分流到 host/fetch 路径(正确但每层重读盘,慢几倍),槽位调小后写死的 chunk 就失效。
    os.environ["PREFILL_CHUNK"] = str(_safe_prefill_chunk(args))

    import uvicorn

    from mlx_streaming.server.app import create_app

    # timeout_graceful_shutdown:SIGTERM 后最多等 3 秒就强制关闭连接。
    # 外部客户端的轮询持有 keep-alive 长连接,不设超时 uvicorn 会永远
    # 卡在 "Waiting for connections to close",进程不退、内存不释放。
    uvicorn.run(create_app(args), host=args.host, port=args.port,
                timeout_graceful_shutdown=3)


if __name__ == "__main__":
    main()