import random
from pathlib import Path

import pytest

from grimboost.tokenization.byte_bpe import ByteBPETokenizer, train_byte_bpe

CORPUS = [
    "The quick brown fox jumps over the lazy dog. ",
    "def f(x):\n    return x + 1\n",
    "中文测试：字节级分词必须精确往返。",
    "Emoji 🙂🚀 and accents café naïve. ",
    "Special markers <|im_start|>user<|im_end|> are ordinary text here. ",
    "Tabs\tand   multiple   spaces and\nnewlines.\n",
]


def _iter_corpus():
    for _ in range(200):
        for line in CORPUS:
            yield line


@pytest.fixture(scope="module")
def tokenizer(tmp_path_factory: pytest.TempPathFactory) -> ByteBPETokenizer:
    path = tmp_path_factory.mktemp("tok") / "bpe.json"
    return train_byte_bpe(
        _iter_corpus(),
        vocab_size=320,
        save_path=path,
        min_frequency=1,
    )


def test_roundtrip_all_examples(tokenizer: ByteBPETokenizer):
    for text in CORPUS:
        ids = tokenizer.encode(text)
        assert tokenizer.decode_bytes(ids) == text.encode("utf-8")
        assert tokenizer.decode(ids) == text


def test_encode_with_spans(tokenizer: ByteBPETokenizer):
    for text in CORPUS:
        data = text.encode("utf-8")
        ids, spans = tokenizer.encode_with_spans(data)
        assert spans[0][0] == 0
        assert spans[-1][1] == len(data)
        for (a, b), (c, _) in zip(spans, spans[1:]):
            assert b == c
        assert b"".join(tokenizer.token_bytes(i) for i in ids) == data


def test_random_unicode_roundtrip(tokenizer: ByteBPETokenizer):
    rng = random.Random(0)
    alphabet = "abc XYZ\n\t.,!?é中🙂🚀0123456789"
    for _ in range(50):
        text = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 80)))
        ids = tokenizer.encode(text)
        assert tokenizer.decode_bytes(ids) == text.encode("utf-8")


def test_save_load(tokenizer: ByteBPETokenizer, tmp_path: Path):
    path = tmp_path / "saved.json"
    tokenizer.save(path)
    loaded = ByteBPETokenizer.load(path)
    text = "roundtrip café 中文 🙂"
    assert loaded.decode_bytes(loaded.encode(text)) == text.encode("utf-8")
    assert loaded.vocab_size == tokenizer.vocab_size


def test_special_tokens_have_no_bytes(tokenizer: ByteBPETokenizer):
    for idx in tokenizer.special_ids:
        assert tokenizer.token_bytes(idx) is None


def test_encode_bytes_exact_with_invalid_tail(tokenizer: ByteBPETokenizer):
    samples = [
        b"hello \xe4\xb8",  # truncated CJK
        b"abc\xff\xfe",     # invalid bytes
        b"abc",
        b"\xf0\x9f\x99",    # truncated emoji
    ]
    for data in samples:
        ids, spans = tokenizer.encode_bytes_exact(data)
        recon = b"".join(tokenizer.token_bytes(i) for i in ids)
        assert recon == data
        assert spans[0][0] == 0
        assert spans[-1][1] == len(data)