"""Model and expert registry.

Parses a Hugging Face / MLX MoE checkpoint directory (config.json + safetensors) into an explicit
description of layers, experts, routing semantics and the *byte ranges* of every expert tensor,
so that the expert store can read a single expert from disk with `pread` without touching the rest
of the file. Only formats that we have verified are accepted (see `SUPPORTED`).
"""
from __future__ import annotations

import json
import os
import struct
from dataclasses import dataclass, field
from typing import Dict, List, Tuple

DTYPE_ITEMSIZE = {"F32": 4, "F16": 2, "BF16": 2, "U32": 4, "I32": 4, "U8": 1, "I8": 1, "U16": 2}

# model_type -> (expert tensor key template, matrices, routing description)
SUPPORTED = {
    "olmoe": {
        "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",
        "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",
        "matrices": ("gate", "up", "down"),
        "router_key": "model.layers.{l}.mlp.gate.{s}",
        "routing": "softmax_then_topk",  # softmax over all experts, then top-k, weights = softmax probs (renormalised iff norm_topk_prob)
        "shared_experts": 0,
    },
    "qwen3_moe": {
        "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",
        "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",
        "matrices": ("gate", "up", "down"),
        "router_key": "model.layers.{l}.mlp.gate.{s}",
        "routing": "softmax_then_topk",
        "shared_experts": 0,
    },
}


@dataclass(frozen=True)
class TensorLoc:
    key: str
    dtype: str
    shape: Tuple[int, ...]
    file: str
    offset: int  # absolute byte offset in file
    nbytes: int


@dataclass
class ExpertSpec:
    layer: int
    expert: int
    tensors: Dict[str, TensorLoc]  # e.g. "gate.weight" -> loc

    @property
    def nbytes(self) -> int:
        return sum(t.nbytes for t in self.tensors.values())


@dataclass
class ModelSpec:
    path: str
    model_type: str
    num_layers: int
    num_experts: int
    top_k: int
    norm_topk_prob: bool
    hidden_size: int
    intermediate_size: int
    quantization: Dict[str, int]
    routing: str
    shared_experts: int
    experts: Dict[Tuple[int, int], ExpertSpec] = field(default_factory=dict)
    dense: Dict[str, TensorLoc] = field(default_factory=dict)
    layout: str = "per_expert"  # or "stacked"

    @property
    def bytes_per_expert(self) -> int:
        return next(iter(self.experts.values())).nbytes

    @property
    def expert_bytes_total(self) -> int:
        return sum(e.nbytes for e in self.experts.values())

    @property
    def dense_bytes_total(self) -> int:
        return sum(t.nbytes for t in self.dense.values())

    def summary(self) -> dict:
        return {
            "path": self.path,
            "model_type": self.model_type,
            "num_layers": self.num_layers,
            "num_experts": self.num_experts,
            "top_k": self.top_k,
            "norm_topk_prob": self.norm_topk_prob,
            "hidden_size": self.hidden_size,
            "intermediate_size": self.intermediate_size,
            "quantization": self.quantization,
            "routing": self.routing,
            "shared_experts": self.shared_experts,
            "layout": self.layout,
            "bytes_per_expert": self.bytes_per_expert,
            "expert_bytes_total": self.expert_bytes_total,
            "dense_bytes_total": self.dense_bytes_total,
            "experts_per_token": self.num_layers * self.top_k,
            "bytes_per_token_if_uncached": self.num_layers * self.top_k * self.bytes_per_expert,
        }


def read_safetensors_header(path: str) -> Dict[str, TensorLoc]:
    """Parse the safetensors header and return absolute byte locations for every tensor."""
    locs = {}
    with open(path, "rb") as f:
        (n,) = struct.unpack("<Q", f.read(8))
        header = json.loads(f.read(n))
    base = 8 + n
    for key, meta in header.items():
        if key == "__metadata__":
            continue
        a, b = meta["data_offsets"]
        locs[key] = TensorLoc(key, meta["dtype"], tuple(meta["shape"]), path, base + a, b - a)
    return locs


def titan_layout_dir(model_dir: str):
    import hashlib
    d = os.path.expanduser(f"~/.cache/titan/{hashlib.sha1(os.path.realpath(model_dir).encode()).hexdigest()[:12]}")
    return d if os.path.exists(os.path.join(d, "experts.titan.json")) else None


def apply_titan_layout(spec: "ModelSpec", layout_dir: str) -> "ModelSpec":
    """Point every expert's tensors at the contiguous blob produced by scripts/convert_layout.py."""
    idx = json.load(open(os.path.join(layout_dir, "experts.titan.json")))
    blob = os.path.join(layout_dir, "experts.titan")
    for key, ent in idx["experts"].items():
        l, e = map(int, key.split(","))
        es = spec.experts[(l, e)]
        es.tensors = {name: TensorLoc(f"titan:{l},{e}:{name}", p["dtype"], tuple(p["shape"]), blob, ent["offset"] + p["off"], p["n"])
                      for name, p in ent["tensors"].items()}
    spec.layout = spec.layout + "+titan_contiguous"
    return spec


def load_spec(model_dir: str, use_titan_layout: bool = True) -> ModelSpec:
    model_dir = os.path.realpath(model_dir)
    with open(os.path.join(model_dir, "config.json")) as f:
        cfg = json.load(f)
    mt = cfg["model_type"]
    if mt not in SUPPORTED:
        raise ValueError(f"model_type {mt!r} is not in the verified set {sorted(SUPPORTED)}")
    fmt = SUPPORTED[mt]
    locs: Dict[str, TensorLoc] = {}
    for fn in sorted(os.listdir(model_dir)):
        if fn.endswith(".safetensors"):
            locs.update(read_safetensors_header(os.path.join(model_dir, fn)))
    q = cfg.get("quantization") or cfg.get("quantization_config") or {}
    spec = ModelSpec(
        path=model_dir, model_type=mt, num_layers=cfg["num_hidden_layers"], num_experts=cfg["num_experts"],
        top_k=cfg["num_experts_per_tok"], norm_topk_prob=bool(cfg.get("norm_topk_prob", False)),
        hidden_size=cfg["hidden_size"], intermediate_size=cfg["intermediate_size"],
        quantization={"group_size": q.get("group_size"), "bits": q.get("bits")}, routing=fmt["routing"],
        shared_experts=fmt["shared_experts"],
    )
    per_expert_keys = set()
    for l in range(spec.num_layers):
        for e in range(spec.num_experts):
            tensors = {}
            for m in fmt["matrices"]:
                for s in ("weight", "scales", "biases"):
                    k = fmt["expert_key"].format(l=l, e=e, m=m, s=s)
                    if k in locs:
                        tensors[f"{m}.{s}"] = locs[k]
                        per_expert_keys.add(k)
            if not tensors:
                # stacked layout: switch_mlp.<m>_proj.<s> has shape [E, ...]; expert e is the contiguous slice e of the first axis
                for m in fmt["matrices"]:
                    for s in ("weight", "scales", "biases"):
                        k = fmt["stacked_key"].format(l=l, m=m, s=s)
                        if k in locs:
                            loc = locs[k]
                            if loc.shape[0] != spec.num_experts:
                                raise ValueError(f"{k}: first axis {loc.shape[0]} != num_experts {spec.num_experts}")
                            per = loc.nbytes // spec.num_experts
                            tensors[f"{m}.{s}"] = TensorLoc(f"{k}[{e}]", loc.dtype, tuple(loc.shape[1:]), loc.file, loc.offset + e * per, per)
                            per_expert_keys.add(k)
                spec.layout = "stacked"
            if not tensors:
                raise ValueError(f"no expert tensors found for layer {l} expert {e} (neither per-expert nor stacked keys)")
            spec.experts[(l, e)] = ExpertSpec(l, e, tensors)
    spec.dense = {k: v for k, v in locs.items() if k not in per_expert_keys}
    if use_titan_layout:
        d = titan_layout_dir(model_dir)
        if d:
            apply_titan_layout(spec, d)
    return spec


if __name__ == "__main__":
    import sys
    s = load_spec(sys.argv[1])
    print(json.dumps(s.summary(), indent=2))