"""Tokenizer-dependent helpers shared by both stages."""
from __future__ import annotations

from dataclasses import dataclass

from .bio import content_offsets
from .spans import Span


def special_ids(tokenizer) -> tuple[int, int, int]:
    cls = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id
    sep = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id
    pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    if cls is None or sep is None:
        raise ValueError("Tokenizer needs CLS/SEP (or BOS/EOS) tokens")
    return cls, sep, pad


def require_fast(tokenizer) -> None:
    if not getattr(tokenizer, "is_fast", False):
        raise RuntimeError("A fast tokenizer is required for offset mapping "
                           "(pip install sentencepiece protobuf tokenizers)")


@dataclass
class ArticleEncoding:
    article_id: str
    ids: list[int]
    offsets: list[Span]   # whitespace-stripped char offsets, one per token


def encode_article(tokenizer, article_id: str, text: str) -> ArticleEncoding:
    """Tokenize the whole article once, without special tokens or truncation.

    Windows are cut from this later, so token indices are global and window
    predictions can be merged unambiguously.
    """
    enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True,
                    truncation=False, verbose=False)
    return ArticleEncoding(article_id, list(enc["input_ids"]),
                           content_offsets(text, enc["offset_mapping"]))