"""Reading checkpoints and adapters, and reading/writing llama.cpp control-vector GGUFs.

Tensors are read straight from safetensors by byte offset, so a 60 GB checkpoint is never loaded whole:
the fit reads one layer at a time. bf16 is upcast to fp32 exactly (the bf16 bits are the high half of
the fp32 word).
"""
from __future__ import annotations

import json
import struct
from pathlib import Path

import numpy as np

PREFIXES = ("model.language_model.", "model.")   # multimodal exports, then text-only exports
N_LAYERS = 60


def _header(path: Path) -> tuple[int, dict]:
    with open(path, "rb") as fh:
        n = struct.unpack("<Q", fh.read(8))[0]
        return n, json.loads(fh.read(n))


_HEADERS: dict[str, tuple[int, dict]] = {}


def read_tensor(path: Path, name: str) -> np.ndarray:
    """One tensor from a safetensors file as fp32."""
    key = str(path)
    if key not in _HEADERS:
        _HEADERS[key] = _header(path)
    n, hd = _HEADERS[key]
    meta = hd[name]
    a, b = meta["data_offsets"]
    dtype = {"BF16": np.uint16, "F16": np.float16, "F32": np.float32}[meta["dtype"]]
    raw = np.fromfile(path, dtype=dtype, count=(b - a) // np.dtype(dtype).itemsize, offset=8 + n + a)
    raw = raw.reshape(meta["shape"])
    if meta["dtype"] == "BF16":
        return (raw.astype(np.uint32) << 16).view(np.float32)
    return raw.astype(np.float32)


def raw_bytes_equal(path_a: Path, name_a: str, path_b: Path, name_b: str, chunk: int = 1 << 26) -> bool:
    """Whether two stored tensors are bit-identical (same dtype, shape and bytes), read in chunks."""
    na, ha = _header(path_a)
    nb, hb = _header(path_b)
    ma, mb = ha[name_a], hb[name_b]
    if ma["dtype"] != mb["dtype"] or ma["shape"] != mb["shape"]:
        return False
    (a0, a1), (b0, _) = ma["data_offsets"], mb["data_offsets"]
    with open(path_a, "rb") as fa, open(path_b, "rb") as fb:
        fa.seek(8 + na + a0)
        fb.seek(8 + nb + b0)
        left = a1 - a0
        while left:
            k = min(chunk, left)
            if fa.read(k) != fb.read(k):
                return False
            left -= k
    return True


class Checkpoint:
    """Tensor access by name over one safetensors checkpoint directory (sharded or not)."""

    def __init__(self, directory: str | Path):
        self.dir = Path(directory)
        index = self.dir / "model.safetensors.index.json"
        if index.exists():
            self.map = json.loads(index.read_text())["weight_map"]
        else:
            single = self.dir / "model.safetensors"
            self.map = {k: single.name for k in _header(single)[1] if k != "__metadata__"}
        prefix = next((p for p in PREFIXES if any(k.startswith(p + "layers.") for k in self.map)), None)
        assert prefix, f"{self.dir}: no decoder layers found"
        self.prefix: str = prefix
        layers = {int(k[len(self.prefix) + 7:].split(".")[0]) for k in self.map if k.startswith(self.prefix + "layers.")}
        # the vision tower's blocks share module names with the language stack; count only this prefix
        assert layers == set(range(N_LAYERS)), f"{self.dir}: expected {N_LAYERS} decoder layers, found {len(layers)}"

    def file(self, name: str) -> Path:
        return self.dir / self.map[self.prefix + name]

    def has(self, name: str) -> bool:
        return self.prefix + name in self.map

    def get(self, name: str) -> np.ndarray:
        """name relative to the language-model prefix, e.g. 'layers.3.mlp.up_proj.weight'."""
        return read_tensor(self.file(name), self.prefix + name)

    def layer(self, L: int) -> dict[str, np.ndarray]:
        stem = f"{self.prefix}layers.{L}."
        return {k[len(stem):]: read_tensor(self.dir / f, k) for k, f in self.map.items() if k.startswith(stem)}

    def same_as(self, other: "Checkpoint", name: str) -> bool:
        return raw_bytes_equal(self.file(name), self.prefix + name, other.file(name), other.prefix + name)

    def embedding_rows(self, token_ids: np.ndarray) -> np.ndarray:
        """Rows of embed_tokens for the given ids as fp32, read row by row (only the rows used)."""
        name = self.prefix + "embed_tokens.weight"
        path = self.file("embed_tokens.weight")
        n, hd = _header(path)
        meta = hd[name]
        assert meta["dtype"] == "BF16", meta["dtype"]
        width = meta["shape"][1]
        start = 8 + n + meta["data_offsets"][0]
        table = {}
        with open(path, "rb") as fh:
            for t in np.unique(token_ids):
                fh.seek(start + int(t) * width * 2)
                table[int(t)] = (np.frombuffer(fh.read(width * 2), np.uint16).astype(np.uint32) << 16).view(np.float32)
        return np.stack([np.stack([table[int(t)] for t in row]) for row in token_ids])


class Adapter:
    """A PEFT LoRA adapter directory (adapter_config.json + adapter_model.safetensors)."""

    def __init__(self, directory: str | Path, base_prefix: str = "model.language_model."):
        self.dir = Path(directory)
        cfg = json.loads((self.dir / "adapter_config.json").read_text())
        assert not cfg.get("use_rslora") and not cfg.get("use_dora"), "rsLoRA / DoRA adapters are not handled"
        self.rank = int(cfg["r"])
        self.scale = cfg["lora_alpha"] / self.rank
        self.file = self.dir / "adapter_model.safetensors"
        self.keys = set(_header(self.file)[1])
        self.stem = f"base_model.model.{base_prefix}layers."

    def families(self, L: int) -> list[str]:
        """Projection names this adapter trains in layer L, e.g. 'self_attn.q_proj'."""
        stem = f"{self.stem}{L}."
        return sorted({k[len(stem):-len(".lora_A.weight")] for k in self.keys
                       if k.startswith(stem) and k.endswith(".lora_A.weight")})

    def pair(self, L: int, family: str) -> tuple[np.ndarray, np.ndarray]:
        """(A, scale * B) for one projection: the update is (scale * B) @ A."""
        k = f"{self.stem}{L}.{family}"
        A = read_tensor(self.file, f"{k}.lora_A.weight")
        B = read_tensor(self.file, f"{k}.lora_B.weight")
        return A, self.scale * B


# ---- llama.cpp control-vector GGUF ----------------------------------------------------------------------------
# One 1-D F32 tensor per layer named direction.<L>, L >= 1 (llama.cpp never applies a vector at layer 0), and
# the metadata cvector-generator writes (common/common.cpp, common_control_vector_load_one).

def write_gguf(path: str | Path, vectors: np.ndarray, name: str | None = None, alpha: float = 1.0) -> None:
    """Write alpha * vectors[1:] as a control-vector GGUF, so it loads at strength 1.0."""
    from gguf import GGUFWriter

    vectors = np.asarray(vectors, np.float32)
    assert vectors.shape[0] == N_LAYERS, vectors.shape
    w = GGUFWriter(str(path), arch="controlvector")
    w.add_string("controlvector.model_hint", "gemma4")
    w.add_int32("controlvector.layer_count", N_LAYERS - 1)
    w.add_float32("controlvector.alpha", float(alpha))
    if name:
        w.add_string("general.name", name)
    for L in range(1, N_LAYERS):
        w.add_tensor(f"direction.{L}", np.ascontiguousarray(alpha * vectors[L], np.float32))
    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file()
    w.close()


def read_gguf(path: str | Path) -> np.ndarray:
    """A control-vector GGUF's direction.L tensors as an (N_LAYERS, d) array; layer 0 stays zero."""
    from gguf import GGUFReader

    r = GGUFReader(str(path))
    t = {x.name: np.asarray(x.data, np.float32) for x in r.tensors if x.name.startswith("direction.")}
    assert set(t) == {f"direction.{L}" for L in range(1, N_LAYERS)}, (path, len(t))
    out = np.zeros((N_LAYERS, len(t["direction.1"])), np.float32)
    for k, v in t.items():
        out[int(k.split(".")[1])] = v
    return out