"""AUTOPIN：专家热度持久化 + 启动预热钉死。

两部分：
- **热度计数与持久化**：搭已有的物化点零同步计频——block host 路径的 flat、
  acquire_gpu 的 LFU piggyback(全命中)与 miss 回退 flat；dual decode/verify 无 Python
  物化点，由 C++ g_real 的 LFU freq 在落盘时按增量差分合并。节流写
  {EXPERT_DIR}/pool_usage.json(同目录临时文件 + rename 防半写；读不到/损坏静默从零)。
- **启动预热钉死**：build_streaming_model 末尾按历史热度把每层 top-N 热专家
  (N = floor(cap_for × AUTOPIN_BUDGET_FRAC))经 blob 并行读预填进常驻池并 pin 住
  (不参与任何驱逐)，消灭冷启动慢热。

覆盖率口径：prefill/host 路径与 GPU remap miss 回退全计；非 dual 的 decode 全命中
快路径搭 LFU piggyback 计(inds 已物化，零额外同步)；dual decode/verify 靠 C++ freq
兜底(仅 lfu 策略计频)。EVICT_POLICY≠lfu 时两条 GPU 快路径均不计——绝不为计数给
decode 热路径新增 GPU→host 同步。
"""
import atexit
import json
import os
import tempfile
from collections import Counter
from typing import Dict, List

from mlx_streaming import config

_USAGE_NAME = "pool_usage.json"


def usage_path() -> str:
    """热度文件路径：AUTOPIN_USAGE_FILE 覆盖，缺省 {EXPERT_DIR}/pool_usage.json。"""
    p = config.autopin_usage_file()
    return p or os.path.join(config.expert_dir(), _USAGE_NAME)


def load_usage(path: str) -> "Dict[int, Counter]":
    """读热度文件；不存在/损坏/结构非法一律静默返回空(从零开始)。"""
    try:
        with open(path) as f:
            raw = json.load(f)
        return {int(l): Counter({int(e): int(c) for e, c in d.items()})
                for l, d in raw.items()}
    except Exception:
        return {}


def save_usage(path: str, counts: "Dict[int, Counter]") -> None:
    """原子写：先写同目录临时文件再 rename，防半写损坏。失败静默(不阻断推理)。"""
    try:
        d = os.path.dirname(path) or "."
        fd, tmp = tempfile.mkstemp(dir=d, prefix=".pool_usage.", suffix=".tmp")
        with os.fdopen(fd, "w") as f:
            json.dump({str(l): {str(e): c for e, c in cnt.items()}
                       for l, cnt in counts.items()}, f)
        os.replace(tmp, path)
    except Exception:
        pass


class HeatCounter:
    """进程内热度累计器：note() 增量计频，节流 + atexit 落盘。

    前向边界用「MoE 层号回绕」近似(与 VirtualPool.begin_forward 同判据)：MoE 块按层号
    递增被调，出现 layer <= 上次即新前向。路径在创建时快照，避免进程退出时 env 已变
    导致 atexit 写错位置。
    """

    def __init__(self, path: str):
        self.path = path
        self.counts: "Dict[int, Counter]" = {}
        self._last_layer = -1
        self._forwards_since_save = 0
        # C++ 累计频次的已合并基线(real_freq_dump 是进程累计值，按差分合并防重复计)
        self._cpp_baseline: "Dict[int, Dict[int, int]]" = {}

    def note(self, layer: int, ids: List[int]) -> None:
        layer = int(layer)
        if layer <= self._last_layer:
            self._forwards_since_save += 1
        self._last_layer = layer
        c = self.counts.get(layer)
        if c is None:
            c = self.counts[layer] = Counter()
        for e in ids:
            c[int(e)] += 1
        if self._forwards_since_save >= config.autopin_save_every():
            self.flush()

    def tick(self) -> None:
        """显式前向边界(dual decode 无 Python 计频点，由 begin_forward 驱动周期落盘)。"""
        self._forwards_since_save += 1
        if self._forwards_since_save >= config.autopin_save_every():
            self.flush()

    def merge_cpp_freq(self) -> None:
        """把 dual decode 的 C++ LFU 累计频次按增量差分合并进来。"""
        try:
            import mlx_streaming.native_moe_ext as _N
            flat = _N.real_freq_dump()
        except Exception:
            return
        cur: "Dict[int, Dict[int, int]]" = {}
        for i in range(0, len(flat), 3):
            l, e, c = int(flat[i]), int(flat[i + 1]), int(flat[i + 2])
            cur.setdefault(l, {})[e] = c
        for l, em in cur.items():
            base = self._cpp_baseline.get(l, {})
            cnt = self.counts.setdefault(l, Counter())
            for e, c in em.items():
                b = base.get(e, 0)
                cnt[e] += c - b if c >= b else c   # real_reset 后 c<b：以 c 为新基线
        self._cpp_baseline = cur

    def flush(self) -> None:
        self.merge_cpp_freq()
        save_usage(self.path, self.counts)
        self._forwards_since_save = 0


_counter: "HeatCounter | None" = None


def _get_counter() -> HeatCounter:
    global _counter
    if _counter is None:
        _counter = HeatCounter(usage_path())
        atexit.register(_counter.flush)
    return _counter


def note(layer: int, ids: List[int]) -> None:
    """路由热度计数入口(热路径调用)：AUTOPIN=0 立即返回，零行为。"""
    if not config.autopin():
        return
    _get_counter().note(layer, ids)


def tick() -> None:
    """前向边界入口(virtual_pool begin_forward 调用)：AUTOPIN=0 立即返回。"""
    if not config.autopin():
        return
    _get_counter().tick()


def counter() -> "HeatCounter | None":
    """测试/诊断用：当前计数器(未触发过计数则为 None)。"""
    return _counter


def warm_start_pin(store) -> "dict":
    """启动预热钉死：按 usage 文件每层取 top-N 热专家，经 store.pin_batch 预填进池并 pin 住。

    N = floor(该层 cap_for × AUTOPIN_BUDGET_FRAC)。pin 加载由 pin_batch 收口：有 blob_loader
    走并行 pread(生产唯一数据源)，无则退回逐专家 per-expert 文件(兼容旧部署)。
    usage 缺失/为空 → 返回零摘要(后台计数仍在累计，供下次启动用)；单层失败不阻断整体。
    """
    summary = {"layers": 0, "pinned": 0, "skipped_layers": 0, "seconds": 0.0}
    usage = load_usage(usage_path())
    if not usage:
        return summary
    import time
    t0 = time.perf_counter()
    frac = config.autopin_budget_frac()
    for layer in sorted(usage):
        n = int(store.cap_for(layer) * frac)
        if n <= 0:
            continue
        top = [e for e, _ in usage[layer].most_common(n)]
        try:
            summary["pinned"] += store.pin_batch(layer, top)
            summary["layers"] += 1
        except Exception:
            summary["skipped_layers"] += 1
    summary["seconds"] = round(time.perf_counter() - t0, 2)
    return summary