runtime/src/resynthesis/packed_token_rows.py

"""Shared mmap-backed packed training row boundary.

This module is import-safe for both the learn loop and the corpus-training
readers: it owns no state and imports nothing from either, so neither side can
create an import cycle by returning these rows across the I/O boundary.
"""

from __future__ import annotations

import hashlib
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast

import torch


_TARGET_FREE_CALIBRATION_INPUT_IDS_DOMAIN = (
    b"nnf.resynthesis.target_free_calibration_input_ids.v1\x00"
)
_TARGET_FREE_CALIBRATION_INPUT_MASK_DOMAIN = (
    b"nnf.resynthesis.target_free_calibration_input_mask.v1\x00"
)
_TARGET_FREE_CALIBRATION_PACKET_DOMAIN = (
    b"nnf.resynthesis.target_free_calibration_prompt_packet.v1\x00"
)


def _domain_separated_tensor_sha256_t_boundary(
    value_t: torch.Tensor,
    *,
    domain: bytes,
) -> torch.Tensor:
    """Hash one exact tensor at an explicit authority boundary.

    Calibration, cold proof, and acceptance must compare one identity domain.
    Including dtype and geometry prevents equal raw bytes under a different
    tensor interpretation from acquiring the same authority.
    """

    if (
        not isinstance(value_t, torch.Tensor)
        or value_t.device.type != "cpu"
        or not value_t.is_contiguous()
        or not isinstance(domain, bytes)
        or not domain.endswith(b"\x00")
    ):
        raise ValueError("target-free calibration tensor authority differs")
    digest = hashlib.sha256()
    digest.update(domain)
    dtype = str(value_t.dtype).encode("ascii")
    digest.update(struct.pack("<Q", len(dtype)))
    digest.update(dtype)
    digest.update(struct.pack("<Q", value_t.ndim))
    for extent in value_t.shape:
        digest.update(struct.pack("<Q", extent))
    digest.update(value_t.reshape(-1).view(torch.uint8).numpy().tobytes())
    return torch.frombuffer(
        bytearray(digest.digest()),
        dtype=torch.uint8,
    ).clone()


def target_free_calibration_input_ids_sha256_t_boundary(
    input_ids_t: torch.Tensor,
) -> torch.Tensor:
    """Return the shared exact identity for calibration prompt token IDs."""

    return _domain_separated_tensor_sha256_t_boundary(
        input_ids_t,
        domain=_TARGET_FREE_CALIBRATION_INPUT_IDS_DOMAIN,
    )


def target_free_calibration_input_mask_sha256_t_boundary(
    input_mask_t: torch.Tensor,
) -> torch.Tensor:
    """Return the shared exact identity for a calibration prompt mask."""

    return _domain_separated_tensor_sha256_t_boundary(
        input_mask_t,
        domain=_TARGET_FREE_CALIBRATION_INPUT_MASK_DOMAIN,
    )


@dataclass(frozen=True, slots=True)
class AllKnowledgeTargetFreeCalibrationPromptPacket:
    """Tensor-only authority for one immutable, independently selected prompt.

    The prompt comes from a packed remaining-corpus snapshot, never historical
    branch receipts, stores, scopes, prompts, or tensors. Filesystem paths and
    JSON records remain outside this packet at explicit external-I/O
    boundaries; every content identity needed by the model path is a tensor.
    """

    input_ids_t: torch.Tensor
    input_mask_t: torch.Tensor
    global_cursor_t: torch.Tensor
    prompt_sha256_t: torch.Tensor
    window_sha256_t: torch.Tensor
    row_authority_sha256_t: torch.Tensor
    snapshot_file_sha256_t: torch.Tensor
    collection_authority_sha256_t: torch.Tensor
    federation_authority_sha256_t: torch.Tensor
    exclusion_collection_file_sha256_t: torch.Tensor
    excluded_prior_work_ids_sha256_t: torch.Tensor
    input_ids_sha256_t: torch.Tensor
    input_mask_sha256_t: torch.Tensor
    physical_source_ranges_sha256_t: torch.Tensor
    authority_sha256_t: torch.Tensor
    target_entered_forward_t: torch.Tensor

    def __post_init__(self) -> None:
        tensors = (
            self.input_ids_t,
            self.input_mask_t,
            self.global_cursor_t,
            self.prompt_sha256_t,
            self.window_sha256_t,
            self.row_authority_sha256_t,
            self.snapshot_file_sha256_t,
            self.collection_authority_sha256_t,
            self.federation_authority_sha256_t,
            self.exclusion_collection_file_sha256_t,
            self.excluded_prior_work_ids_sha256_t,
            self.input_ids_sha256_t,
            self.input_mask_sha256_t,
            self.physical_source_ranges_sha256_t,
            self.authority_sha256_t,
            self.target_entered_forward_t,
        )
        digest_tensors = (
            self.prompt_sha256_t,
            self.window_sha256_t,
            self.snapshot_file_sha256_t,
            self.collection_authority_sha256_t,
            self.federation_authority_sha256_t,
            self.exclusion_collection_file_sha256_t,
            self.excluded_prior_work_ids_sha256_t,
            self.input_ids_sha256_t,
            self.input_mask_sha256_t,
            self.physical_source_ranges_sha256_t,
            self.authority_sha256_t,
        )
        if (
            any(
                not isinstance(value_t, torch.Tensor)
                or value_t.device.type != "cpu"
                or not value_t.is_contiguous()
                for value_t in tensors
            )
            or self.input_ids_t.dtype != torch.long
            or self.input_ids_t.ndim != 2
            or self.input_ids_t.shape[0] != 1
            or self.input_ids_t.numel() < 1
            or bool(self.input_ids_t.lt(0).any())
            or self.input_mask_t.dtype != torch.bool
            or self.input_mask_t.shape != self.input_ids_t.shape
            or not bool(self.input_mask_t.all())
            or self.global_cursor_t.dtype != torch.int64
            or self.global_cursor_t.shape != (1,)
            or bool(self.global_cursor_t.lt(0).any())
            or any(
                value_t.dtype != torch.uint8
                or value_t.shape != (32,)
                for value_t in digest_tensors
            )
            or self.row_authority_sha256_t.dtype != torch.uint8
            or self.row_authority_sha256_t.shape != (2, 32)
            or self.target_entered_forward_t.dtype != torch.bool
            or self.target_entered_forward_t.shape != ()
            or bool(self.target_entered_forward_t)
            or not torch.equal(
                self.input_ids_sha256_t,
                target_free_calibration_input_ids_sha256_t_boundary(
                    self.input_ids_t
                ),
            )
            or not torch.equal(
                self.input_mask_sha256_t,
                target_free_calibration_input_mask_sha256_t_boundary(
                    self.input_mask_t
                ),
            )
            or not torch.equal(
                self.authority_sha256_t,
                target_free_calibration_prompt_authority_sha256_t_boundary(
                    self
                ),
            )
        ):
            raise ValueError(
                "all-knowledge target-free calibration prompt differs"
            )


def target_free_calibration_prompt_authority_sha256_t_boundary(
    packet: AllKnowledgeTargetFreeCalibrationPromptPacket,
) -> torch.Tensor:
    """Recompute one prompt packet's content authority without host metadata."""

    if not isinstance(
        packet,
        AllKnowledgeTargetFreeCalibrationPromptPacket,
    ):
        raise TypeError("target-free calibration prompt packet is malformed")
    return _target_free_calibration_prompt_fields_authority_sha256_t_boundary(
        (
            packet.input_ids_t,
            packet.input_mask_t,
            packet.global_cursor_t,
            packet.prompt_sha256_t,
            packet.window_sha256_t,
            packet.row_authority_sha256_t,
            packet.snapshot_file_sha256_t,
            packet.collection_authority_sha256_t,
            packet.federation_authority_sha256_t,
            packet.exclusion_collection_file_sha256_t,
            packet.excluded_prior_work_ids_sha256_t,
            packet.input_ids_sha256_t,
            packet.input_mask_sha256_t,
            packet.physical_source_ranges_sha256_t,
            packet.target_entered_forward_t,
        )
    )


def validated_target_free_calibration_prompt_packet_boundary(
    packet: AllKnowledgeTargetFreeCalibrationPromptPacket,
) -> AllKnowledgeTargetFreeCalibrationPromptPacket:
    """Return an isolated, fully revalidated copy of one prompt packet."""

    if not isinstance(
        packet,
        AllKnowledgeTargetFreeCalibrationPromptPacket,
    ):
        raise TypeError("target-free calibration prompt packet is malformed")
    return AllKnowledgeTargetFreeCalibrationPromptPacket(
        input_ids_t=packet.input_ids_t.detach().cpu().clone().contiguous(),
        input_mask_t=packet.input_mask_t.detach().cpu().clone().contiguous(),
        global_cursor_t=(
            packet.global_cursor_t.detach().cpu().clone().contiguous()
        ),
        prompt_sha256_t=(
            packet.prompt_sha256_t.detach().cpu().clone().contiguous()
        ),
        window_sha256_t=(
            packet.window_sha256_t.detach().cpu().clone().contiguous()
        ),
        row_authority_sha256_t=(
            packet.row_authority_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        snapshot_file_sha256_t=(
            packet.snapshot_file_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        collection_authority_sha256_t=(
            packet.collection_authority_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        federation_authority_sha256_t=(
            packet.federation_authority_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        exclusion_collection_file_sha256_t=(
            packet.exclusion_collection_file_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        excluded_prior_work_ids_sha256_t=(
            packet.excluded_prior_work_ids_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        input_ids_sha256_t=(
            packet.input_ids_sha256_t.detach().cpu().clone().contiguous()
        ),
        input_mask_sha256_t=(
            packet.input_mask_sha256_t.detach().cpu().clone().contiguous()
        ),
        physical_source_ranges_sha256_t=(
            packet.physical_source_ranges_sha256_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
        authority_sha256_t=(
            packet.authority_sha256_t.detach().cpu().clone().contiguous()
        ),
        target_entered_forward_t=(
            packet.target_entered_forward_t.detach()
            .cpu()
            .clone()
            .contiguous()
        ),
    )


def _target_free_calibration_prompt_fields_authority_sha256_t_boundary(
    fields_t: tuple[torch.Tensor, ...],
) -> torch.Tensor:
    """Hash the exact ordered tensor fields that form prompt authority."""

    digest = hashlib.sha256()
    digest.update(_TARGET_FREE_CALIBRATION_PACKET_DOMAIN)
    for value_t in fields_t:
        tensor_digest_t = _domain_separated_tensor_sha256_t_boundary(
            value_t,
            domain=_TARGET_FREE_CALIBRATION_PACKET_DOMAIN,
        )
        digest.update(tensor_digest_t.numpy().tobytes())
    return torch.frombuffer(
        bytearray(digest.digest()),
        dtype=torch.uint8,
    ).clone()


def build_all_knowledge_target_free_calibration_prompt_packet_boundary(
    *,
    input_ids_t: torch.Tensor,
    input_mask_t: torch.Tensor,
    global_cursor_t: torch.Tensor,
    prompt_sha256_t: torch.Tensor,
    window_sha256_t: torch.Tensor,
    row_authority_sha256_t: torch.Tensor,
    snapshot_file_sha256_t: torch.Tensor,
    collection_authority_sha256_t: torch.Tensor,
    federation_authority_sha256_t: torch.Tensor,
    exclusion_collection_file_sha256_t: torch.Tensor,
    excluded_prior_work_ids_sha256_t: torch.Tensor,
    physical_source_ranges_sha256_t: torch.Tensor,
    target_entered_forward_t: torch.Tensor,
) -> AllKnowledgeTargetFreeCalibrationPromptPacket:
    """Seal one already validated prompt selection into its tensor packet."""

    input_ids_t = input_ids_t.detach().cpu().clone().contiguous()
    input_mask_t = input_mask_t.detach().cpu().clone().contiguous()
    global_cursor_t = global_cursor_t.detach().cpu().clone().contiguous()
    prompt_sha256_t = prompt_sha256_t.detach().cpu().clone().contiguous()
    window_sha256_t = window_sha256_t.detach().cpu().clone().contiguous()
    row_authority_sha256_t = (
        row_authority_sha256_t.detach().cpu().clone().contiguous()
    )
    snapshot_file_sha256_t = (
        snapshot_file_sha256_t.detach().cpu().clone().contiguous()
    )
    collection_authority_sha256_t = (
        collection_authority_sha256_t.detach().cpu().clone().contiguous()
    )
    federation_authority_sha256_t = (
        federation_authority_sha256_t.detach().cpu().clone().contiguous()
    )
    exclusion_collection_file_sha256_t = (
        exclusion_collection_file_sha256_t.detach()
        .cpu()
        .clone()
        .contiguous()
    )
    excluded_prior_work_ids_sha256_t = (
        excluded_prior_work_ids_sha256_t.detach()
        .cpu()
        .clone()
        .contiguous()
    )
    physical_source_ranges_sha256_t = (
        physical_source_ranges_sha256_t.detach()
        .cpu()
        .clone()
        .contiguous()
    )
    target_entered_forward_t = (
        target_entered_forward_t.detach().cpu().clone().contiguous()
    )
    input_ids_sha256_t = (
        target_free_calibration_input_ids_sha256_t_boundary(input_ids_t)
    )
    input_mask_sha256_t = (
        target_free_calibration_input_mask_sha256_t_boundary(input_mask_t)
    )
    fields_t = (
        input_ids_t,
        input_mask_t,
        global_cursor_t,
        prompt_sha256_t,
        window_sha256_t,
        row_authority_sha256_t,
        snapshot_file_sha256_t,
        collection_authority_sha256_t,
        federation_authority_sha256_t,
        exclusion_collection_file_sha256_t,
        excluded_prior_work_ids_sha256_t,
        input_ids_sha256_t,
        input_mask_sha256_t,
        physical_source_ranges_sha256_t,
        target_entered_forward_t,
    )
    return AllKnowledgeTargetFreeCalibrationPromptPacket(
        input_ids_t=input_ids_t,
        input_mask_t=input_mask_t,
        global_cursor_t=global_cursor_t,
        prompt_sha256_t=prompt_sha256_t,
        window_sha256_t=window_sha256_t,
        row_authority_sha256_t=row_authority_sha256_t,
        snapshot_file_sha256_t=snapshot_file_sha256_t,
        collection_authority_sha256_t=collection_authority_sha256_t,
        federation_authority_sha256_t=federation_authority_sha256_t,
        exclusion_collection_file_sha256_t=(
            exclusion_collection_file_sha256_t
        ),
        excluded_prior_work_ids_sha256_t=(
            excluded_prior_work_ids_sha256_t
        ),
        input_ids_sha256_t=input_ids_sha256_t,
        input_mask_sha256_t=input_mask_sha256_t,
        physical_source_ranges_sha256_t=(
            physical_source_ranges_sha256_t
        ),
        authority_sha256_t=(
            _target_free_calibration_prompt_fields_authority_sha256_t_boundary(
                fields_t
            )
        ),
        target_entered_forward_t=target_entered_forward_t,
    )


@dataclass(frozen=True, slots=True)
class PackedTokenPhysicalSourceRange:
    """Cold provenance for bytes copied from one immutable packed object.

    ``object_compressed_*`` identifies the physical zstd frame range that had
    to be read. ``source_uncompressed_*`` identifies the exact bytes selected
    from that frame's logical token stream, and ``packet_*`` binds those bytes
    to their position in the coalesced ``PackedTokenBatchPacket`` arena. This
    is external-I/O evidence only; it never enters model forward or routing.
    """

    component_sha256: str
    payload_work_id: str
    object_path: str
    object_sha256: str
    object_byte_count: int
    object_compressed_byte_start: int
    object_compressed_byte_end: int
    source_uncompressed_byte_start: int
    source_uncompressed_byte_end: int
    packet_byte_start: int
    packet_byte_end: int
    global_cursor: int

    def __post_init__(self) -> None:
        for digest, label in (
            (self.component_sha256, "component"),
            (self.payload_work_id, "payload WorkID"),
            (self.object_sha256, "object"),
        ):
            try:
                decoded = bytes.fromhex(digest)
            except ValueError as error:
                raise ValueError(
                    f"packed token physical {label} digest differs"
                ) from error
            if len(decoded) != 32:
                raise ValueError(
                    f"packed token physical {label} digest differs"
                )
        if not self.object_path or not Path(self.object_path).is_absolute():
            raise ValueError("packed token physical object path differs")
        integer_values = (
            self.object_byte_count,
            self.object_compressed_byte_start,
            self.object_compressed_byte_end,
            self.source_uncompressed_byte_start,
            self.source_uncompressed_byte_end,
            self.packet_byte_start,
            self.packet_byte_end,
            self.global_cursor,
        )
        if any(
            not isinstance(value, int)
            or isinstance(value, bool)
            or value < 0
            for value in integer_values
        ):
            raise ValueError("packed token physical byte geometry differs")
        if (
            self.object_byte_count < 1
            or self.object_compressed_byte_end
            <= self.object_compressed_byte_start
            or self.object_compressed_byte_end > self.object_byte_count
            or self.source_uncompressed_byte_end
            <= self.source_uncompressed_byte_start
            or self.packet_byte_end <= self.packet_byte_start
            or (
                self.source_uncompressed_byte_end
                - self.source_uncompressed_byte_start
            )
            != self.packet_byte_end - self.packet_byte_start
            or self.source_uncompressed_byte_start % 4
            or self.source_uncompressed_byte_end % 4
            or self.packet_byte_start % 4
            or self.packet_byte_end % 4
        ):
            raise ValueError("packed token physical byte ranges differ")


@dataclass(frozen=True, slots=True)
class PackedTokenShardGeometryPacket:
    """Tensor-native single-window geometry cached once per sealed shard."""

    token_offsets_t: torch.Tensor
    prompt_lengths_t: torch.Tensor
    prompt_sha256_t: torch.Tensor
    window_sha256_t: torch.Tensor

    def __post_init__(self) -> None:
        row_count = self.prompt_lengths_t.numel()
        if (
            self.token_offsets_t.device.type != "cpu"
            or self.token_offsets_t.dtype != torch.int64
            or self.token_offsets_t.ndim != 1
            or self.token_offsets_t.numel() != row_count + 1
            or not self.token_offsets_t.is_contiguous()
            or not torch.equal(
                self.token_offsets_t[:1],
                torch.zeros((1,), dtype=torch.int64),
            )
            or not bool(
                torch.all(
                    self.token_offsets_t[1:]
                    > self.token_offsets_t[:-1]
                )
            )
            or self.prompt_lengths_t.device.type != "cpu"
            or self.prompt_lengths_t.dtype != torch.int32
            or self.prompt_lengths_t.ndim != 1
            or not self.prompt_lengths_t.is_contiguous()
            or not bool(torch.all(self.prompt_lengths_t > 0))
            or not bool(
                torch.all(
                    self.prompt_lengths_t.to(dtype=torch.int64)
                    < (
                        self.token_offsets_t[1:]
                        - self.token_offsets_t[:-1]
                    )
                )
            )
        ):
            raise ValueError(
                "packed token shard geometry differs"
            )
        for digest_t, label in (
            (self.prompt_sha256_t, "prompt"),
            (self.window_sha256_t, "window"),
        ):
            if (
                digest_t.device.type != "cpu"
                or digest_t.dtype != torch.uint8
                or digest_t.shape != (row_count, 32)
                or not digest_t.is_contiguous()
            ):
                raise ValueError(
                    f"packed token shard {label} authority differs"
                )

    def __len__(self) -> int:
        return self.prompt_lengths_t.numel()


@dataclass(frozen=True, slots=True)
class PackedTokenBatchPacket:
    """Tensor-native packed rows at the immutable corpus I/O boundary.

    ``token_ids_t`` stores each row as its prompt immediately followed by its
    answer. ``row_offsets_t`` and ``prompt_lengths_t`` recover those two spans
    without per-row Python objects. Authority rows store the component SHA-256
    followed by the payload WorkID SHA-256; ``row_authority_index_t`` binds
    each token row to one such pair.
    """

    token_ids_t: torch.Tensor
    row_offsets_t: torch.Tensor
    prompt_lengths_t: torch.Tensor
    prompt_sha256_t: torch.Tensor
    window_sha256_t: torch.Tensor
    authority_sha256_t: torch.Tensor
    row_authority_index_t: torch.Tensor
    global_cursor_start_t: torch.Tensor
    task_intent_targets_t: torch.Tensor | None = None

    def __post_init__(self) -> None:
        tensors = (
            self.token_ids_t,
            self.row_offsets_t,
            self.prompt_lengths_t,
            self.prompt_sha256_t,
            self.window_sha256_t,
            self.authority_sha256_t,
            self.row_authority_index_t,
            self.global_cursor_start_t,
        )
        if any(
            value.device.type != "cpu" or not value.is_contiguous()
            for value in tensors
        ):
            raise ValueError(
                "packed token batch packet tensors must be contiguous CPU tensors"
            )
        if self.token_ids_t.dtype != torch.int32 or self.token_ids_t.ndim != 1:
            raise ValueError("packed token batch token arena differs")
        if (
            self.row_offsets_t.dtype != torch.int64
            or self.row_offsets_t.ndim != 1
            or self.row_offsets_t.numel() < 2
        ):
            raise ValueError("packed token batch row offsets differ")
        if not torch.equal(
            self.row_offsets_t[:1],
            torch.zeros((1,), dtype=torch.int64),
        ) or not torch.equal(
            self.row_offsets_t[-1:],
            torch.tensor(
                [self.token_ids_t.numel()],
                dtype=torch.int64,
            ),
        ):
            raise ValueError("packed token batch token frontier differs")
        row_widths_t = self.row_offsets_t[1:] - self.row_offsets_t[:-1]
        if not bool(torch.all(row_widths_t > 0)):
            raise ValueError("packed token batch row widths differ")
        row_count = self.row_offsets_t.numel() - 1
        if (
            self.prompt_lengths_t.dtype != torch.int32
            or self.prompt_lengths_t.shape != (row_count,)
            or not bool(torch.all(self.prompt_lengths_t > 0))
            or not bool(
                torch.all(
                    self.prompt_lengths_t.to(dtype=torch.int64)
                    < row_widths_t
                )
            )
        ):
            raise ValueError("packed token batch prompt geometry differs")
        for digest_t, label in (
            (self.prompt_sha256_t, "prompt"),
            (self.window_sha256_t, "window"),
        ):
            if digest_t.dtype != torch.uint8 or digest_t.shape != (row_count, 32):
                raise ValueError(f"packed token batch {label} authority differs")
        if (
            self.authority_sha256_t.dtype != torch.uint8
            or self.authority_sha256_t.ndim != 3
            or self.authority_sha256_t.shape[0] < 1
            or self.authority_sha256_t.shape[1:] != (2, 32)
        ):
            raise ValueError("packed token batch authority table differs")
        if (
            self.row_authority_index_t.dtype != torch.int32
            or self.row_authority_index_t.shape != (row_count,)
            or not bool(torch.all(self.row_authority_index_t >= 0))
            or not bool(
                torch.all(
                    self.row_authority_index_t
                    < self.authority_sha256_t.shape[0]
                )
            )
        ):
            raise ValueError("packed token batch row authority indices differ")
        if (
            self.global_cursor_start_t.dtype != torch.int64
            or self.global_cursor_start_t.shape != (1,)
            or not bool(torch.all(self.global_cursor_start_t >= 0))
        ):
            raise ValueError("packed token batch global cursor differs")
        if self.task_intent_targets_t is not None and (
            self.task_intent_targets_t.device.type != "cpu"
            or self.task_intent_targets_t.dtype != torch.float32
            or self.task_intent_targets_t.ndim != 2
            or self.task_intent_targets_t.shape[0] != row_count
            or self.task_intent_targets_t.shape[1] < 1
            or not self.task_intent_targets_t.is_contiguous()
            or not bool(torch.isfinite(self.task_intent_targets_t).all())
        ):
            raise ValueError(
                "packed token batch task-intent arena differs"
            )

    def __len__(self) -> int:
        return self.row_offsets_t.numel() - 1

    def prompt_ids_boundary(self, row_index: int) -> torch.Tensor:
        """Return one prompt view at an explicit Python consumer boundary."""

        if (
            isinstance(row_index, bool)
            or row_index < 0
            or row_index >= len(self)
        ):
            raise IndexError("packed token batch row is unavailable")
        row_start = int(self.row_offsets_t[row_index])
        prompt_end = row_start + int(self.prompt_lengths_t[row_index])
        return self.token_ids_t[row_start:prompt_end]

    def answer_ids_boundary(self, row_index: int) -> torch.Tensor:
        """Return one answer view at an explicit Python consumer boundary."""

        if (
            isinstance(row_index, bool)
            or row_index < 0
            or row_index >= len(self)
        ):
            raise IndexError("packed token batch row is unavailable")
        row_start = int(self.row_offsets_t[row_index])
        answer_start = row_start + int(self.prompt_lengths_t[row_index])
        row_end = int(self.row_offsets_t[row_index + 1])
        return self.token_ids_t[answer_start:row_end]

    def task_intent_targets_boundary(
        self,
        row_index: int,
    ) -> torch.Tensor | None:
        """Return one loss-side intent view at an explicit consumer boundary.

        Intent labels remain separate from ``token_ids_t`` and therefore
        cannot enter the prompt forward.  Keeping their rows in this packet
        still gives the loss boundary the same cursor, authority, and reorder
        semantics as the token arena.
        """

        if (
            isinstance(row_index, bool)
            or row_index < 0
            or row_index >= len(self)
        ):
            raise IndexError("packed token batch row is unavailable")
        if self.task_intent_targets_t is None:
            return None
        return self.task_intent_targets_t[row_index]


class PackedTokenTrainingRow(dict[str, Any]):
    """One metadata row backed by mmap token tensors at the I/O boundary."""

    batch_packet: PackedTokenBatchPacket | None
    batch_packet_row_index: int | None

    def __init__(
        self,
        metadata: dict[str, Any],
        *,
        prompt_ids: torch.Tensor,
        answer_ids: torch.Tensor,
        corrected_prompt_ids: torch.Tensor,
        task_intent_targets: torch.Tensor | None,
        source_row_sha256: str,
        batch_packet: PackedTokenBatchPacket | None = None,
        batch_packet_row_index: int | None = None,
    ) -> None:
        super().__init__(metadata)
        if (batch_packet is None) != (batch_packet_row_index is None):
            raise ValueError("packed token row batch binding is incomplete")
        if batch_packet is not None and (
            not isinstance(batch_packet_row_index, int)
            or isinstance(batch_packet_row_index, bool)
            or batch_packet_row_index < 0
            or batch_packet_row_index >= len(batch_packet)
        ):
            raise ValueError("packed token row batch index differs")
        self["prompt_ids"] = prompt_ids
        self["answer_ids"] = answer_ids
        if task_intent_targets is not None:
            self["task_intent_targets"] = task_intent_targets
        self.corrected_prompt_ids_t = corrected_prompt_ids
        self.source_row_sha256 = source_row_sha256
        self.batch_packet = batch_packet
        self.batch_packet_row_index = batch_packet_row_index

    @staticmethod
    def _integer_list(value: torch.Tensor) -> list[int]:
        return [
            int(token)
            for token in value.to(device="cpu", dtype=torch.long).tolist()
        ]

    def legacy_row_boundary(self) -> dict[str, Any]:
        """Materialize Python sequences only for legacy single-row paths."""

        prompt_ids = self._integer_list(cast(torch.Tensor, self["prompt_ids"]))
        answer_ids = self._integer_list(cast(torch.Tensor, self["answer_ids"]))
        corrected_prompt_ids = self._integer_list(self.corrected_prompt_ids_t)
        row = {
            key: value
            for key, value in self.items()
            if key not in {"prompt_ids", "answer_ids", "task_intent_targets"}
        }
        row.update(
            {
                "prompt_ids": prompt_ids,
                "answer_ids": answer_ids,
                "input_ids": [*prompt_ids, *answer_ids],
                "target_ids": [-100] * len(prompt_ids) + answer_ids,
                "corrected_prompt_ids": corrected_prompt_ids,
                "corrected_input_ids": [
                    *corrected_prompt_ids,
                    *answer_ids,
                ],
                "corrected_target_ids": (
                    [-100] * len(corrected_prompt_ids) + answer_ids
                ),
            }
        )
        task_intent_targets = self.get("task_intent_targets")
        if isinstance(task_intent_targets, torch.Tensor):
            row["task_intent_targets"] = [
                float(value)
                for value in task_intent_targets.to(
                    device="cpu",
                    dtype=torch.float32,
                ).tolist()
            ]
        return row

    def bmpas_tensor_row_boundary(self) -> dict[str, Any]:
        """Tensor-native row for BMPAS train_fn — no Python list materialization."""

        prompt_t = cast(torch.Tensor, self["prompt_ids"]).to(dtype=torch.long)
        answer_t = cast(torch.Tensor, self["answer_ids"]).to(dtype=torch.long)
        corrected_t = self.corrected_prompt_ids_t.to(dtype=torch.long)
        input_ids = torch.cat((prompt_t, answer_t), dim=0)
        target_ids = torch.cat(
            (
                torch.full((prompt_t.numel(),), -100, dtype=torch.long),
                answer_t,
            ),
            dim=0,
        )
        corrected_input = torch.cat((corrected_t, answer_t), dim=0)
        corrected_target = torch.cat(
            (
                torch.full((corrected_t.numel(),), -100, dtype=torch.long),
                answer_t,
            ),
            dim=0,
        )
        row: dict[str, Any] = {
            key: value
            for key, value in self.items()
            if key not in {"prompt_ids", "answer_ids", "task_intent_targets"}
        }
        row.update(
            {
                "prompt_ids": prompt_t,
                "answer_ids": answer_t,
                "input_ids": input_ids,
                "target_ids": target_ids,
                "corrected_prompt_ids": corrected_t,
                "corrected_input_ids": corrected_input,
                "corrected_target_ids": corrected_target,
                "_bmpas_tensor_native": True,
            }
        )
        task_intent_targets = self.get("task_intent_targets")
        if isinstance(task_intent_targets, torch.Tensor):
            row["task_intent_targets"] = task_intent_targets
        return row

    def __eq__(self, other: object) -> bool:
        if isinstance(other, PackedTokenTrainingRow):
            return self.source_row_sha256 == other.source_row_sha256
        if isinstance(other, dict):
            return self.legacy_row_boundary() == other
        return False