"""Cite input format: ``[CLS] question [SEP] [BLOCK] text [BLOCK] text ... [SEP]``."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any

BLOCK_TOKEN = "[BLOCK]"


@dataclass(frozen=True)
class CiteInput:
    input_ids: list[int]
    block_positions: list[int]
    n_blocks: int

    @property
    def truncated_blocks(self) -> int:
        """Blocks left out because they did not fit in ``max_tokens``."""

        return self.n_blocks - len(self.block_positions)


def add_block_token(tokenizer: Any) -> int:
    tokenizer.add_special_tokens({"additional_special_tokens": [BLOCK_TOKEN]})
    return int(tokenizer.convert_tokens_to_ids(BLOCK_TOKEN))


def encode(
    tokenizer: Any,
    block_token_id: int,
    question: str,
    blocks: Sequence[str],
    *,
    max_tokens: int,
) -> CiteInput:
    """Encode whole blocks in order until the next one would not fit."""

    ids = [tokenizer.cls_token_id, *tokenizer.encode(question, add_special_tokens=False)]
    ids.append(tokenizer.sep_token_id)
    if len(ids) + 1 > max_tokens:
        raise ValueError(f"question alone needs {len(ids) + 1} tokens, over {max_tokens}")
    positions = []
    for text in blocks:
        piece = [block_token_id, *tokenizer.encode(text, add_special_tokens=False)]
        if len(ids) + len(piece) + 1 > max_tokens:
            break
        positions.append(len(ids))
        ids.extend(piece)
    ids.append(tokenizer.sep_token_id)
    return CiteInput(input_ids=ids, block_positions=positions, n_blocks=len(blocks))