"""Hub format of a DSSD head set: `heads.safetensors` + `dssd_config.json`.

heads.safetensors, for every head k (bf16 except the token ids):
  h{k}.norm.weight [d], h{k}.proj.weight [B_k + 1, d] (last row = reject class ⊥),
  h{k}.reject_bias [1], h{k}.vocab_ids [B_k] (int32).
dssd_config.json: format_version, base_model, base_revision, head_layers, budgets,
reject_class, dtype, default_domain, domains {name: {policy, stop}}, provenance.
A policy uses the bench format (scripts/compose_final.py): type "threshold", thresholds,
cum_gamma, acc_maps {k: [knots, values]}, max_draft, tree.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import torch

from .heads import ExitHead

FORMAT_VERSION = 1
HEADS_FILE = "heads.safetensors"
CONFIG_FILE = "dssd_config.json"
BASE_DIMS = ("hidden_size", "vocab_size", "num_hidden_layers")


def save_heads(out_dir: str | Path, heads: list[ExitHead], layers: list[int], config: dict,
               model_config=None) -> None:
    """Write heads.safetensors and dssd_config.json (format, layers, budgets and reject class
    are filled in here; every other key comes from `config`). With the base model's config, its
    dimensions are recorded and checked at load time."""
    from safetensors.torch import save_file
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    tensors = {}
    for k, h in enumerate(heads):
        if h.adapter is not None or h.vocab_ids is None or not h.reject:
            raise ValueError("only pruned heads with a reject class and no adapter are supported")
        tensors[f"h{k}.norm.weight"] = h.norm.weight.detach().to(torch.bfloat16).contiguous()
        tensors[f"h{k}.proj.weight"] = h.proj.weight.detach().to(torch.bfloat16).contiguous()
        tensors[f"h{k}.reject_bias"] = h.reject_bias.detach().to(torch.bfloat16).contiguous()
        tensors[f"h{k}.vocab_ids"] = h.vocab_ids.to(torch.int32).contiguous()
    save_file({n: t.cpu() for n, t in tensors.items()}, str(out / HEADS_FILE))
    conf = {"format_version": FORMAT_VERSION, **config, "head_layers": [int(x) for x in layers],
            "budgets": [h.n_vocab for h in heads], "reject_class": True, "dtype": "bfloat16"}
    if model_config is not None:
        conf["base_dims"] = {k: int(getattr(model_config, k)) for k in BASE_DIMS}
    (out / CONFIG_FILE).write_text(json.dumps(conf, indent=1, ensure_ascii=False))


def _resolve(path_or_repo: str, revision: str | None, token: str | None) -> Path:
    if os.path.isdir(path_or_repo):
        return Path(path_or_repo)
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(path_or_repo, revision=revision, token=token,
                                  allow_patterns=[HEADS_FILE, CONFIG_FILE]))


def read_config(path_or_repo: str, revision: str | None = None, token: str | None = None) -> dict:
    return json.loads((_resolve(path_or_repo, revision, token) / CONFIG_FILE).read_text())


def load_heads(path_or_repo: str, model_config, device="cpu", dtype=torch.bfloat16,
               revision: str | None = None, token: str | None = None) -> tuple[list[ExitHead], list[int], dict]:
    """Heads, their layers and the config from a local directory or a Hub repo. Checks that
    they fit the base model (hidden size, vocabulary, number of layers)."""
    from safetensors.torch import load_file
    root = _resolve(path_or_repo, revision, token)
    conf = json.loads((root / CONFIG_FILE).read_text())
    if conf.get("format_version") != FORMAT_VERSION:
        raise ValueError(f"unsupported format_version {conf.get('format_version')!r} (expected {FORMAT_VERSION})")
    for k, v in conf.get("base_dims", {}).items():  # heads trained on another base model
        if int(getattr(model_config, k)) != v:
            raise ValueError(f"base model {k} {getattr(model_config, k)} != {v} recorded with the heads")
    layers = [int(x) for x in conf["head_layers"]]
    n_layers = model_config.num_hidden_layers
    if any(not 0 < x < n_layers for x in layers):
        raise ValueError(f"head layer out of range for a {n_layers}-layer base model: {layers}")
    t = load_file(str(root / HEADS_FILE))
    heads = []
    for k in range(len(layers)):
        w, ids = t[f"h{k}.proj.weight"], t[f"h{k}.vocab_ids"].long()
        if w.shape[1] != model_config.hidden_size:
            raise ValueError(f"head {k}: hidden_size {w.shape[1]} != base model {model_config.hidden_size}")
        if int(ids.max()) >= model_config.vocab_size:
            raise ValueError(f"head {k}: token id {int(ids.max())} >= base model vocab_size {model_config.vocab_size}")
        h = ExitHead(model_config.hidden_size, model_config.vocab_size, model_config.rms_norm_eps,
                     vocab_ids=ids, reject=True)
        h.load_state_dict({"norm.weight": t[f"h{k}.norm.weight"], "proj.weight": w,
                           "reject_bias": t[f"h{k}.reject_bias"], "vocab_ids": ids})
        heads.append(h.to(device=device, dtype=dtype).eval())
    return heads, layers, conf


def heads_from_pt(path: str, model_config) -> tuple[list[ExitHead], list[int]]:
    """Pruned heads from a training checkpoint {'heads': [state_dict], 'meta': {'layers'}}
    (scripts/train_pruned.py, merge_pruned_heads.py), dtype preserved."""
    ck = torch.load(path, map_location="cpu", weights_only=False)
    layers = [int(x) for x in ck["meta"]["layers"]]
    heads = []
    for k in range(len(layers)):
        sd = ck["heads"][k]
        if any(n.startswith("adapter.") for n in sd) or sd.get("vocab_ids") is None or "reject_bias" not in sd:
            raise ValueError(f"head {k}: only pruned heads with a reject class and no adapter are supported")
        n_out = sd["proj.weight"].shape[0] - 1
        h = ExitHead(model_config.hidden_size, model_config.vocab_size, model_config.rms_norm_eps,
                     vocab_ids=sd["vocab_ids"][:n_out], reject=True)
        h.load_state_dict(sd)
        heads.append(h.to(sd["proj.weight"].dtype).eval())
    return heads, layers