from tokenizers import pre_tokenizers
from tokenizers.trainers import BpeTrainer

from korbyte.challenges import CHALLENGE_TEXTS
from korbyte.special_tokens import added_special_tokens, special_token_strings
from korbyte.train import build_tokenizer
from korbyte.validate import _random_challenges


def _small_tokenizer():
    tokenizer = build_tokenizer()
    trainer = BpeTrainer(
        vocab_size=2_000,
        min_frequency=1,
        special_tokens=[],
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
    )
    tokenizer.train_from_iterator(CHALLENGE_TEXTS * 4, trainer=trainer)
    tokenizer.add_special_tokens(added_special_tokens())
    return tokenizer


def test_byte_level_pipeline_is_exactly_reversible() -> None:
    tokenizer = _small_tokenizer()
    for text in [*CHALLENGE_TEXTS, *_random_challenges(count=50)]:
        encoded = tokenizer.encode(text, add_special_tokens=False)
        assert tokenizer.decode(encoded.ids, skip_special_tokens=False) == text


def test_special_tokens_are_added_in_declared_order() -> None:
    tokenizer = _small_tokenizer()
    vocab = tokenizer.get_vocab(with_added_tokens=True)
    ids = [vocab[token] for token in special_token_strings()]
    assert ids == list(range(ids[0], ids[0] + 256))


def test_pre_tokenizer_keeps_word_boundaries_and_six_digit_numbers() -> None:
    segments = build_tokenizer().pre_tokenizer.pre_tokenize_str(
        "가나다 라마바 123456 1234567, 테스트"
    )
    assert [offset for _, offset in segments] == [
        (0, 3),
        (3, 7),
        (7, 14),
        (14, 21),
        (21, 22),
        (22, 23),
        (23, 27),
    ]