import json
import os
import re
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import load_file

from models.demos.blackhole.qwen36.tt.model_config import Qwen36ModelArgs
from models.demos.blackhole.qwen36.tt.weight_mapping import remap_qwen36_state_dict

LANGUAGE_PREFIX = "model.language_model."
VISUAL_PREFIX = "model.visual."
LM_HEAD_KEY = "lm_head.weight"
INDEX_FILE = "model.safetensors.index.json"
HEAD_WEIGHTS = "joint_head.safetensors"
HEAD_CONFIG = "joint_head_config.json"
_HEX40 = re.compile(r"^[0-9a-f]{40}$")
_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
PRECISION_TAG_KEYS = (
    ("gu", "QWEN36_MLP_GATE_UP_DTYPE", "bfp4"),
    ("dn", "QWEN36_MLP_DOWN_DTYPE", "bfp8"),
    ("pj", "QWEN36_PROJ_DTYPE", "bfp8"),
)


def resolve_snapshot():
    for env in ("CLEF_MODEL", "HF_MODEL"):
        value = os.environ.get(env)
        if value and os.path.isfile(os.path.join(value, "config.json")):
            return str(Path(value).resolve())
    raise FileNotFoundError("set CLEF_MODEL (or HF_MODEL) to the local Cloudflare/clef snapshot directory")


def revision_sha8(snapshot):
    name = Path(snapshot).name
    if _HEX40.match(name):
        return name[:8]
    import hashlib

    return hashlib.sha1((Path(snapshot) / INDEX_FILE).read_bytes()).hexdigest()[:8]


def precision_cache_tag():
    return "_".join(f"{short}-{os.environ.get(env, default)}" for short, env, default in PRECISION_TAG_KEYS)


def read_weight_map(snapshot):
    return json.loads((Path(snapshot) / INDEX_FILE).read_text())["weight_map"]


def language_keys(weight_map, layer_filter=None):
    keys = []
    for key in weight_map:
        if not key.startswith(LANGUAGE_PREFIX):
            continue
        match = _LAYER_RE.match(key)
        if match is not None and layer_filter is not None and int(match.group(1)) not in layer_filter:
            continue
        keys.append(key)
    return sorted(keys)


def read_tensors(snapshot, weight_map, keys):
    by_file = {}
    for key in keys:
        by_file.setdefault(weight_map[key], []).append(key)
    out = {}
    for filename in sorted(by_file):
        with safe_open(str(Path(snapshot) / filename), framework="pt") as shard:
            for key in by_file[filename]:
                out[key] = shard.get_tensor(key)
    return out


class LmHeadRows:
    def __init__(self, snapshot):
        self.snapshot = str(snapshot)
        weight_map = read_weight_map(snapshot)
        self.path = str(Path(snapshot) / weight_map[LM_HEAD_KEY])
        with safe_open(self.path, framework="pt") as shard:
            self.shape = tuple(shard.get_slice(LM_HEAD_KEY).get_shape())
        self._cache = {}

    def __call__(self, token_ids):
        ids = [int(t) for t in (token_ids.tolist() if torch.is_tensor(token_ids) else token_ids)]
        missing = sorted({t for t in ids if t not in self._cache})
        if missing:
            with safe_open(self.path, framework="pt") as shard:
                view = shard.get_slice(LM_HEAD_KEY)
                for t in missing:
                    self._cache[t] = view[t : t + 1][0]
        if not ids:
            return torch.empty((0, self.shape[1]), dtype=torch.bfloat16)
        return torch.stack([self._cache[t] for t in ids])

    def __getitem__(self, token_ids):
        return self(token_ids)


class ClefModelArgs(Qwen36ModelArgs):
    def __init__(self, mesh_device=None, max_batch_size=1, max_seq_len=2048, **kwargs):
        self.snapshot = resolve_snapshot()
        os.environ["HF_MODEL"] = self.snapshot
        super().__init__(mesh_device, max_batch_size=max_batch_size, max_seq_len=max_seq_len, **kwargs)
        self.revision_sha8 = revision_sha8(self.snapshot)
        if mesh_device is not None and self.num_devices > 1:
            self.gdn_gate_fp32 = os.environ.get("QWEN36_GDN_GATE_FP32", "1") == "1"
        self.weight_map = read_weight_map(self.snapshot)
        self.layer_indices = None
        self._lm_head_rows = None

    def weight_cache_path(self, dtype=None):
        base = super().weight_cache_path(dtype)
        return base.with_name(f"{base.name}_clef_{self.revision_sha8}_{precision_cache_tag()}")

    def active_layers(self):
        if self.layer_indices is not None:
            return set(int(i) for i in self.layer_indices)
        return set(range(self.n_layers))

    def raw_language_state_dict(self, layer_filter=None):
        if layer_filter is None:
            layer_filter = self.active_layers()
        keys = language_keys(self.weight_map, layer_filter)
        return read_tensors(self.snapshot, self.weight_map, keys)

    def load_state_dict(self):
        raw = self.raw_language_state_dict()
        state_dict = remap_qwen36_state_dict(raw)
        assert "output.weight" not in state_dict
        assert "tok_embeddings.weight" in state_dict and "norm.weight" in state_dict
        return state_dict

    def mapped_keys(self, layer_filter):
        keys = language_keys(self.weight_map, layer_filter)
        placeholders = {k: torch.empty(0, 0, 0) for k in keys}
        return set(remap_qwen36_state_dict(placeholders))

    @property
    def lm_head_rows(self):
        if self._lm_head_rows is None:
            self._lm_head_rows = LmHeadRows(self.snapshot)
        return self._lm_head_rows

    def load_lm_head_rows(self, token_ids):
        return self.lm_head_rows(token_ids)

    def vision_state_dict(self):
        keys = sorted(k for k in self.weight_map if k.startswith(VISUAL_PREFIX))
        raw = read_tensors(self.snapshot, self.weight_map, keys)
        return {k[len(VISUAL_PREFIX) :]: v for k, v in raw.items()}

    def head_state_dict(self):
        return load_file(str(Path(self.snapshot) / HEAD_WEIGHTS))

    def head_config(self):
        return json.loads((Path(self.snapshot) / HEAD_CONFIG).read_text())