"""Causal segment boundaries and visibility masks."""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass(frozen=True, slots=True)
class SegmentLayout:
    token_segment_ids: torch.Tensor
    segment_token_mask: torch.Tensor
    segment_valid: torch.Tensor
    starts: torch.Tensor
    ends: torch.Tensor

    @property
    def segment_count(self) -> int:
        return int(self.segment_token_mask.shape[1])

    def historical_memory_mask(self, memory_segment_ids: torch.Tensor) -> torch.Tensor:
        """Return [batch, token, memory] visibility for completed history only."""
        if memory_segment_ids.ndim != 2:
            raise ValueError("memory_segment_ids must have shape [batch, memory]")
        token_segments = self.token_segment_ids.unsqueeze(-1)
        return memory_segment_ids.unsqueeze(1) < token_segments


def fixed_segment_layout(
    attention_mask: torch.Tensor,
    *,
    segment_size: int,
) -> SegmentLayout:
    """Create bounded contiguous segments without using future content.

    Sentence and paragraph boundary commits can be supplied later as explicit
    boundary metadata. Fixed maximum-length commits are always available and
    preserve the same completed-history causality contract.
    """
    if attention_mask.ndim != 2:
        raise ValueError("attention_mask must have shape [batch, sequence]")
    batch, seq_len = attention_mask.shape
    device = attention_mask.device
    segment_count = (seq_len + segment_size - 1) // segment_size
    positions = torch.arange(seq_len, device=device)
    token_segment_ids = torch.div(positions, segment_size, rounding_mode="floor")
    token_segment_ids = token_segment_ids.unsqueeze(0).expand(batch, -1)
    segment_ids = torch.arange(segment_count, device=device)
    segment_token_mask = token_segment_ids.unsqueeze(1) == segment_ids.view(1, -1, 1)
    segment_token_mask &= attention_mask.to(torch.bool).unsqueeze(1)
    segment_valid = segment_token_mask.any(dim=-1)
    starts = segment_ids * segment_size
    ends = torch.minimum(starts + segment_size, torch.tensor(seq_len, device=device))
    return SegmentLayout(
        token_segment_ids=token_segment_ids,
        segment_token_mask=segment_token_mask,
        segment_valid=segment_valid,
        starts=starts,
        ends=ends,
    )


def assert_no_open_segment_visibility(layout: SegmentLayout, memory_segment_ids: torch.Tensor) -> None:
    visible = layout.historical_memory_mask(memory_segment_ids)
    same_segment = memory_segment_ids.unsqueeze(1) == layout.token_segment_ids.unsqueeze(-1)
    if bool((visible & same_segment).any()):
        raise AssertionError("open segment can read its own compiled state")


__all__ = ["SegmentLayout", "assert_no_open_segment_visibility", "fixed_segment_layout"]