# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""Lazy, shard-aware reader for the GR00T sharded-safetensors checkpoints (``weight-loader-spec.md`` §1).

* Every shard header is parsed once with the plain safetensors framing (8-byte little-endian header length followed
  by a JSON header; ``data_offsets`` are relative to the byte after the header) so that shape / dtype / shard / byte
  questions never touch tensor data (§1.1).
* The shard is resolved **per key**, never per layer or category: a shard boundary falls inside a unit in every
  version (N1.5 LLM layer 5, N1.6 DiT block 18, N1.7 DiT block 19; §1.1 consequence (b)).
* Reads go through ``safetensors.safe_open`` (mmap) and return the tensor as stored (bf16) unless a dtype is
  requested; ``slot()`` / ``rows()`` use ``get_slice`` so that one embodiment slot of a ``[32, in, out]`` tensor costs
  ~10 MB instead of 300 MB (§1.2).

No ``ttnn`` import anywhere in this module.
"""

from __future__ import annotations

import json
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, Tuple, Union

import torch
from safetensors import safe_open

from models.experimental.gr00t.common.configs import GR00T_VERSIONS, get_config, snapshot_dir

# safetensors dtype tag -> torch dtype (only tags that can appear in HF checkpoints).
SAFETENSORS_DTYPES: Mapping[str, torch.dtype] = {
    "BF16": torch.bfloat16,
    "F16": torch.float16,
    "F32": torch.float32,
    "F64": torch.float64,
    "I8": torch.int8,
    "U8": torch.uint8,
    "I16": torch.int16,
    "I32": torch.int32,
    "I64": torch.int64,
    "BOOL": torch.bool,
}
_DTYPE_BYTES: Mapping[str, int] = {
    "BF16": 2,
    "F16": 2,
    "F32": 4,
    "F64": 8,
    "I8": 1,
    "U8": 1,
    "I16": 2,
    "I32": 4,
    "I64": 8,
    "BOOL": 1,
}


@dataclass(frozen=True)
class TensorInfo:
    """Header entry of one tensor (no data read)."""

    key: str
    shard: str  # shard file name (e.g. 'model-00002-of-00003.safetensors')
    dtype: str  # safetensors tag, e.g. 'BF16'
    shape: Tuple[int, ...]
    data_offsets: Tuple[int, int]  # [start, end) relative to the first data byte of the shard

    @property
    def nbytes(self) -> int:
        return self.data_offsets[1] - self.data_offsets[0]

    @property
    def numel(self) -> int:
        n = 1
        for d in self.shape:
            n *= d
        return n

    @property
    def torch_dtype(self) -> torch.dtype:
        return SAFETENSORS_DTYPES[self.dtype]


@dataclass(frozen=True)
class ShardInfo:
    """Per-shard framing facts (``weight-loader-spec.md`` §1.1 table)."""

    name: str
    path: Path
    header_bytes: int  # JSON header length (the 8-byte length prefix is not included)
    n_tensors: int
    data_bytes: int  # sum of tensor byte ranges == file size - 8 - header_bytes
    file_size: int


def read_safetensors_header(path: Union[str, Path]) -> Tuple[int, Dict[str, dict]]:
    """Return ``(header_len, header_json)`` of a safetensors file without reading any tensor data.

    Framing: 8-byte little-endian unsigned header length, then a UTF-8 JSON object whose keys are tensor names plus an
    optional ``__metadata__`` entry (safetensors format spec; ``weight-loader-spec.md`` §1.1).
    """
    path = Path(path)
    with open(path, "rb") as fh:
        prefix = fh.read(8)
        if len(prefix) != 8:
            raise ValueError(f"{path}: truncated safetensors header prefix")
        (header_len,) = struct.unpack("<Q", prefix)
        raw = fh.read(header_len)
        if len(raw) != header_len:
            raise ValueError(f"{path}: truncated safetensors header ({len(raw)} of {header_len} bytes)")
    header = json.loads(raw.decode("utf-8"))
    if not isinstance(header, dict):
        raise ValueError(f"{path}: safetensors header is not a JSON object")
    return header_len, header


class LazyCheckpoint:
    """Header-indexed, lazily read view of one GR00T checkpoint snapshot.

    ``LazyCheckpoint("n17")`` resolves the pinned HF snapshot (``configs.snapshot_dir``); a directory path is also
    accepted for ad-hoc snapshots.  Public API::

        ck = LazyCheckpoint("n15")
        ck.keys()                         # all 899 keys, shard-sorted then name-sorted
        ck.info(key)                      # TensorInfo(shape, dtype, shard, data_offsets)
        ck.shape(key), ck.dtype(key)      # convenience
        ck.shard_of(key)                  # 'model-00002-of-00003.safetensors'
        ck.shards_for_prefix(prefix)      # set of shards touched by keys with this prefix
        ck.get(key)                       # torch tensor as stored (bf16)
        ck.get(key, torch.float32)        # upcast on read
        ck.slot(key, e)                   # W[e] of a [32, ...] embodiment tensor, reads only that slot
        ck.rows(key, lo, hi)              # rows lo:hi of a 2-D table (embedding / position rows)
        ck.list_by_prefix(prefix)         # sorted keys starting with prefix
        ck.iter_by_prefix(prefix)         # (key, tensor) pairs, shard-sorted (never holds more than one tensor)
        ck.verify_headers()               # T2 of weight-loader-spec.md §9: byte accounting must close
    """

    def __init__(self, version_or_dir: Union[str, Path]):
        if isinstance(version_or_dir, str) and version_or_dir in GR00T_VERSIONS:
            self.version: Optional[str] = version_or_dir
            self.dir = snapshot_dir(version_or_dir)
        else:
            self.version = None
            self.dir = Path(version_or_dir)
            if not (self.dir / "model.safetensors.index.json").is_file():
                raise FileNotFoundError(f"{self.dir}: no model.safetensors.index.json")
        index_path = self.dir / "model.safetensors.index.json"
        with open(index_path) as fh:
            index = json.load(fh)
        if "weight_map" not in index:
            raise ValueError(f"{index_path}: missing weight_map")
        self.weight_map: Dict[str, str] = dict(index["weight_map"])
        self.metadata: Dict[str, int] = dict(index.get("metadata", {}))

        self._infos: Dict[str, TensorInfo] = {}
        self._shards: Dict[str, ShardInfo] = {}
        for shard_name in sorted(set(self.weight_map.values())):
            shard_path = self.dir / shard_name
            if not shard_path.is_file():
                raise FileNotFoundError(f"{self.dir}: shard {shard_name} listed in the index is missing")
            header_len, header = read_safetensors_header(shard_path)
            data_bytes = 0
            n_tensors = 0
            for key, entry in header.items():
                if key == "__metadata__":
                    continue
                if key in self._infos:
                    raise ValueError(f"{key}: present in two shards ({self._infos[key].shard}, {shard_name})")
                if self.weight_map.get(key) != shard_name:
                    raise ValueError(f"{key}: header says shard {shard_name}, index says {self.weight_map.get(key)}")
                if entry["dtype"] not in SAFETENSORS_DTYPES:
                    raise ValueError(f"{key}: unsupported safetensors dtype {entry['dtype']!r}")
                start, end = entry["data_offsets"]
                info = TensorInfo(
                    key=key,
                    shard=shard_name,
                    dtype=entry["dtype"],
                    shape=tuple(int(d) for d in entry["shape"]),
                    data_offsets=(int(start), int(end)),
                )
                expected = info.numel * _DTYPE_BYTES[info.dtype]
                if info.nbytes != expected:
                    raise ValueError(f"{key}: data_offsets span {info.nbytes} bytes, shape*dtype needs {expected}")
                self._infos[key] = info
                data_bytes += info.nbytes
                n_tensors += 1
            self._shards[shard_name] = ShardInfo(
                name=shard_name,
                path=shard_path,
                header_bytes=header_len,
                n_tensors=n_tensors,
                data_bytes=data_bytes,
                file_size=shard_path.stat().st_size,
            )
        missing = set(self.weight_map) - set(self._infos)
        if missing:
            raise ValueError(f"{len(missing)} index keys have no header entry, e.g. {sorted(missing)[:3]}")
        self._handles: Dict[str, object] = {}

    # ------------------------------------------------------------------ header-only queries
    def __len__(self) -> int:
        return len(self._infos)

    def __contains__(self, key: object) -> bool:
        return key in self._infos

    def keys(self) -> List[str]:
        """All keys, ordered by shard then name (the order a streaming loader should read them in, §1.2)."""
        return sorted(self._infos, key=lambda k: (self._infos[k].shard, k))

    def __iter__(self) -> Iterator[str]:
        return iter(self.keys())

    def info(self, key: str) -> TensorInfo:
        try:
            return self._infos[key]
        except KeyError:
            raise KeyError(f"{key!r} is not in the checkpoint ({self.dir.name})") from None

    def shape(self, key: str) -> Tuple[int, ...]:
        return self.info(key).shape

    def dtype(self, key: str) -> torch.dtype:
        return self.info(key).torch_dtype

    def shard_of(self, key: str) -> str:
        return self.info(key).shard

    @property
    def shards(self) -> Tuple[ShardInfo, ...]:
        return tuple(self._shards[n] for n in sorted(self._shards))

    @property
    def shard_names(self) -> Tuple[str, ...]:
        return tuple(sorted(self._shards))

    @property
    def total_bytes(self) -> int:
        return sum(s.data_bytes for s in self.shards)

    def list_by_prefix(self, prefix: str) -> List[str]:
        """Sorted keys that start with ``prefix``; raises if none match (typos must not pass silently)."""
        out = sorted(k for k in self._infos if k.startswith(prefix))
        if not out:
            raise KeyError(f"no key starts with {prefix!r} in {self.dir.name}")
        return out

    def shards_for_prefix(self, prefix: str) -> Tuple[str, ...]:
        """Shards touched by the keys under ``prefix`` (two for the straddling unit of each version)."""
        return tuple(sorted({self._infos[k].shard for k in self.list_by_prefix(prefix)}))

    def verify_headers(self) -> Dict[str, int]:
        """Byte accounting (``weight-loader-spec.md`` §9 T2): per shard ``8 + header + data == file size`` and the sum
        of all data bytes equals ``metadata.total_size``.  Returns the per-shard data byte counts."""
        for s in self.shards:
            if 8 + s.header_bytes + s.data_bytes != s.file_size:
                raise ValueError(f"{s.name}: 8 + {s.header_bytes} + {s.data_bytes} != file size {s.file_size}")
        total = self.total_bytes
        if "total_size" in self.metadata and int(self.metadata["total_size"]) != total:
            raise ValueError(f"sum of data bytes {total} != metadata.total_size {self.metadata['total_size']}")
        return {s.name: s.data_bytes for s in self.shards}

    # ------------------------------------------------------------------ data reads
    def _handle(self, shard: str):
        h = self._handles.get(shard)
        if h is None:
            h = safe_open(str(self._shards[shard].path), framework="pt", device="cpu")
            self._handles[shard] = h
        return h

    def get(self, key: str, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """Full tensor, as stored (bf16) or converted to ``dtype``."""
        info = self.info(key)
        t = self._handle(info.shard).get_tensor(key)
        if tuple(t.shape) != info.shape:
            raise ValueError(f"{key}: header shape {info.shape} != tensor shape {tuple(t.shape)}")
        return t if dtype is None else t.to(dtype)

    def get_fp32(self, key: str) -> torch.Tensor:
        return self.get(key, torch.float32)

    def slot(self, key: str, e: int, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """``W[e]`` of an embodiment-indexed ``[32, ...]`` tensor, reading only that slot (``weight-loader-spec.md``
        §1.2/§4.7: the slot is ``[in, out]`` and is consumed as ``x @ W[e] + b[e]`` -- do not transpose)."""
        info = self.info(key)
        if len(info.shape) < 2:
            raise ValueError(f"{key}: shape {info.shape} has no leading slot dimension")
        if not 0 <= e < info.shape[0]:
            raise IndexError(f"{key}: slot {e} out of range [0, {info.shape[0]})")
        t = self._handle(info.shard).get_slice(key)[e]
        return t if dtype is None else t.to(dtype)

    def rows(self, key: str, lo: int, hi: int, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """Rows ``lo:hi`` of a 2-D table (e.g. ``position_embedding.weight[0:H]``, ``embed_tokens`` gathers)."""
        info = self.info(key)
        if not (0 <= lo <= hi <= info.shape[0]):
            raise IndexError(f"{key}: rows [{lo}, {hi}) out of range for shape {info.shape}")
        t = self._handle(info.shard).get_slice(key)[lo:hi]
        return t if dtype is None else t.to(dtype)

    def iter_by_prefix(self, prefix: str, dtype: Optional[torch.dtype] = None) -> Iterator[Tuple[str, torch.Tensor]]:
        """Yield ``(key, tensor)`` for every key under ``prefix`` in shard order (one tensor in flight at a time)."""
        keys = sorted(self.list_by_prefix(prefix), key=lambda k: (self._infos[k].shard, k))
        for k in keys:
            yield k, self.get(k, dtype)

    def nbytes_by_prefix(self, prefix: str) -> int:
        return sum(self._infos[k].nbytes for k in self.list_by_prefix(prefix))

    def close(self) -> None:
        self._handles.clear()

    # ------------------------------------------------------------------ version-level facts
    def expected_counts(self) -> Optional[Tuple[int, int, int]]:
        """``(n_tensors, n_shards, n_bytes)`` from ``configs`` when the checkpoint is a known version, else None."""
        if self.version is None:
            return None
        c = get_config(self.version).checkpoint
        return c.n_tensors, c.n_shards, c.n_bytes

    def __repr__(self) -> str:
        return f"LazyCheckpoint({self.version or self.dir}, {len(self)} tensors, {len(self._shards)} shards)"


def straddling_prefix(version: str) -> str:
    """The layer/block whose tensors span two shards (``weight-loader-spec.md`` §1.1)."""
    return get_config(version).checkpoint.straddling_prefix


def canonical_key(version: str, key: str) -> str:
    """Strip the version-specific container prefix to the HF-relative canonical name used by the research tables
    (``weight-loader-spec.md`` §2 'strip prefix -> canonical' column).  Unknown prefixes raise."""
    c = get_config(version).checkpoint
    candidates: Sequence[Tuple[str, str]] = [
        (c.vision_prefix, ""),
        (c.llm_prefix, ""),
        (c.action_head_prefix, ""),
    ]
    if c.projector_prefix is not None:
        candidates = [(c.projector_prefix, "")] + list(candidates)
    for prefix, repl in candidates:
        if key.startswith(prefix):
            return repl + key[len(prefix) :]
    raise KeyError(f"{version}: {key!r} does not start with a known container prefix")


def iter_layer_keys(ck: LazyCheckpoint, prefix: str, indices: Iterable[int]) -> Iterator[Tuple[int, List[str]]]:
    """``(layer_index, sorted keys)`` for ``f"{prefix}{i}."`` -- the shard set is looked up per key by ``ck``."""
    for i in indices:
        yield i, ck.list_by_prefix(f"{prefix}{i}.")