"""Lazy, read-only Release 188 NoNE page storage.

The public release is one standard safetensors weight map.  This module reads
only that small JSON index and the published direct-page index at startup.
Multi-gigabyte page shards are opened only when the model routes one of their
page IDs.
"""

from __future__ import annotations

import mmap
import re
from pathlib import Path, PurePosixPath
from typing import Final, Iterator

import torch
from safetensors import safe_open

from resynthesis.none_paging import (
    NoNEPageRequestPacket,
    NoNEPageWeights,
    validate_page_weights,
)

RELEASE_188_GENERATION: Final[int] = 188
_DIRECT_INDEX_PATH: Final[str] = (
    "weights/safetensors/direct-page-index.safetensors"
)
_PAGE_SHARD_PATTERN: Final[re.Pattern[str]] = re.compile(
    r"weights/safetensors/pages-(\d{5})-of-(\d{5})\.safetensors"
)
_PAGE_TENSOR_COMPONENTS: Final[tuple[str, ...]] = (
    "down_t",
    "ffn_mode_t",
    "format_revision_t",
    "gate_t",
    "glyph_down_t",
    "glyph_up_t",
    "outcome_memory_t",
    "page_ids_t",
    "repair_memory_t",
    "transfer_memory_t",
    "translation_gate_t",
    "up_t",
)
_PAGE_COMPONENT_BITS: Final[dict[str, int]] = {
    name: 1 << index for index, name in enumerate(_PAGE_TENSOR_COMPONENTS)
}
_COMPLETE_PAGE_COMPONENT_MASK: Final[int] = (
    (1 << len(_PAGE_TENSOR_COMPONENTS)) - 1
)
_SELF_CONTAINED_PAGE_REVISIONS: Final[frozenset[int]] = frozenset((2, 3, 7))
_RELEASE_SHARD_SCHEMA: Final[str] = (
    "nnf.resynthesis.release-page-shard.v1"
)
_WHITESPACE: Final[bytes] = b" \t\r\n"


def _skip_whitespace(payload: mmap.mmap, position: int) -> int:
    while position < len(payload) and payload[position] in _WHITESPACE:
        position += 1
    return position


def _ascii_json_string(
    payload: mmap.mmap,
    position: int,
) -> tuple[str, int]:
    """Read one unescaped ASCII string from the generated public index."""

    if position >= len(payload) or payload[position] != ord('"'):
        raise RuntimeError("Release 188 weight-map string is malformed")
    start = position + 1
    position = start
    while position < len(payload):
        value = payload[position]
        if value == ord('"'):
            try:
                return payload[start:position].decode("ascii"), position + 1
            except UnicodeDecodeError as error:
                raise RuntimeError(
                    "Release 188 weight-map string is not ASCII"
                ) from error
        if value == ord("\\") or value < 0x20 or value > 0x7E:
            raise RuntimeError(
                "Release 188 weight-map contains an escaped or non-ASCII string"
            )
        position += 1
    raise RuntimeError("Release 188 weight-map string is truncated")


def _weight_map_pairs(index_path: Path) -> Iterator[tuple[str, str]]:
    """Stream the flat weight map without constructing a million-entry dict."""

    with index_path.open("rb") as handle:
        with mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) as payload:
            marker = b'"weight_map"'
            marker_position = payload.find(marker)
            if marker_position < 0:
                raise RuntimeError("Release 188 index omits its weight map")
            position = _skip_whitespace(
                payload,
                marker_position + len(marker),
            )
            if position >= len(payload) or payload[position] != ord(":"):
                raise RuntimeError("Release 188 weight map is malformed")
            position = _skip_whitespace(payload, position + 1)
            if position >= len(payload) or payload[position] != ord("{"):
                raise RuntimeError("Release 188 weight map is not an object")
            position += 1
            while True:
                position = _skip_whitespace(payload, position)
                if position >= len(payload):
                    raise RuntimeError("Release 188 weight map is truncated")
                if payload[position] == ord("}"):
                    break
                name, position = _ascii_json_string(payload, position)
                position = _skip_whitespace(payload, position)
                if position >= len(payload) or payload[position] != ord(":"):
                    raise RuntimeError(
                        "Release 188 weight-map entry is malformed"
                    )
                position = _skip_whitespace(payload, position + 1)
                shard, position = _ascii_json_string(payload, position)
                yield name, shard
                position = _skip_whitespace(payload, position)
                if position >= len(payload):
                    raise RuntimeError("Release 188 weight map is truncated")
                if payload[position] == ord(","):
                    position += 1
                    continue
                if payload[position] == ord("}"):
                    break
                raise RuntimeError("Release 188 weight-map separator differs")


def _resolved_repo_file(repository_root: Path, relative_path: str) -> Path:
    relative = PurePosixPath(relative_path)
    if relative.is_absolute() or ".." in relative.parts:
        raise RuntimeError("Release 188 weight path escapes its repository")
    resolved = repository_root.joinpath(*relative.parts).resolve()
    try:
        resolved.relative_to(repository_root)
    except ValueError as error:
        raise RuntimeError(
            "Release 188 weight path escapes its repository"
        ) from error
    return resolved


class Release188SafetensorsPageStore:
    """Read-only page store compatible with ``NoNEPagedExpertRuntime``."""

    def __init__(self, index_path: Path) -> None:
        resolved_index = index_path.expanduser().resolve()
        if (
            resolved_index.name != "model.safetensors.index.json"
            or not resolved_index.is_file()
        ):
            raise RuntimeError(
                "Release 188 repository-root safetensors index is missing"
            )
        self.index_path = resolved_index
        self.root = resolved_index.parent
        self._page_shard_by_id: dict[int, Path] = {}
        component_masks: dict[int, int] = {}
        direct_index_names: set[str] = set()
        shard_ordinals: set[int] = set()
        shard_count: int | None = None

        for tensor_name, relative_shard in _weight_map_pairs(resolved_index):
            if relative_shard == _DIRECT_INDEX_PATH:
                direct_index_names.add(tensor_name)
            page_stem, separator, component = tensor_name.partition(".")
            page_digits = page_stem.removeprefix("page_")
            if (
                not tensor_name.startswith("page_")
                or not separator
                or not page_digits.isdigit()
            ):
                continue
            component_bit = _PAGE_COMPONENT_BITS.get(component)
            shard_match = _PAGE_SHARD_PATTERN.fullmatch(relative_shard)
            if (
                component_bit is None
                or shard_match is None
            ):
                raise RuntimeError(
                    "Release 188 page tensor mapping is malformed"
                )
            page_id = int(page_digits)
            ordinal = int(shard_match.group(1))
            observed_count = int(shard_match.group(2))
            if (
                ordinal < 1
                or observed_count < 1
                or ordinal > observed_count
                or (
                    shard_count is not None
                    and observed_count != shard_count
                )
            ):
                raise RuntimeError("Release 188 page shard sequence differs")
            shard_count = observed_count
            shard_ordinals.add(ordinal)
            shard_path = _resolved_repo_file(self.root, relative_shard)
            previous_shard = self._page_shard_by_id.setdefault(
                page_id,
                shard_path,
            )
            if previous_shard != shard_path:
                raise RuntimeError(
                    "Release 188 page tensors cross physical shards"
                )
            prior_mask = component_masks.get(page_id, 0)
            if prior_mask & component_bit:
                raise RuntimeError(
                    "Release 188 page tensor mapping contains a duplicate"
                )
            component_masks[page_id] = prior_mask | component_bit

        required_direct_names = {
            "format_revisions_t",
            "page_count_t",
            "page_ids_t",
        }
        if not required_direct_names.issubset(direct_index_names):
            raise RuntimeError(
                "Release 188 weight map omits its direct page index"
            )
        if (
            not component_masks
            or any(
                mask != _COMPLETE_PAGE_COMPONENT_MASK
                for mask in component_masks.values()
            )
            or shard_count is None
            or shard_ordinals != set(range(1, shard_count + 1))
        ):
            raise RuntimeError(
                "Release 188 weight map has incomplete page tensors"
            )

        direct_index_path = _resolved_repo_file(
            self.root,
            _DIRECT_INDEX_PATH,
        )
        if not direct_index_path.is_file():
            raise RuntimeError("Release 188 direct page index is missing")
        with safe_open(  # type: ignore[no-untyped-call]
            str(direct_index_path),
            framework="pt",
            device="cpu",
        ) as direct_index:
            page_ids_t: torch.Tensor = direct_index.get_tensor(
                "page_ids_t"
            ).long()
            revisions_t: torch.Tensor = direct_index.get_tensor(
                "format_revisions_t"
            ).long()
            page_count_t: torch.Tensor = direct_index.get_tensor(
                "page_count_t"
            ).long()
        if (
            page_ids_t.ndim != 1
            or page_ids_t.numel() < 1
            or revisions_t.shape != page_ids_t.shape
            or page_count_t.shape != (1,)
            or int(page_count_t[0]) != page_ids_t.shape[0]
            or torch.unique(page_ids_t).numel() != page_ids_t.numel()
            or page_ids_t.lt(0).any()
        ):
            raise RuntimeError("Release 188 direct page index is malformed")
        page_ids = tuple(int(value) for value in page_ids_t.tolist())
        revisions = tuple(int(value) for value in revisions_t.tolist())
        if (
            set(page_ids) != set(self._page_shard_by_id)
            or any(
                revision not in _SELF_CONTAINED_PAGE_REVISIONS
                for revision in revisions
            )
        ):
            raise RuntimeError(
                "Release 188 direct index and page shards differ"
            )
        self._page_catalog_ids_t = page_ids_t.detach().cpu().clone()
        self._format_revision_by_id = dict(
            zip(page_ids, revisions, strict=True)
        )
        self._session_id_t: torch.Tensor | None = None
        self._accepted_generation_t = torch.tensor(
            RELEASE_188_GENERATION,
            dtype=torch.long,
        )

    @property
    def page_catalog_ids_t(self) -> torch.Tensor:
        """Return the complete public page catalog as a tensor."""

        return self._page_catalog_ids_t.clone()

    def begin_session(self, session_id_t: torch.Tensor) -> torch.Tensor:
        """Bind caller-owned inference session identity to immutable release 188."""

        resolved = session_id_t.detach().cpu().long().reshape(-1)
        if (
            session_id_t.dtype != torch.long
            or session_id_t.ndim != 1
            or resolved.numel() < 1
        ):
            raise ValueError("Release 188 session identity is malformed")
        self._session_id_t = resolved.clone()
        return self._accepted_generation_t.clone()

    def accepted_generation_t(self) -> torch.Tensor:
        """Return the immutable accepted generation for the active session."""

        if self._session_id_t is None:
            raise RuntimeError("Release 188 page store has no active session")
        return self._accepted_generation_t.clone()

    def accepted_page_ids_t_boundary(self) -> torch.Tensor:
        """Return the complete accepted physical page identity in sorted order."""

        if self._session_id_t is None:
            raise RuntimeError("Release 188 page store has no active session")
        return torch.sort(self._page_catalog_ids_t).values.clone()

    def _validated_request_page_ids_boundary(
        self,
        *,
        session_id_t: torch.Tensor,
        generation_t: torch.Tensor,
        page_ids_t: torch.Tensor,
    ) -> tuple[int, ...]:
        active_session_t = self._session_id_t
        requested_session_t = session_id_t.detach().cpu().long().reshape(-1)
        requested_generation_t = (
            generation_t.detach().cpu().long().reshape(-1)
        )
        requested_page_ids_t = (
            page_ids_t.detach().cpu().long().reshape(-1)
        )
        if (
            active_session_t is None
            or session_id_t.dtype != torch.long
            or session_id_t.ndim != 1
            or not torch.equal(active_session_t, requested_session_t)
        ):
            raise RuntimeError("Release 188 page request crossed its session")
        if (
            generation_t.dtype != torch.long
            or requested_generation_t.shape != (1,)
            or not torch.equal(
                requested_generation_t,
                self._accepted_generation_t.reshape(1),
            )
        ):
            raise RuntimeError("Release 188 page request generation differs")
        if (
            page_ids_t.dtype != torch.long
            or page_ids_t.ndim != 1
            or requested_page_ids_t.numel() < 1
            or torch.unique(requested_page_ids_t).numel()
            != requested_page_ids_t.numel()
        ):
            raise ValueError("Release 188 page request is malformed")
        requested_page_ids = tuple(
            int(value) for value in requested_page_ids_t.tolist()
        )
        if any(
            page_id not in self._page_shard_by_id
            for page_id in requested_page_ids
        ):
            raise KeyError("Release 188 page request is outside its catalog")
        return requested_page_ids

    def _load_page_row_boundary(
        self,
        handle: object,
        *,
        page_id: int,
    ) -> NoNEPageWeights:
        prefix = f"page_{page_id:06d}."
        keys = set(handle.keys())  # type: ignore[attr-defined]
        expected_keys = {
            f"{prefix}{component}" for component in _PAGE_TENSOR_COMPONENTS
        }
        if {name for name in keys if name.startswith(prefix)} != expected_keys:
            raise RuntimeError("Release 188 routed page tensor schema differs")

        def tensor(component: str) -> torch.Tensor:
            return handle.get_tensor(  # type: ignore[attr-defined,no-any-return]
                f"{prefix}{component}"
            )

        stored_page_ids_t = tensor("page_ids_t").long()
        stored_revision_t = tensor("format_revision_t").long()
        if (
            stored_page_ids_t.shape != (1,)
            or int(stored_page_ids_t[0]) != page_id
            or stored_revision_t.shape != (1,)
            or int(stored_revision_t[0])
            != self._format_revision_by_id[page_id]
        ):
            raise RuntimeError("Release 188 routed page identity differs")
        weights = NoNEPageWeights(
            page_ids_t=stored_page_ids_t,
            ffn_mode_t=tensor("ffn_mode_t"),
            gate_t=tensor("gate_t"),
            up_t=tensor("up_t"),
            down_t=tensor("down_t"),
            glyph_down_t=tensor("glyph_down_t"),
            glyph_up_t=tensor("glyph_up_t"),
            translation_gate_t=tensor("translation_gate_t"),
            outcome_memory_t=tensor("outcome_memory_t"),
            repair_memory_t=tensor("repair_memory_t"),
            transfer_memory_t=tensor("transfer_memory_t"),
        )
        validate_page_weights(weights, synchronize_tensor_values=False)
        return weights

    def materialize_page_ids_boundary(
        self,
        *,
        session_id_t: torch.Tensor,
        generation_t: torch.Tensor,
        page_ids_t: torch.Tensor,
        device: torch.device,
        dtype: torch.dtype,
        trainable: bool,
    ) -> NoNEPageWeights:
        """Materialize only model-routed page rows from their mapped shards."""

        if trainable:
            raise RuntimeError(
                "Release 188 public safetensors page store is read-only"
            )
        if not isinstance(device, torch.device):
            raise TypeError("Release 188 page device must be torch.device")
        if not dtype.is_floating_point:
            raise TypeError("Release 188 page dtype must be floating point")
        requested_page_ids = self._validated_request_page_ids_boundary(
            session_id_t=session_id_t,
            generation_t=generation_t,
            page_ids_t=page_ids_t,
        )
        positions_by_shard: dict[Path, list[tuple[int, int]]] = {}
        for position, page_id in enumerate(requested_page_ids):
            positions_by_shard.setdefault(
                self._page_shard_by_id[page_id],
                [],
            ).append((position, page_id))
        rows: list[NoNEPageWeights | None] = [None] * len(requested_page_ids)
        for shard_path, routed_rows in positions_by_shard.items():
            if not shard_path.is_file():
                raise RuntimeError(
                    f"Release 188 routed page shard is missing: {shard_path}"
                )
            with safe_open(  # type: ignore[no-untyped-call]
                str(shard_path),
                framework="pt",
                device="cpu",
            ) as handle:
                metadata = handle.metadata()
                if (
                    metadata is None
                    or metadata.get("schema") != _RELEASE_SHARD_SCHEMA
                    or metadata.get("release") != str(RELEASE_188_GENERATION)
                ):
                    raise RuntimeError(
                        "Release 188 routed page shard metadata differs"
                    )
                for position, page_id in routed_rows:
                    rows[position] = self._load_page_row_boundary(
                        handle,
                        page_id=page_id,
                    )
        if any(row is None for row in rows):
            raise RuntimeError("Release 188 page materialization is incomplete")
        resolved_rows = tuple(row for row in rows if row is not None)
        weights = (
            resolved_rows[0]
            if len(resolved_rows) == 1
            else NoNEPageWeights(
                page_ids_t=torch.cat(
                    tuple(row.page_ids_t for row in resolved_rows)
                ),
                ffn_mode_t=torch.cat(
                    tuple(row.ffn_mode_t for row in resolved_rows)
                ),
                gate_t=torch.cat(tuple(row.gate_t for row in resolved_rows)),
                up_t=torch.cat(tuple(row.up_t for row in resolved_rows)),
                down_t=torch.cat(tuple(row.down_t for row in resolved_rows)),
                glyph_down_t=torch.cat(
                    tuple(row.glyph_down_t for row in resolved_rows)
                ),
                glyph_up_t=torch.cat(
                    tuple(row.glyph_up_t for row in resolved_rows)
                ),
                translation_gate_t=torch.cat(
                    tuple(row.translation_gate_t for row in resolved_rows)
                ),
                outcome_memory_t=torch.cat(
                    tuple(row.outcome_memory_t for row in resolved_rows)
                ),
                repair_memory_t=torch.cat(
                    tuple(row.repair_memory_t for row in resolved_rows)
                ),
                transfer_memory_t=torch.cat(
                    tuple(row.transfer_memory_t for row in resolved_rows)
                ),
            )
        )
        return weights.to(
            device=device,
            dtype=dtype,
            trainable=False,
        )

    def materialize_weights(
        self,
        request: NoNEPageRequestPacket,
        *,
        device: torch.device,
        dtype: torch.dtype,
        trainable: bool,
    ) -> NoNEPageWeights:
        """Compatibility adapter for the runtime's tensor request packet."""

        return self.materialize_page_ids_boundary(
            session_id_t=request.session_id_t,
            generation_t=request.generation_t,
            page_ids_t=request.unique_page_ids_t,
            device=device,
            dtype=dtype,
            trainable=trainable,
        )