"""引擎单例管理:加载状态、生成串行锁、运行期调参与后台重建。"""
from __future__ import annotations

import asyncio
import resource
import sys

from mlx_streaming.tui.backend import FakeBackend


def rss_bytes() -> int:
    """当前进程常驻内存峰值(字节)。macOS 的 ru_maxrss 直接是字节;Linux 是 KB。"""
    ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return int(ru if sys.platform == "darwin" else ru * 1024)


class EngineManager:
    """持有一个 ChatBackend 单例;引擎是单对话的,同一时刻只允许一轮生成。

    asyncio.Lock 做排队:第二个并发请求等锁而非直接失败。k / max_tokens / expert_slots
    是运行期可调参数,由 /api/engine/config 修改;k 与 max_tokens 即时生效(作为后续
    OpenAI 请求的缺省值),expert_slots 变化触发后台重建。
    """

    def __init__(self, args, backend):
        self.args = args  # argparse Namespace 风格,重建 MLXBackend 时复用
        self.backend = backend
        self.loaded = False
        self.reloading = False
        self.load_error: str | None = None
        self.lock = asyncio.Lock()
        self.last_tok_per_s = 0.0
        self.k = getattr(args, "k", 3)
        self.max_tokens = getattr(args, "max_tokens", 4096)
        self.expert_slots = getattr(args, "expert_slots", 32)

    def load_initial(self, on_status=lambda m: None) -> None:
        """启动加载(放线程里跑,真模型以分钟计);就绪前 loaded=False。"""
        try:
            self.backend.load(on_status)
            self.loaded = True
        except Exception as e:  # noqa: BLE001  加载失败不拖垮 server,状态暴露在 /api/stats
            self.load_error = f"{type(e).__name__}: {e}"

    def rebuild(self, on_status=lambda m: None) -> None:
        """后台重建引擎(fake 模式只换实例;真后端按更新后的 args 重新装配 + 预热)。"""
        self.reloading = True
        self.loaded = False
        try:
            self.args.expert_slots = self.expert_slots
            old = self.backend
            if isinstance(old, FakeBackend):
                self.backend = FakeBackend(reply=old.reply)
            else:
                from mlx_streaming.tui.backend import MLXBackend

                self.backend = MLXBackend(self.args)
            self.backend.load(on_status)
            self.loaded = True
        except Exception as e:  # noqa: BLE001
            self.load_error = f"{type(e).__name__}: {e}"
        finally:
            self.reloading = False

    def _resident_pool(self):
        """拿常驻专家池(fake 后端 / 非双源路径下为 None)。"""
        store = getattr(getattr(self.backend, "_model", None), "_expert_store", None)
        return getattr(store, "_resident", None)

    def unplaced_experts(self) -> int:
        """专家池累计「分不到槽落 0 号槽」数(见 ResidentExpertPool.unplaced)。

        取不到(fake 后端 / 非双源路径)返回 0。这是数值正确性的看门指标:非 0 就说明
        某些前向是拿错专家权重算的,回答不可信。
        """
        return int(getattr(self._resident_pool(), "unplaced", 0) or 0)

    def pool_metrics(self) -> dict:
        """专家池命中/读盘指标：定位「慢」到底是池不够大还是走了慢路径。

        decode 每步要取 层数×top_k 个专家，miss 一个就得读一份专家权重(8bit ≈3.3MB)，
        所以 hit_rate 直接决定速度。gpu_fallback 记的是本该零 host 往返的快路径退回读盘的次数。
        """
        rp = self._resident_pool()
        if rp is None:
            return {}
        hits, misses = int(getattr(rp, "hits", 0)), int(getattr(rp, "misses", 0))
        return {"hits": hits, "misses": misses,
                "hit_rate": round(hits / (hits + misses), 4) if hits + misses else 0.0,
                "gpu_fastpath": int(getattr(rp, "gpu_fastpath", 0)),
                "gpu_fallback": int(getattr(rp, "gpu_fallback", 0)),
                "prefetch_hits": int(getattr(rp, "prefetch_hits", 0)),
                "prefetch_loads": int(getattr(rp, "prefetch_loads", 0))}

    def apply_gen_params(self, req_max_tokens=None) -> int:
        """把当前调参写入后端 args(单飞前提下安全),返回本轮有效 max_tokens。"""
        eff = req_max_tokens or self.max_tokens
        if hasattr(self.backend, "args"):
            self.backend.args.k = self.k
            self.backend.args.max_tokens = eff
        return eff