"""Read the explicit shared-bank aliases and architecture blocks of a v2 checkpoint."""

import json
import math
import struct
from pathlib import Path

import torch
from safetensors import safe_open


class SharedBankCheckpoint:
    def __init__(self, directory):
        self.directory = Path(directory).resolve()
        index = json.loads((self.directory / "model.safetensors.index.json").read_text())
        self.metadata = index["metadata"]
        if self.metadata["format"] != "the-net-shared-aliases-v2":
            raise ValueError("Expected the-net-shared-aliases-v2")
        self.weight_map = index["weight_map"]
        self.headers = {}

    def _resolve(self, name):
        filename = self.weight_map[name]
        key = name
        if name.startswith(self.metadata["alias_prefix"]):
            filename = self.metadata["canonical_bank"]
            key = self.metadata["alias_tensor"]
        path = (self.directory / filename).resolve()
        if not path.is_relative_to(self.directory):
            raise ValueError("Checkpoint paths must stay inside the checkpoint directory")
        if path not in self.headers:
            with safe_open(path, framework="pt"):
                pass
            with path.open("rb") as stream:
                length = struct.unpack("<Q", stream.read(8))[0]
                self.headers[path] = (8 + length, json.loads(stream.read(length)))
        base, header = self.headers[path]
        return path, base, header[key]

    def tensor_info(self, name):
        _, _, entry = self._resolve(name)
        return {"dtype": entry["dtype"], "shape": list(entry["shape"])}

    def stored_slice(self, name, start, length):
        path, base, entry = self._resolve(name)
        size = math.prod(entry["shape"])
        if type(start) is not int or type(length) is not int:
            raise ValueError("Slice bounds must be integers")
        if start < 0 or length < 0 or start + length > size or length > 1_000_000:
            raise ValueError("Invalid slice or slice exceeds one million stored elements")
        dtype, width = {"U8": (torch.uint8, 1), "F16": (torch.float16, 2)}[entry["dtype"]]
        if length == 0:
            return torch.empty(0, dtype=dtype)
        with path.open("rb") as stream:
            stream.seek(base + entry["data_offsets"][0] + start * width)
            data = bytearray(stream.read(length * width))
        if len(data) != length * width:
            raise ValueError("Truncated checkpoint payload")
        return torch.frombuffer(data, dtype=dtype).clone()

    def bit_slice(self, name, start, length):
        info = self.tensor_info(name)
        if info["dtype"] != "U8":
            raise ValueError("Bit slices require a U8 shared bank")
        if type(start) is not int or type(length) is not int:
            raise ValueError("Slice bounds must be integers")
        if start < 0 or length < 0 or start + length > math.prod(info["shape"]) * 8:
            raise ValueError("Invalid bit slice")
        if length > 1_000_000:
            raise ValueError("Read at most one million bits per slice")
        if length == 0:
            return torch.empty(0, dtype=torch.uint8)
        offset = start % 8
        data = self.stored_slice(name, start // 8, (offset + length + 7) // 8)
        positions = torch.arange(offset, offset + length)
        return ((data[positions // 8].long() >> (positions % 8)) & 1).to(torch.uint8)