"""Subword tokenization and alignment utilities."""

from dataclasses import dataclass
from typing import Optional

from transformers import PreTrainedTokenizer


@dataclass
class AlignedTokens:
    """Represents alignment between words and subwords."""

    input_ids: list[int]  # Subword token IDs
    attention_mask: list[int]  # Attention mask
    word_ids: list[Optional[int]]  # Maps subword idx -> word idx (None for special tokens)
    word_start_mask: list[bool]  # True for first subword of each word

    def __len__(self) -> int:
        return len(self.input_ids)


def tokenize_and_align(
    tokenizer: PreTrainedTokenizer,
    words: list[str],
    max_length: int = 512,
) -> AlignedTokens:
    """Tokenize pre-tokenized words and track alignment.

    Args:
        tokenizer: HuggingFace tokenizer.
        words: List of words (pre-tokenized).
        max_length: Maximum sequence length.

    Returns:
        AlignedTokens with subword IDs and alignment information.
    """
    encoding = tokenizer(
        words,
        is_split_into_words=True,
        max_length=max_length,
        truncation=True,
        padding=False,
        return_offsets_mapping=False,
    )

    word_ids = encoding.word_ids()
    word_start_mask = []
    prev_word_id = None

    for word_id in word_ids:
        if word_id is None:
            # Special token (CLS, SEP, PAD)
            word_start_mask.append(False)
        elif word_id != prev_word_id:
            # First subword of a new word
            word_start_mask.append(True)
        else:
            # Continuation subword
            word_start_mask.append(False)
        prev_word_id = word_id

    return AlignedTokens(
        input_ids=encoding["input_ids"],
        attention_mask=encoding["attention_mask"],
        word_ids=word_ids,
        word_start_mask=word_start_mask,
    )


def get_word_count(aligned: AlignedTokens) -> int:
    """Get the number of words in aligned tokens.

    Args:
        aligned: AlignedTokens object.

    Returns:
        Number of words (count of True in word_start_mask).
    """
    return sum(aligned.word_start_mask)


def get_word_to_subword_map(aligned: AlignedTokens) -> dict[int, list[int]]:
    """Get mapping from word index to subword indices.

    Args:
        aligned: AlignedTokens object.

    Returns:
        Dictionary mapping word index to list of subword indices.
    """
    word_to_subwords: dict[int, list[int]] = {}
    for subword_idx, word_id in enumerate(aligned.word_ids):
        if word_id is not None:
            if word_id not in word_to_subwords:
                word_to_subwords[word_id] = []
            word_to_subwords[word_id].append(subword_idx)
    return word_to_subwords
