#!/usr/bin/env python3
"""Minimal safetensors and GGUF v3 I/O used by the reproducible build."""

from __future__ import annotations

import hashlib
import json
import mmap
import os
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Callable, Iterable

import numpy as np

from recipe import (
    GGML_F16,
    GGML_F32,
    GGML_MXFP4,
    GGML_Q2_K,
    GGML_Q8_0,
    GGUF_METADATA,
    SOURCE_FILES,
    TensorRecipe,
    source_expectations,
)


ALIGNMENT = 32
GGUF_VALUE_U32 = 4
GGUF_VALUE_I32 = 5
GGUF_VALUE_STRING = 8
GGUF_VALUE_ARRAY = 9

_DTYPE_BYTES = {
    "F8_E4M3": 1,
    "F8_E8M0": 1,
    "I8": 1,
    "BF16": 2,
    "F16": 2,
    "F32": 4,
}


def align_up(value: int, alignment: int = ALIGNMENT) -> int:
    return (value + alignment - 1) // alignment * alignment


def sha256_path(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(chunk_size):
            digest.update(chunk)
    return digest.hexdigest()


def verify_source_files(source_dir: Path) -> dict[str, dict[str, int | str]]:
    observed: dict[str, dict[str, int | str]] = {}
    for name, expected in SOURCE_FILES.items():
        path = source_dir / name
        if not path.is_file():
            raise FileNotFoundError(f"missing pinned source shard: {path}")
        size = path.stat().st_size
        if size != expected["size"]:
            raise ValueError(f"{name}: expected {expected['size']} bytes, found {size}")
        digest = sha256_path(path)
        if digest != expected["sha256"]:
            raise ValueError(
                f"{name}: expected SHA-256 {expected['sha256']}, found {digest}"
            )
        observed[name] = {"size": size, "sha256": digest}
    return observed


@dataclass(frozen=True)
class SafeTensorInfo:
    dtype: str
    shape: tuple[int, ...]
    start: int
    end: int
    shard_index: int


class SafeTensorSet:
    """Read-only mmap index spanning the three pinned safetensors shards."""

    def __init__(self, paths: Iterable[Path]):
        self._files: list[BinaryIO] = []
        self._maps: list[mmap.mmap] = []
        self._index: dict[str, SafeTensorInfo] = {}
        try:
            for shard_index, path in enumerate(paths):
                handle = path.open("rb")
                mapping = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ)
                self._files.append(handle)
                self._maps.append(mapping)
                self._index_shard(path, shard_index, mapping)
        except Exception:
            self.close()
            raise

    def _index_shard(self, path: Path, shard_index: int, mapping: mmap.mmap) -> None:
        if len(mapping) < 8:
            raise ValueError(f"{path}: shorter than safetensors header length")
        header_len = struct.unpack_from("<Q", mapping, 0)[0]
        data_start = 8 + header_len
        if data_start > len(mapping):
            raise ValueError(f"{path}: safetensors header overruns file")
        try:
            header = json.loads(mapping[8:data_start])
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise ValueError(f"{path}: invalid safetensors header: {error}") from error
        if not isinstance(header, dict):
            raise ValueError(f"{path}: safetensors header is not an object")
        for name, raw in header.items():
            if name == "__metadata__" or not name.startswith("mtp."):
                continue
            if name in self._index:
                raise ValueError(f"duplicate tensor across shards: {name}")
            if not isinstance(raw, dict):
                raise ValueError(f"{path}: tensor {name} metadata is not an object")
            dtype = raw.get("dtype")
            shape = raw.get("shape")
            offsets = raw.get("data_offsets")
            if (
                not isinstance(dtype, str)
                or not isinstance(shape, list)
                or not all(isinstance(dim, int) and dim >= 0 for dim in shape)
                or not isinstance(offsets, list)
                or len(offsets) != 2
                or not all(
                    isinstance(offset, int) and offset >= 0 for offset in offsets
                )
            ):
                raise ValueError(f"{path}: malformed metadata for tensor {name}")
            relative_start, relative_end = offsets
            start = data_start + relative_start
            end = data_start + relative_end
            if relative_start > relative_end or end > len(mapping):
                raise ValueError(f"{path}: tensor {name} data range overruns shard")
            if dtype in _DTYPE_BYTES:
                count = 1
                for dim in shape:
                    count *= dim
                expected_bytes = count * _DTYPE_BYTES[dtype]
                if end - start != expected_bytes:
                    raise ValueError(
                        f"{path}: tensor {name} needs {expected_bytes} bytes, "
                        f"data range contains {end - start}"
                    )
            self._index[name] = SafeTensorInfo(
                dtype=dtype,
                shape=tuple(shape),
                start=start,
                end=end,
                shard_index=shard_index,
            )

    def close(self) -> None:
        for mapping in reversed(getattr(self, "_maps", [])):
            try:
                mapping.close()
            except BufferError:
                pass
        for handle in reversed(getattr(self, "_files", [])):
            handle.close()
        self._maps = []
        self._files = []

    def __enter__(self) -> "SafeTensorSet":
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    def info(self, name: str) -> SafeTensorInfo:
        try:
            return self._index[name]
        except KeyError as error:
            raise KeyError(f"source tensor not found: {name}") from error

    def raw(self, name: str) -> memoryview:
        info = self.info(name)
        return memoryview(self._maps[info.shard_index])[info.start : info.end]

    def array(self, name: str, dtype: np.dtype | str) -> np.ndarray:
        info = self.info(name)
        return np.frombuffer(
            self._maps[info.shard_index],
            dtype=dtype,
            count=(info.end - info.start) // np.dtype(dtype).itemsize,
            offset=info.start,
        ).reshape(info.shape)

    def validate_plan(self, plan: Iterable[TensorRecipe]) -> None:
        expected_names: set[str] = set()
        for tensor in plan:
            for name, dtype, shape in source_expectations(tensor):
                if name in expected_names:
                    raise ValueError(
                        f"recipe consumes source tensor more than once: {name}"
                    )
                expected_names.add(name)
                info = self.info(name)
                if info.dtype != dtype:
                    raise ValueError(
                        f"{name}: expected dtype {dtype}, found {info.dtype}"
                    )
                if info.shape != shape:
                    raise ValueError(
                        f"{name}: expected shape {shape}, found {info.shape}"
                    )
        observed_names = self.names()
        if observed_names != expected_names:
            missing = sorted(expected_names - observed_names)
            unexpected = sorted(observed_names - expected_names)
            raise ValueError(
                "source tensor inventory differs from the recipe: "
                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
            )

    def names(self) -> frozenset[str]:
        return frozenset(self._index)


def _write_string(handle: BinaryIO, value: str) -> None:
    encoded = value.encode("utf-8")
    handle.write(struct.pack("<Q", len(encoded)))
    handle.write(encoded)


def write_gguf(
    path: Path,
    plan: tuple[TensorRecipe, ...],
    produce: Callable[[TensorRecipe, BinaryIO], int],
) -> None:
    """Write deterministic GGUF v3 data, streaming each tensor in plan order."""

    with path.open("wb", buffering=8 * 1024 * 1024) as handle:
        handle.write(b"GGUF")
        handle.write(struct.pack("<IQQ", 3, len(plan), len(GGUF_METADATA)))
        for key, kind, value in GGUF_METADATA:
            _write_string(handle, key)
            if kind == "u32":
                handle.write(struct.pack("<II", GGUF_VALUE_U32, int(value)))
            elif kind == "string":
                handle.write(struct.pack("<I", GGUF_VALUE_STRING))
                _write_string(handle, str(value))
            elif kind == "array_i32":
                values = tuple(int(item) for item in value)
                handle.write(
                    struct.pack("<IIQ", GGUF_VALUE_ARRAY, GGUF_VALUE_I32, len(values))
                )
                for item in values:
                    handle.write(struct.pack("<i", item))
            else:
                raise ValueError(f"unsupported GGUF metadata kind: {kind}")

        relative_offset = 0
        relative_offsets: list[int] = []
        for tensor in plan:
            relative_offsets.append(relative_offset)
            relative_offset = align_up(relative_offset + tensor.byte_len)
        for tensor, offset in zip(plan, relative_offsets, strict=True):
            _write_string(handle, tensor.out_name)
            handle.write(struct.pack("<I", len(tensor.dims)))
            for dim in tensor.dims:
                handle.write(struct.pack("<Q", dim))
            handle.write(struct.pack("<IQ", tensor.ggml_type, offset))

        data_start = align_up(handle.tell())
        handle.write(b"\0" * (data_start - handle.tell()))
        relative_written = 0
        for tensor in plan:
            before = handle.tell()
            produced = produce(tensor, handle)
            observed = handle.tell() - before
            if produced != observed or observed != tensor.byte_len:
                raise ValueError(
                    f"{tensor.out_name}: expected {tensor.byte_len} bytes, "
                    f"producer reported {produced} and wrote {observed}"
                )
            relative_written += observed
            padded = align_up(relative_written)
            handle.write(b"\0" * (padded - relative_written))
            relative_written = padded
        handle.flush()
        os.fsync(handle.fileno())


@dataclass(frozen=True)
class GgufTensor:
    name: str
    dims: tuple[int, ...]
    ggml_type: int
    relative_offset: int
    byte_len: int
    absolute_offset: int


class GgufReader:
    def __init__(self, path: Path):
        self.path = path
        self._file = path.open("rb")
        self._map = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ)
        self.metadata: dict[str, object] = {}
        self.tensors: tuple[GgufTensor, ...] = ()
        self.directory_end = 0
        self.data_start = 0
        try:
            self._parse()
        except Exception:
            self.close()
            raise

    def close(self) -> None:
        try:
            self._map.close()
        except BufferError:
            pass
        self._file.close()

    def __enter__(self) -> "GgufReader":
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    def _unpack(self, fmt: str, cursor: int) -> tuple[tuple[object, ...], int]:
        size = struct.calcsize(fmt)
        if cursor + size > len(self._map):
            raise ValueError("GGUF directory is truncated")
        return struct.unpack_from(fmt, self._map, cursor), cursor + size

    def _string(self, cursor: int) -> tuple[str, int]:
        (length,), cursor = self._unpack("<Q", cursor)
        end = cursor + int(length)
        if end > len(self._map):
            raise ValueError("GGUF string overruns file")
        try:
            value = self._map[cursor:end].decode("utf-8")
        except UnicodeDecodeError as error:
            raise ValueError(f"GGUF string is not UTF-8: {error}") from error
        return value, end

    def _parse(self) -> None:
        if self._map[:4] != b"GGUF":
            raise ValueError("not a GGUF file")
        (version, tensor_count, kv_count), cursor = self._unpack("<IQQ", 4)
        if version != 3:
            raise ValueError(f"expected GGUF v3, found v{version}")
        for _ in range(int(kv_count)):
            key, cursor = self._string(cursor)
            (value_type,), cursor = self._unpack("<I", cursor)
            if value_type == GGUF_VALUE_U32:
                (value,), cursor = self._unpack("<I", cursor)
            elif value_type == GGUF_VALUE_STRING:
                value, cursor = self._string(cursor)
            elif value_type == GGUF_VALUE_ARRAY:
                (element_type, length), cursor = self._unpack("<IQ", cursor)
                if element_type != GGUF_VALUE_I32:
                    raise ValueError(f"unsupported GGUF array type {element_type}")
                values = []
                for _ in range(int(length)):
                    (item,), cursor = self._unpack("<i", cursor)
                    values.append(item)
                value = tuple(values)
            else:
                raise ValueError(f"unsupported GGUF metadata type {value_type}")
            if key in self.metadata:
                raise ValueError(f"duplicate GGUF metadata key {key}")
            self.metadata[key] = value

        directory: list[tuple[str, tuple[int, ...], int, int]] = []
        for _ in range(int(tensor_count)):
            name, cursor = self._string(cursor)
            (ndim,), cursor = self._unpack("<I", cursor)
            dims = []
            for _ in range(int(ndim)):
                (dim,), cursor = self._unpack("<Q", cursor)
                dims.append(int(dim))
            (ggml_type, relative_offset), cursor = self._unpack("<IQ", cursor)
            directory.append((name, tuple(dims), int(ggml_type), int(relative_offset)))

        self.directory_end = cursor
        self.data_start = align_up(cursor)
        tensors: list[GgufTensor] = []
        seen: set[str] = set()
        previous_end = self.data_start
        for name, dims, ggml_type, relative_offset in directory:
            if name in seen:
                raise ValueError(f"duplicate GGUF tensor {name}")
            seen.add(name)
            byte_len = ggml_nbytes(dims, ggml_type)
            absolute_offset = self.data_start + relative_offset
            if absolute_offset % ALIGNMENT:
                raise ValueError(f"tensor {name} is not {ALIGNMENT}-byte aligned")
            if absolute_offset < previous_end:
                raise ValueError(f"tensor {name} overlaps the previous tensor")
            if absolute_offset + byte_len > len(self._map):
                raise ValueError(f"tensor {name} overruns GGUF file")
            tensors.append(
                GgufTensor(
                    name=name,
                    dims=dims,
                    ggml_type=ggml_type,
                    relative_offset=relative_offset,
                    byte_len=byte_len,
                    absolute_offset=absolute_offset,
                )
            )
            previous_end = absolute_offset + byte_len
        self.tensors = tuple(tensors)

    def tensor(self, name: str) -> GgufTensor:
        for tensor in self.tensors:
            if tensor.name == name:
                return tensor
        raise KeyError(f"GGUF tensor not found: {name}")

    def tensor_array(self, tensor: GgufTensor, dtype: np.dtype | str) -> np.ndarray:
        return np.frombuffer(
            self._map,
            dtype=dtype,
            count=tensor.byte_len // np.dtype(dtype).itemsize,
            offset=tensor.absolute_offset,
        )

    def bytes_at(self, start: int, end: int) -> bytes:
        if start < 0 or start > end or end > len(self._map):
            raise ValueError("GGUF byte range is outside the file")
        return self._map[start:end]


def ggml_nbytes(dims: tuple[int, ...], ggml_type: int) -> int:
    elements = 1
    for dim in dims:
        elements *= dim
    if ggml_type == GGML_F32:
        return elements * 4
    if ggml_type == GGML_F16:
        return elements * 2
    if ggml_type == GGML_Q8_0:
        if elements % 32:
            raise ValueError("Q8_0 tensor is not block aligned")
        return elements // 32 * 34
    if ggml_type == GGML_Q2_K:
        if elements % 256:
            raise ValueError("Q2_K tensor is not block aligned")
        return elements // 256 * 84
    if ggml_type == GGML_MXFP4:
        if elements % 32:
            raise ValueError("MXFP4 tensor is not block aligned")
        return elements // 32 * 17
    raise ValueError(f"unsupported GGML tensor type {ggml_type}")