"""Tokenizer tests. Run directly: `.venv/bin/python tests/test_tokenizer.py`

No pytest dependency — plain asserts and a tiny runner, to keep the project's
dependencies limited to torch + numpy.
"""

import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from tokenizer import BPETokenizer, pre_tokenize  # noqa: E402

EOT = "<|endoftext|>"

# A small but varied training corpus.
CORPUS = (
    "Once upon a time, there was a little cat. The cat liked to play.\n"
    "The cat ran and ran. Then the cat found a red ball!\n" + EOT + "\n"
    "Tom and Lily went to the park. They were very happy.\n"
) * 50


def _fresh_tokenizer() -> BPETokenizer:
    return BPETokenizer.train(CORPUS, vocab_size=400, special_tokens=[EOT])


def test_pretokenize_is_lossless():
    for text in ["Hello, world!", "  spaces\tand\nnewlines ", "a", "", "!!!??"]:
        assert "".join(pre_tokenize(text)) == text


def test_roundtrip_on_training_text():
    tok = _fresh_tokenizer()
    for line in CORPUS.split("\n"):
        assert tok.decode(tok.encode(line)) == line


def test_roundtrip_on_unseen_text():
    tok = _fresh_tokenizer()
    unseen = "A brave dog jumped over the fence; zebras watched quietly."
    assert tok.decode(tok.encode(unseen)) == unseen


def test_special_token_is_single_id():
    tok = _fresh_tokenizer()
    ids = tok.encode(f"Hello {EOT} world")
    eot_id = tok.special_tokens[EOT]
    assert ids.count(eot_id) == 1
    # The special token id must not appear from encoding ordinary text.
    assert eot_id not in tok.encode("Hello world")
    assert tok.decode(tok.encode(f"a{EOT}b")) == f"a{EOT}b"


def test_unicode_roundtrip():
    tok = _fresh_tokenizer()
    # Multibyte characters must survive being split across byte-level tokens.
    text = "café naïve \U0001f600"
    assert tok.decode(tok.encode(text)) == text


def test_vocab_size_and_layout():
    tok = _fresh_tokenizer()
    # vocab_size is a target ceiling: a tiny, repetitive corpus saturates
    # (every word merges to a single token) before reaching it, so we may get
    # fewer merges. The derived vocab dict must always match the reported size.
    assert tok.vocab_size <= 400
    assert len(tok.vocab) == tok.vocab_size
    # First 256 ids are the raw bytes.
    assert tok.vocab[65] == b"A"


def test_save_load_roundtrip():
    tok = _fresh_tokenizer()
    with tempfile.TemporaryDirectory() as d:
        path = Path(d) / "tok.json"
        tok.save(path)
        loaded = BPETokenizer.load(path)
    sample = f"The cat and Tom. {EOT}"
    assert loaded.encode(sample) == tok.encode(sample)
    assert loaded.vocab_size == tok.vocab_size


def test_determinism():
    a = _fresh_tokenizer()
    b = _fresh_tokenizer()
    assert a.merges == b.merges


def main() -> None:
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for t in tests:
        t()
        print(f"ok  {t.__name__}")
    print(f"\n{len(tests)} passed")


if __name__ == "__main__":
    main()