"""CPython native MoE backend 的 Python 包装层。"""
import importlib
import json
import os
import time
from collections import OrderedDict
from functools import lru_cache

import mlx.core as mx
import numpy as np

from mlx_streaming import config


_EXPERT_STAGE_CACHE: "OrderedDict[tuple, tuple[mx.array, mx.array, mx.array]]" = OrderedDict()
_BUNDLE_STAGE_CACHE: "OrderedDict[tuple, tuple[mx.array, ...]]" = OrderedDict()
_SLOT_POOLS: dict[tuple, "NativeComputeSlotPool"] = {}
_STAGE_STATS = {
    "expert_hits": 0,
    "expert_misses": 0,
    "bundle_hits": 0,
    "bundle_misses": 0,
    "evictions": 0,
    "can_checks": 0,
    "can_true": 0,
    "can_false": 0,
    "calls": 0,
    "route_sync_s": 0.0,
    "stage_s": 0.0,
    "enqueue_s": 0.0,
}


class NativeComputeSlotPool:
    """按层维护 compute-buffer 常驻 slot，hit 时复用 [cap,...] MLX arrays。"""

    def __init__(self, compute_dir: str, layer: int, hidden: int, inter: int,
                 group: int, bits: int, num_experts: int, cap: int):
        self.compute_dir = compute_dir
        self.layer = int(layer)
        self.hidden = int(hidden)
        self.inter = int(inter)
        self.group = int(group)
        self.bits = int(bits)
        self.num_experts = int(num_experts)
        self.cap = int(cap)
        self.slot_of: "OrderedDict[int, int]" = OrderedDict()
        self.expert_of: list[int | None] = [None] * self.cap
        self.slot_arrays = [self._empty_slot() for _ in range(self.cap)]
        self.pool_arrays: tuple[mx.array, ...] | None = None
        self.dirty = True
        self.hits = 0
        self.misses = 0
        self.evictions = 0
        self.rebuilds = 0

    def _empty_slot(self) -> tuple[mx.array, ...]:
        gu_words = self.hidden * self.bits // 32
        gu_groups = self.hidden // self.group
        down_words = self.inter * self.bits // 32
        down_groups = self.inter // self.group
        return (
            mx.zeros((self.inter, gu_words), dtype=mx.uint32),
            mx.zeros((self.inter, gu_groups), dtype=mx.uint16),
            mx.zeros((self.inter, gu_groups), dtype=mx.uint16),
            mx.zeros((self.inter, gu_words), dtype=mx.uint32),
            mx.zeros((self.inter, gu_groups), dtype=mx.uint16),
            mx.zeros((self.inter, gu_groups), dtype=mx.uint16),
            mx.zeros((self.hidden, down_words), dtype=mx.uint32),
            mx.zeros((self.hidden, down_groups), dtype=mx.uint16),
            mx.zeros((self.hidden, down_groups), dtype=mx.uint16),
        )

    def _load_slot(self, expert: int) -> tuple[mx.array, ...]:
        gate = _stage_one_projection_uncached(
            self.compute_dir, self.layer, "gate_proj", expert,
            self.inter, self.hidden, self.group, self.bits, self.num_experts)
        up = _stage_one_projection_uncached(
            self.compute_dir, self.layer, "up_proj", expert,
            self.inter, self.hidden, self.group, self.bits, self.num_experts)
        down = _stage_one_projection_uncached(
            self.compute_dir, self.layer, "down_proj", expert,
            self.hidden, self.inter, self.group, self.bits, self.num_experts)
        return (*gate, *up, *down)

    def _assign_slot(self, expert: int) -> int:
        cached = self.slot_of.get(expert)
        if cached is not None:
            self.slot_of.move_to_end(expert)
            self.hits += 1
            return cached
        self.misses += 1
        if len(self.slot_of) < self.cap:
            slot = next(i for i, e in enumerate(self.expert_of) if e is None)
        else:
            old_expert, slot = self.slot_of.popitem(last=False)
            self.expert_of[slot] = None
            self.evictions += 1
        self.slot_arrays[slot] = self._load_slot(expert)
        self.expert_of[slot] = expert
        self.slot_of[expert] = slot
        self.dirty = True
        return slot

    def _rebuild_pool_arrays(self) -> None:
        self.pool_arrays = tuple(
            mx.stack([slot[i] for slot in self.slot_arrays], axis=0)
            for i in range(9)
        )
        self.rebuilds += 1
        self.dirty = False

    def acquire(self, expert_ids: list[int]) -> tuple[list[int], tuple[mx.array, ...]]:
        local = [self._assign_slot(int(e)) for e in expert_ids]
        if self.pool_arrays is None or self.dirty:
            self._rebuild_pool_arrays()
        assert self.pool_arrays is not None
        return local, self.pool_arrays

    def stats(self) -> dict:
        return {
            "cap": self.cap,
            "entries": len(self.slot_of),
            "hits": self.hits,
            "misses": self.misses,
            "evictions": self.evictions,
            "rebuilds": self.rebuilds,
        }


def _compute_dir() -> str:
    return config.compute_buffer_dir() or os.path.join(
        config.expert_dir(default=""), "compute_buffers")


@lru_cache(maxsize=1)
def _load_ext():
    try:
        return importlib.import_module("mlx_streaming.native_moe_ext")
    except Exception:
        return None


@lru_cache(maxsize=256)
def _has_layer_buffers(compute_dir: str, layer: int) -> bool:
    for proj in ("gate_proj", "up_proj", "down_proj"):
        for name in ("weight", "scales", "biases"):
            path = os.path.join(compute_dir, f"layer{layer:02d}.{proj}.{name}.bin")
            if not os.path.exists(path):
                return False
    return True


@lru_cache(maxsize=1024)
def _buffer_index_matches(compute_dir: str, layer: int, proj: str,
                          out_dim: int, in_dim: int, group: int, bits: int,
                          num_experts: int) -> bool:
    path = os.path.join(compute_dir, f"layer{layer:02d}.{proj}.index.json")
    if not os.path.exists(path):
        return False
    try:
        with open(path) as f:
            index = json.load(f)
    except Exception:
        return False
    tensors = index.get("tensors", {})
    weight_shape = tensors.get("weight", {}).get("shape_per_expert")
    scale_shape = tensors.get("scales", {}).get("shape_per_expert")
    bias_shape = tensors.get("biases", {}).get("shape_per_expert")
    expected_weight = [out_dim, in_dim * bits // 32]
    expected_grouped = [out_dim, in_dim // group]
    return (
        int(index.get("num_experts", -1)) == int(num_experts)
        and weight_shape == expected_weight
        and scale_shape == expected_grouped
        and bias_shape == expected_grouped
    )


@lru_cache(maxsize=512)
def _has_matching_compute_buffers(compute_dir: str, layer: int,
                                  hidden: int, inter: int, group: int,
                                  bits: int, num_experts: int) -> bool:
    if not _has_layer_buffers(compute_dir, layer):
        return False
    # index.json 是防呆线：bit/group 或目录错配时宁可回退 MLX，不冒险读错 mmap shape。
    return (
        _buffer_index_matches(
            compute_dir, layer, "gate_proj", inter, hidden, group, bits, num_experts)
        and _buffer_index_matches(
            compute_dir, layer, "up_proj", inter, hidden, group, bits, num_experts)
        and _buffer_index_matches(
            compute_dir, layer, "down_proj", hidden, inter, group, bits, num_experts)
    )


@lru_cache(maxsize=256)
def _projection_memmaps(compute_dir: str, layer: int, proj: str,
                        out_dim: int, in_dim: int, group: int, bits: int,
                        num_experts: int):
    words = in_dim * bits // 32
    groups = in_dim // group
    base = os.path.join(compute_dir, f"layer{layer:02d}.{proj}")
    weight = np.memmap(
        base + ".weight.bin",
        dtype=np.uint32,
        mode="r",
        shape=(num_experts, out_dim, words),
    )
    scales = np.memmap(
        base + ".scales.bin",
        dtype=np.uint16,
        mode="r",
        shape=(num_experts, out_dim, groups),
    )
    biases = np.memmap(
        base + ".biases.bin",
        dtype=np.uint16,
        mode="r",
        shape=(num_experts, out_dim, groups),
    )
    return weight, scales, biases


def _stage_projection(compute_dir: str, layer: int, proj: str,
                      expert_ids: list[int], out_dim: int, in_dim: int,
                      group: int, bits: int, num_experts: int):
    if config.native_moe_stage_cache():
        staged = [
            _stage_one_projection(
                compute_dir, layer, proj, int(e), out_dim, in_dim, group, bits, num_experts)
            for e in expert_ids
        ]
        return tuple(mx.stack([entry[i] for entry in staged], axis=0) for i in range(3))
    return _stage_projection_uncached(
        compute_dir, layer, proj, expert_ids, out_dim, in_dim, group, bits, num_experts)


def _stage_projection_uncached(compute_dir: str, layer: int, proj: str,
                              expert_ids: list[int], out_dim: int, in_dim: int,
                              group: int, bits: int, num_experts: int):
    weight, scales, biases = _projection_memmaps(
        compute_dir, layer, proj, out_dim, in_dim, group, bits, num_experts)
    ids = np.asarray(expert_ids, dtype=np.int64)
    # advanced indexing 已经产出 numpy 临时数组；mx.array 再接管成 MLX-managed staging arrays。
    return (
        mx.array(np.asarray(weight[ids])),
        mx.array(np.asarray(scales[ids])),
        mx.array(np.asarray(biases[ids])),
    )


def _stage_one_projection_uncached(compute_dir: str, layer: int, proj: str, expert: int,
                                  out_dim: int, in_dim: int, group: int, bits: int,
                                  num_experts: int):
    weight, scales, biases = _projection_memmaps(
        compute_dir, layer, proj, out_dim, in_dim, group, bits, num_experts)
    return (
        mx.array(np.asarray(weight[int(expert)])),
        mx.array(np.asarray(scales[int(expert)])),
        mx.array(np.asarray(biases[int(expert)])),
    )


def _stage_one_projection(compute_dir: str, layer: int, proj: str, expert: int,
                          out_dim: int, in_dim: int, group: int, bits: int,
                          num_experts: int):
    key = (compute_dir, int(layer), proj, int(expert), out_dim, in_dim, group, bits, num_experts)
    cached = _EXPERT_STAGE_CACHE.get(key)
    if cached is not None:
        _EXPERT_STAGE_CACHE.move_to_end(key)
        _STAGE_STATS["expert_hits"] += 1
        return cached
    _STAGE_STATS["expert_misses"] += 1
    staged = _stage_one_projection_uncached(
        compute_dir, layer, proj, expert, out_dim, in_dim, group, bits, num_experts)
    _EXPERT_STAGE_CACHE[key] = staged
    _EXPERT_STAGE_CACHE.move_to_end(key)
    limit = max(0, config.native_moe_stage_cache_experts())
    while limit and len(_EXPERT_STAGE_CACHE) > limit:
        _EXPERT_STAGE_CACHE.popitem(last=False)
        _STAGE_STATS["evictions"] += 1
    return staged


def _slot_pool(compute_dir: str, layer: int, hidden: int, inter: int, group: int,
               bits: int, num_experts: int) -> NativeComputeSlotPool:
    cap = max(1, config.native_moe_slot_cap())
    key = (compute_dir, int(layer), hidden, inter, group, bits, num_experts, cap)
    pool = _SLOT_POOLS.get(key)
    if pool is None:
        pool = NativeComputeSlotPool(
            compute_dir, int(layer), hidden, inter, group, bits, num_experts, cap)
        _SLOT_POOLS[key] = pool
    return pool


def _stage_bundle(compute_dir: str, layer: int, expert_ids: list[int],
                  hidden: int, inter: int, group: int, bits: int,
                  num_experts: int):
    key = (
        compute_dir, int(layer), tuple(int(e) for e in expert_ids),
        hidden, inter, group, bits, num_experts,
    )
    if config.native_moe_stage_bundle_cache():
        cached = _BUNDLE_STAGE_CACHE.get(key)
        if cached is not None:
            _BUNDLE_STAGE_CACHE.move_to_end(key)
            _STAGE_STATS["bundle_hits"] += 1
            return cached
    _STAGE_STATS["bundle_misses"] += 1
    gate = _stage_projection(
        compute_dir, int(layer), "gate_proj", expert_ids, inter, hidden, group, bits, num_experts)
    up = _stage_projection(
        compute_dir, int(layer), "up_proj", expert_ids, inter, hidden, group, bits, num_experts)
    down = _stage_projection(
        compute_dir, int(layer), "down_proj", expert_ids, hidden, inter, group, bits, num_experts)
    staged = (*gate, *up, *down)
    if config.native_moe_stage_bundle_cache():
        _BUNDLE_STAGE_CACHE[key] = staged
        _BUNDLE_STAGE_CACHE.move_to_end(key)
        limit = max(0, config.native_moe_stage_cache_bundles())
        while limit and len(_BUNDLE_STAGE_CACHE) > limit:
            _BUNDLE_STAGE_CACHE.popitem(last=False)
            _STAGE_STATS["evictions"] += 1
    return staged


def prefetch_native_moe_stage(layer: int, expert_ids: list[int],
                              hidden: int, inter: int, group: int, bits: int,
                              num_experts: int) -> bool:
    """提前把预测专家拷入 MLX-managed staging cache，供后续 fused native op 复用。"""
    if not config.native_moe():
        return False
    if not config.native_moe_stage_prefetch():
        return False
    if config.native_moe_synthetic():
        return False
    compute_dir = _compute_dir()
    if not compute_dir or not _has_matching_compute_buffers(
        compute_dir, layer, hidden, inter, group, bits, num_experts):
        return False
    try:
        uniq = list(dict.fromkeys(int(e) for e in expert_ids))
        if not uniq:
            return False
        for proj, out_dim, in_dim in (
            ("gate_proj", inter, hidden),
            ("up_proj", inter, hidden),
            ("down_proj", hidden, inter),
        ):
            for e in uniq:
                _stage_one_projection(
                    compute_dir, int(layer), proj, e, out_dim, in_dim, group, bits, num_experts)
        return True
    except Exception:
        if config.native_moe_raise():
            raise
        return False


def stage_cache_stats() -> dict:
    slot_stats = {
        "pools": len(_SLOT_POOLS),
        "entries": sum(len(pool.slot_of) for pool in _SLOT_POOLS.values()),
        "hits": sum(pool.hits for pool in _SLOT_POOLS.values()),
        "misses": sum(pool.misses for pool in _SLOT_POOLS.values()),
        "evictions": sum(pool.evictions for pool in _SLOT_POOLS.values()),
        "rebuilds": sum(pool.rebuilds for pool in _SLOT_POOLS.values()),
    }
    return {
        **_STAGE_STATS,
        "route_sync_s": round(_STAGE_STATS["route_sync_s"], 6),
        "stage_s": round(_STAGE_STATS["stage_s"], 6),
        "enqueue_s": round(_STAGE_STATS["enqueue_s"], 6),
        "expert_entries": len(_EXPERT_STAGE_CACHE),
        "bundle_entries": len(_BUNDLE_STAGE_CACHE),
        "slot_pool": slot_stats,
    }


def note_route_sync(seconds: float) -> None:
    _STAGE_STATS["route_sync_s"] += float(seconds)


def can_native_moe(layer: int, hidden: int, inter: int, group: int, bits: int,
                   num_experts: int) -> bool:
    """只做缓存化的轻量可用性判断；失败时避免上层触发 route host 同步。"""
    _STAGE_STATS["can_checks"] += 1
    if not config.native_moe():
        _STAGE_STATS["can_false"] += 1
        return False
    if not config.native_moe_mlx_op():
        _STAGE_STATS["can_false"] += 1
        return False
    if config.native_moe_synthetic():
        ok = _load_ext() is not None
        _STAGE_STATS["can_true" if ok else "can_false"] += 1
        return ok
    if (hidden * bits) % 32 != 0 or (inter * bits) % 32 != 0:
        _STAGE_STATS["can_false"] += 1
        return False
    compute_dir = _compute_dir()
    if not compute_dir:
        _STAGE_STATS["can_false"] += 1
        return False
    ok = (
        _load_ext() is not None
        and _has_matching_compute_buffers(
            compute_dir, int(layer), hidden, inter, group, bits, num_experts)
    )
    _STAGE_STATS["can_true" if ok else "can_false"] += 1
    return ok


def try_native_moe(layer: int, expert_ids: list[int], x: mx.array, scores: mx.array,
                   hidden: int, inter: int, group: int, bits: int,
                   num_experts: int) -> "mx.array | None":
    """尝试 native fused MoE，失败返回 None 让调用方回退 MLX 路径。"""
    if not config.native_moe():
        return None
    if not config.native_moe_mlx_op():
        return None
    if (hidden * bits) % 32 != 0 or (inter * bits) % 32 != 0:
        return None
    if not expert_ids:
        return None
    if not can_native_moe(layer, hidden, inter, group, bits, num_experts):
        return None
    compute_dir = _compute_dir()
    ext = _load_ext()
    if ext is None:
        return None
    synthetic = config.native_moe_synthetic()
    if (not synthetic and not _has_matching_compute_buffers(
            compute_dir, layer, hidden, inter, group, bits, num_experts)):
        return None
    try:
        expert_arr = mx.array([int(e) for e in expert_ids], dtype=mx.uint32)
        if synthetic:
            t_enqueue = time.perf_counter()
            y = ext.fused_moe(
                x.astype(mx.float32),
                expert_arr,
                scores.astype(mx.float32),
                compute_dir,
                int(layer),
                hidden,
                inter,
                group,
                bits,
                num_experts,
                True,
            )
            _STAGE_STATS["enqueue_s"] += time.perf_counter() - t_enqueue
        else:
            t_stage = time.perf_counter()
            if config.native_moe_slot_pool():
                pool = _slot_pool(
                    compute_dir, int(layer), hidden, inter, group, bits, num_experts)
                local_slots, pool_arrays = pool.acquire(expert_ids)
                local_arr = mx.array(local_slots, dtype=mx.uint32)
                (
                    gate_w, gate_s, gate_b,
                    up_w, up_s, up_b,
                    down_w, down_s, down_b,
                ) = pool_arrays
            else:
                local_arr = None
                (
                    gate_w, gate_s, gate_b,
                    up_w, up_s, up_b,
                    down_w, down_s, down_b,
                ) = _stage_bundle(
                    compute_dir, int(layer), expert_ids,
                    hidden, inter, group, bits, num_experts)
            _STAGE_STATS["stage_s"] += time.perf_counter() - t_stage
            t_enqueue = time.perf_counter()
            if local_arr is not None:
                y = ext.fused_moe_slots(
                    x.astype(mx.float32),
                    local_arr,
                    scores.astype(mx.float32),
                    gate_w, gate_s, gate_b,
                    up_w, up_s, up_b,
                    down_w, down_s, down_b,
                    hidden, inter, group, bits,
                )
            else:
                y = ext.fused_moe_staged(
                    x.astype(mx.float32),
                    scores.astype(mx.float32),
                    gate_w, gate_s, gate_b,
                    up_w, up_s, up_b,
                    down_w, down_s, down_b,
                    hidden, inter, group, bits,
                )
            _STAGE_STATS["enqueue_s"] += time.perf_counter() - t_enqueue
        _STAGE_STATS["calls"] += 1
        return y
    except Exception:
        if config.native_moe_raise():
            raise
        return None