import torch
from torch.nn.attention.flex_attention import create_block_mask


def build_doc_block_mask(segment_ids: torch.Tensor, seq_len: int, device):
    """Build the FlexAttention block mask for causal + same-document attention.

    segment_ids: (B, T) int32 tensor from the dataset.
    A query at position q attends to a key at position kv only when
    kv <= q (causal) AND both positions belong to the same packed document.
    """
    B = segment_ids.shape[0]

    def doc_mask_mod(b, h, q_idx, kv_idx):
        causal = q_idx >= kv_idx
        same_doc = segment_ids[b, q_idx] == segment_ids[b, kv_idx]
        return causal & same_doc

    # H=None means one mask is shared across all heads
    return create_block_mask(
        doc_mask_mod, B=B, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device
    )