# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: test/test_tokenizer.py
# ─────────────────────────────────────────────────────────

"""
Tests for the custom AXIOM-1 byte-level BPE tokenizer.
"""

import os
import pytest

from axiom1.tokenizer.tokenizer import AXIOM1Tokenizer, SPECIAL_TOKENS


@pytest.fixture(scope="module")
def tokenizer() -> AXIOM1Tokenizer:
    """Train a small tokenizer for testing (cached per module)."""
    texts = [
        "Hello world! AXIOM-1 is the first EGen Core language model.",
        "One block. Twelve iterations. The first axiom.",
        "The quick brown fox jumps over the lazy dog.",
        "Unicode: café, naïve, résumé, 日本語",
        "Numbers: x = 42; y = 3.14; print(x + y)",
    ] * 50
    return AXIOM1Tokenizer.train(
        texts=texts,
        save_dir="/tmp/axiom1_test_tokenizer",
        vocab_size=500,
    )


class TestSpecialTokens:
    """Verify special token IDs are correct."""

    def test_pad_id(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.pad_token_id == 0, f"pad should be 0, got {tokenizer.pad_token_id}"

    def test_bos_id(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.bos_token_id == 1, f"bos should be 1, got {tokenizer.bos_token_id}"

    def test_eos_id(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.eos_token_id == 2, f"eos should be 2, got {tokenizer.eos_token_id}"

    def test_unk_reuses_pad(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.unk_token_id == 0, f"unk should reuse pad (0), got {tokenizer.unk_token_id}"

    def test_bot_id(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.bot_token_id == 3, f"bot should be 3, got {tokenizer.bot_token_id}"

    def test_eot_id(self, tokenizer: AXIOM1Tokenizer) -> None:
        assert tokenizer.eot_token_id == 4, f"eot should be 4, got {tokenizer.eot_token_id}"


class TestEncoding:
    """Tests for encoding behavior."""

    def test_roundtrip(self, tokenizer: AXIOM1Tokenizer) -> None:
        text = "Hello world!"
        ids = tokenizer.encode(text, add_special_tokens=False)
        decoded = tokenizer.decode(ids, skip_special_tokens=False)
        assert decoded.strip() == text, f"Round-trip failed: '{decoded.strip()}' != '{text}'"

    def test_bos_eos_added(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("test", add_special_tokens=True)
        assert ids[0] == 1, f"First token should be BOS (1), got {ids[0]}"
        assert ids[-1] == 2, f"Last token should be EOS (2), got {ids[-1]}"

    def test_no_special_tokens(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("test", add_special_tokens=False)
        assert ids[0] != 1 or len(ids) > 2, "BOS should not be added when disabled"

    def test_returns_list(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("hello")
        assert isinstance(ids, list), f"encode should return list, got {type(ids)}"
        assert all(isinstance(i, int) for i in ids), "All IDs should be integers"


class TestPadding:
    """Tests for padding."""

    def test_pads_to_max_length(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("Hi", add_special_tokens=True,
                               padding=True, max_length=20)
        assert len(ids) == 20, f"Should pad to 20, got {len(ids)}"

    def test_pad_token_is_zero(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("Hi", add_special_tokens=True,
                               padding=True, max_length=20)
        # Last few tokens should be pad (0)
        assert ids[-1] == 0, f"Padding token should be 0, got {ids[-1]}"

    def test_no_padding_without_flag(self, tokenizer: AXIOM1Tokenizer) -> None:
        ids = tokenizer.encode("Hi", add_special_tokens=True,
                               padding=False, max_length=20)
        assert len(ids) < 20, "Should not pad when padding=False"


class TestTruncation:
    """Tests for truncation."""

    def test_truncates_to_max_length(self, tokenizer: AXIOM1Tokenizer) -> None:
        long_text = "word " * 100
        ids = tokenizer.encode(long_text, add_special_tokens=True,
                               truncation=True, max_length=10)
        assert len(ids) == 10, f"Should truncate to 10, got {len(ids)}"

    def test_no_truncation_without_flag(self, tokenizer: AXIOM1Tokenizer) -> None:
        long_text = "word " * 100
        ids = tokenizer.encode(long_text, add_special_tokens=True,
                               truncation=False, max_length=10)
        assert len(ids) > 10, "Should not truncate when truncation=False"


class TestBatchEncoding:
    """Tests for batch encoding."""

    def test_batch_returns_list_of_lists(self, tokenizer: AXIOM1Tokenizer) -> None:
        batch = tokenizer.encode_batch(["hello", "world"])
        assert isinstance(batch, list), "Batch should be a list"
        assert len(batch) == 2, f"Batch size should be 2, got {len(batch)}"

    def test_batch_padding(self, tokenizer: AXIOM1Tokenizer) -> None:
        batch = tokenizer.encode_batch(
            ["hello", "goodbye world"],
            padding=True, max_length=15
        )
        assert all(len(ids) == 15 for ids in batch), \
            "All sequences should be padded to 15"


class TestSaveLoad:
    """Tests for save/load."""

    def test_save_load_roundtrip(self, tokenizer: AXIOM1Tokenizer) -> None:
        tok2 = AXIOM1Tokenizer.load("/tmp/axiom1_test_tokenizer")
        ids1 = tokenizer.encode("Save load test", add_special_tokens=True)
        ids2 = tok2.encode("Save load test", add_special_tokens=True)
        assert ids1 == ids2, "Loaded tokenizer should produce identical IDs"

    def test_config_file_exists(self) -> None:
        assert os.path.exists("/tmp/axiom1_test_tokenizer/config.json"), \
            "config.json should exist after save"

    def test_tokenizer_file_exists(self) -> None:
        assert os.path.exists("/tmp/axiom1_test_tokenizer/tokenizer.json"), \
            "tokenizer.json should exist after save"
