"""Build a TITAN model from an MLX checkpoint: the host architecture (attention, norms, router) is
mlx-lm's own module tree; every layer's `switch_mlp` is replaced by a TitanSwitchGLU over a slot table.
Expert weights are never materialised from the checkpoint by mlx-lm: they are read on demand by the
ExpertStore. `capacity` = number of resident experts per layer (num_experts = everything resident).
"""
from __future__ import annotations

import os
from typing import Dict, List, Optional

import mlx.core as mx
import mlx.nn as nn
from mlx_lm.utils import load_model, load_tokenizer

from .moe import TitanMoEBlock, TitanSwitchGLU
from .prefetch import Prefetcher
from . import predictors as _pred
from .registry import ModelSpec, load_spec
from .slots import EvictionPolicy, LFUPolicy, SlotTable
from .store import ExpertStore

POLICIES = {"lru": EvictionPolicy, "lfu": LFUPolicy}


class TitanModel:
    def __init__(self, model_dir: str, capacity: Optional[int] = None, policy: str = "lru", nocache: bool = False,
                 on_route=None, predictor: str = "none", m: int = 8, prefetch_workers: int = 2, traces_dir: Optional[str] = None):
        self.model_dir = os.path.realpath(model_dir)
        self.spec: ModelSpec = load_spec(self.model_dir)
        self.store = ExpertStore(self.spec, nocache=nocache)
        self.capacity = capacity or self.spec.num_experts
        self.policy_name = policy
        # Load the host model lazily so the stacked expert tensors are never evaluated.
        from pathlib import Path
        self.model, cfg = load_model(Path(self.model_dir), lazy=True)
        self.tokenizer = load_tokenizer(Path(self.model_dir))
        q = cfg.get("quantization") or cfg.get("quantization_config")
        self.tables: List[SlotTable] = []
        self.blocks: List[TitanSwitchGLU] = []
        for l, layer in enumerate(self.model.model.layers):
            table = SlotTable(l, self.capacity, self.store, POLICIES[policy]())
            block = TitanSwitchGLU(table, q["group_size"], q["bits"], on_route=on_route)
            # drop the lazy stacked weights and splice in the TITAN block
            mlp = layer.mlp
            del mlp.switch_mlp  # removes the lazy stacked SwitchGLU (never evaluated)
            mlp.switch_mlp = block  # plain attribute: not an nn.Module, not in parameters()
            self.tables.append(table)
            self.blocks.append(block)
        mx.eval(self.model.parameters())  # dense weights only
        # optional predictive prefetcher: replaces each layer's MoE block with the hooked exact-routing block
        self.prefetcher = None
        self.predictor_name = predictor
        if predictor != "none":
            pred = _pred.make(predictor, self.spec.num_layers, self.spec.num_experts)
            if traces_dir is not None and hasattr(pred, "fit"):
                pred.fit(_load_fit_traces(traces_dir))
            self.prefetcher = Prefetcher(self.tables, self.store, pred, m=m, workers=prefetch_workers)
            for l, layer in enumerate(self.model.model.layers):
                self.blocks[l].prefetcher = self.prefetcher
                layer.mlp = TitanMoEBlock(layer.mlp, self.blocks[l], self.spec.model_type, self.prefetcher)

    # convenience -----------------------------------------------------------------
    @property
    def layers(self):
        return self.model.layers

    def __call__(self, inputs, cache=None):
        return self.model(inputs, cache=cache)

    def attach_workers(self, workers: Dict[str, "WorkerClient"], placement: Dict, fallback_local: bool = True):
        """Enable distributed execution (Gate E). placement: {(layer, expert): worker_name or None}."""
        from .distributed import RemoteExecutor
        self.remote = RemoteExecutor(workers, placement, fallback_local)
        for b in self.blocks:
            b.remote = self.remote
        # a layer with every expert resident locally would take the sync-free fast path and bypass the remote split:
        # the capacity is bounded by construction in distributed runs (capacity < num_experts), so no change is needed here.
        return self.remote

    def warm_all(self):
        """Load every expert (only meaningful when capacity == num_experts)."""
        for t in self.tables:
            t.ensure(list(range(self.spec.num_experts)))
            t.reset_stats()
        self.store.reset_stats()

    def cache_stats(self) -> Dict:
        per = [t.stats() for t in self.tables]
        if self.prefetcher is not None:
            per_pf = self.prefetcher.stats()
        else:
            per_pf = None
        hits = sum(p["hits"] + p["hits_device"] for p in per)
        misses = sum(p["misses"] for p in per)
        return {"layers": per, "hits": hits, "misses": misses,
                "hit_rate": hits / (hits + misses) if hits + misses else None,
                "stall_time_s": sum(p["stall_time_s"] for p in per), "store": self.store.stats(), "prefetch": per_pf}

    def reset_stats(self):
        for t in self.tables:
            t.reset_stats()
        self.store.reset_stats()
        if self.prefetcher is not None:
            self.prefetcher.reset_stats()


def _load_fit_traces(traces_dir: str):
    import json, numpy as np
    idx = json.load(open(os.path.join(traces_dir, "index.json")))
    out = []
    for p in idx["prompts"]:
        if p["split"] != "fit":
            continue
        z = np.load(os.path.join(traces_dir, p["name"] + ".npz"))
        out.append((z["topk"].astype(np.int64), z["probs"]))
    return out