"""Tensor-at-a-time safetensors loading; never load a whole checkpoint to RAM."""

from __future__ import annotations

import gc
import hashlib
import json
import struct
from collections import OrderedDict
from pathlib import Path

import mlx.core as mx
import numpy as np


class WeightStore:
    def __init__(self, path, *, cache_gb=0.0, cast=None, reader="mlx"):
        self.path = Path(path)
        files = (
            [self.path]
            if self.path.is_file()
            else sorted(self.path.glob("*.safetensors"))
        )
        if not files:
            raise FileNotFoundError(f"No safetensors weights at {self.path}")
        self.index = {}
        self.headers = {}
        self.data_start = {}
        self.cache = OrderedDict()
        self.cache_limit = int(cache_gb * 1024**3)
        self.cached_bytes = 0
        self.cast = cast
        if reader not in ("mlx", "memmap"):
            raise ValueError("reader must be mlx or memmap")
        self.reader = reader
        self.loads = 0
        self.bytes_read = 0
        self.files = files
        for file in files:
            with file.open("rb") as handle:
                raw = handle.read(8)
                if len(raw) != 8:
                    raise ValueError(f"Truncated safetensors header: {file}")
                length = struct.unpack("<Q", raw)[0]
                if length > 100_000_000:
                    raise ValueError(f"Invalid safetensors header size: {file}")
                header = json.loads(handle.read(length))
                self.data_start[file] = 8 + length
            for name, spec in header.items():
                if name == "__metadata__":
                    continue
                if name in self.index:
                    raise ValueError(f"Duplicate tensor {name} in {self.path}")
                self.index[name] = file
                self.headers[name] = spec

    def __contains__(self, key):
        return key in self.index

    def fingerprint(self):
        # Source identity for cache invalidation, not a content-integrity checksum.
        info = [
            (str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns)
            for p in self.files
        ]
        return hashlib.sha256(json.dumps(info).encode()).hexdigest()

    def get(self, key, *, rows=None, cache=True):
        if rows is None and key in self.cache:
            self.cache.move_to_end(key)
            return self.cache[key]
        if key not in self.index:
            raise KeyError(f"Missing checkpoint tensor {key} in {self.path}")
        # safetensors' Python MLX adapter routes BF16 through NumPy, which has
        # no BF16 dtype. Map its uint16 payload directly and bitcast in MLX.
        spec, file = self.headers[key], self.index[key]
        dtypes = {
            "BF16": "<u2",
            "F16": "<f2",
            "F32": "<f4",
            "U32": "<u4",
            "I32": "<i4",
            "I64": "<i8",
            "U8": "u1",
            "BOOL": "?",
        }
        if spec["dtype"] not in dtypes:
            raise ValueError(f"Unsupported tensor dtype {spec['dtype']} for {key}")
        if self.reader == "mlx" and rows is None:
            # MLX creates lazy per-tensor Load nodes. Discard the dictionary
            # immediately: retaining it would also retain every evaluated weight
            # in the shard, bypassing this store's explicit cache budget.
            value = mx.load(str(file))[key]
        else:
            # A full-vocabulary MLX Load is evaluated before a row gather.
            # Read requested rows directly from the safetensors mapping instead;
            # this preserves order/duplicates and the original BF16 bit payload.
            mapped = np.memmap(
                file,
                dtype=dtypes[spec["dtype"]],
                mode="r",
                offset=self.data_start[file] + spec["data_offsets"][0],
                shape=tuple(spec["shape"]),
            )
            value = mx.array(mapped if rows is None else mapped[rows])
            if spec["dtype"] == "BF16":
                value = value.view(mx.bfloat16)
            del mapped
        if self.cast is not None and mx.issubdtype(value.dtype, mx.floating):
            value = value.astype(self.cast)
        mx.eval(value)
        self.loads += 1
        self.bytes_read += value.nbytes
        if rows is None and cache and value.nbytes <= self.cache_limit:
            while self.cache and self.cached_bytes + value.nbytes > self.cache_limit:
                _, old = self.cache.popitem(last=False)
                self.cached_bytes -= old.nbytes
            self.cache[key] = value
            self.cached_bytes += value.nbytes
        return value

    def group(self, prefix, *, exclude=()):
        return {
            k[len(prefix) :]: self.get(k)
            for k in self.index
            if k.startswith(prefix) and not any(s in k for s in exclude)
        }

    def clear(self):
        self.cache.clear()
        self.cached_bytes = 0
        gc.collect()
        mx.clear_cache()


class Linear:
    def __init__(
        self, weight, bias=None, scales=None, qbiases=None, bits=8, group_size=64
    ):
        self.weight = weight
        self.bias = bias
        self.scales = scales
        self.qbiases = qbiases
        self.bits = bits
        self.group_size = group_size

    @classmethod
    def load(cls, store, prefix, *, bits=None, group_size=64, grouped_heads=None):
        # Converted MLX quantized tensors use .scales/.biases, independently of bias.
        if prefix + ".scales" in store:
            meta = json.loads((store.path / "quantization.json").read_text())
            return cls(
                store.get(prefix + ".weight"),
                store.get(prefix + ".bias") if prefix + ".bias" in store else None,
                store.get(prefix + ".scales"),
                store.get(prefix + ".biases"),
                meta["bits"],
                meta["group_size"],
            )
        w = store.get(prefix + ".weight")
        b = store.get(prefix + ".bias") if prefix + ".bias" in store else None
        # Original H3 QKV is [head, qkv, dim, input], unlike MLX [qkv, head, dim, input].
        if grouped_heads is not None:
            heads, dim = grouped_heads
            w = w.reshape(heads, 3, dim, -1).transpose(1, 0, 2, 3).reshape(w.shape)
            if b is not None:
                b = b.reshape(heads, 3, dim).transpose(1, 0, 2).reshape(b.shape)
        if bits is not None and w.dtype != mx.float32:
            q, s, z = mx.quantize(w, bits=bits, group_size=group_size)
            mx.eval(q, s, z)
            return cls(q, b, s, z, bits, group_size)
        mx.eval(w)
        return cls(w, b)

    @property
    def dtype(self):
        return self.scales.dtype if self.scales is not None else self.weight.dtype

    @property
    def nbytes(self):
        return sum(
            x.nbytes
            for x in (self.weight, self.bias, self.scales, self.qbiases)
            if x is not None
        )

    def __call__(self, x):
        x = x.astype(self.dtype)
        if self.scales is None:
            # One BF16 rounding after the bias epilogue, as in torch F.linear.
            return (
                x @ self.weight.T
                if self.bias is None
                else mx.addmm(self.bias, x, self.weight.T)
            )
        else:
            out = mx.quantized_matmul(
                x,
                self.weight,
                self.scales,
                self.qbiases,
                transpose=True,
                group_size=self.group_size,
                bits=self.bits,
            )
        return out if self.bias is None else out + self.bias