from __future__ import annotations

from .benchmark import (
    JAPANESE_RE,
    BenchmarkExample,
    CorpusToken,
    PreparedSplit,
    SplitCoverage,
)
from .domain import RerankRequest
from .mozc import MozcDictionaryIndex


def prepare_mozc_examples(
    sentences: tuple[tuple[CorpusToken, ...], ...],
    index: MozcDictionaryIndex,
    *,
    pool_size: int,
    context_mode: str,
) -> PreparedSplit:
    if pool_size < 2:
        raise ValueError("pool_size must be at least 2")
    if context_mode not in {"left_only", "bidirectional"}:
        raise ValueError("context_mode must be left_only or bidirectional")
    eligible_tokens = 0
    ambiguous = 0
    oracle_in_pool = 0
    oracle_miss = 0
    examples: list[BenchmarkExample] = []
    for sentence in sentences:
        surfaces = tuple(token.surface for token in sentence)
        for target_index, token in enumerate(sentence):
            if not token.reading or not JAPANESE_RE.search(token.surface):
                continue
            eligible_tokens += 1
            candidates = index.lookup(token.reading, limit=pool_size)
            if len(candidates) < 2:
                continue
            ambiguous += 1
            if token.surface not in {candidate.surface for candidate in candidates}:
                oracle_miss += 1
                continue
            oracle_in_pool += 1
            examples.append(
                BenchmarkExample(
                    request=RerankRequest(
                        reading=token.reading,
                        candidates=candidates,
                        left_context=surfaces[:target_index],
                        right_context=(
                            surfaces[target_index + 1 :]
                            if context_mode == "bidirectional"
                            else ()
                        ),
                    ),
                    expected=token.surface,
                )
            )
    return PreparedSplit(
        examples=tuple(examples),
        coverage=SplitCoverage(
            eligible_tokens=eligible_tokens,
            ambiguous_known_reading=ambiguous,
            oracle_in_pool=oracle_in_pool,
            oracle_miss=oracle_miss,
        ),
    )