from tokenizers import pre_tokenizers
from tokenizers.trainers import BpeTrainer
from transformers import PreTrainedTokenizerFast

from korbyte.comparison import Document, _kmmlu_test_files, _ranks, evaluate_tokenizer
from korbyte.train import build_tokenizer


def test_comparison_metrics_preserve_exact_text() -> None:
    tokenizer = build_tokenizer()
    tokenizer.train_from_iterator(
        ["한국어 테스트", "abc 123456"] * 4,
        trainer=BpeTrainer(
            vocab_size=1_000,
            min_frequency=1,
            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        ),
    )
    fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer)
    metrics = evaluate_tokenizer(
        fast,
        [Document("한국어 테스트", 2), Document("abc 123456", 2)],
        repeats=2,
    )
    assert metrics.documents == 2
    assert metrics.exact_document_ratio == 1.0
    assert metrics.byte_fidelity == 1.0
    assert metrics.unknown_token_ratio == 0.0


def test_ranks_lower_metrics_first() -> None:
    results = {
        "winner": {"metrics": {"fertility": 1.0}},
        "runner-up": {"metrics": {"fertility": 1.2}},
        "unavailable": {"error": "unsupported"},
    }
    assert _ranks(results, "fertility") == {"winner": 1, "runner-up": 2}


def test_kmmlu_loader_requires_all_subjects() -> None:
    class Api:
        def list_repo_files(self, *args, **kwargs):
            return [f"data/subject-{index}-test.csv" for index in range(45)]

    files = _kmmlu_test_files(Api())
    assert len(files) == 45
    assert files[0] == "data/subject-0-test.csv"