# tests/test_config.py
import pytest
from bclm.config import (
    BCLM1ModelConfig,
    EOT_TOKEN_ID,
    PAD_TOKEN_ID,
    USER_TOKEN_ID,
    ASSISTANT_TOKEN_ID,
    SYSTEM_TOKEN_ID,
    CONFIG_FILENAME,
    WEIGHTS_FILENAME,
)


class TestSpecialTokens:
    def test_token_ids_are_unique(self):
        ids = [
            EOT_TOKEN_ID,
            PAD_TOKEN_ID,
            USER_TOKEN_ID,
            ASSISTANT_TOKEN_ID,
            SYSTEM_TOKEN_ID,
        ]
        assert len(ids) == len(set(ids))

    def test_token_ids_above_base_vocab(self):
        for tid in [EOT_TOKEN_ID, PAD_TOKEN_ID, USER_TOKEN_ID,
                     ASSISTANT_TOKEN_ID, SYSTEM_TOKEN_ID]:
            assert tid >= 32000

    def test_filenames(self):
        assert CONFIG_FILENAME == "config.json"
        assert WEIGHTS_FILENAME == "model.safetensors"


class TestBCLM1ModelConfig:
    def test_defaults(self):
        cfg = BCLM1ModelConfig()
        assert cfg.vocab_size == 32768
        assert cfg.embed_dim == 384
        assert cfg.n_layers == 10
        assert cfg.attn_heads == 6
        assert cfg.attn_kv_heads == 2
        assert isinstance(cfg.local_attn_layers, tuple)
        assert isinstance(cfg.global_attn_layers, tuple)

    def test_from_dict_basic(self):
        d = {
            "vocab_size": 65536,
            "embed_dim": 768,
            "n_layers": 24,
            "attn_heads": 12,
            "attn_kv_heads": 4,
            "local_attn_layers": [1, 3, 5],
            "global_attn_layers": [7, 11],
        }
        cfg = BCLM1ModelConfig.from_dict(d)
        assert cfg.vocab_size == 65536
        assert cfg.embed_dim == 768
        assert cfg.n_layers == 24
        assert cfg.local_attn_layers == (1, 3, 5)
        assert cfg.global_attn_layers == (7, 11)

    def test_from_dict_legacy_alias(self):
        d = {"attn_block_size": 2048}
        cfg = BCLM1ModelConfig.from_dict(d)
        assert cfg.attn_window_size == 2048

    def test_from_dict_strips_unknown_keys(self):
        d = {
            "vocab_size": 32768,
            "some_future_key": True,
            "another_unknown": 42,
        }
        cfg = BCLM1ModelConfig.from_dict(d)
        assert cfg.vocab_size == 32768
        assert not hasattr(cfg, "some_future_key")

    def test_to_dict_roundtrip(self):
        original = BCLM1ModelConfig(
            vocab_size=32768,
            local_attn_layers=(1, 5, 9),
            global_attn_layers=(3,),
        )
        d = original.to_dict()
        assert isinstance(d["local_attn_layers"], list)
        assert isinstance(d["global_attn_layers"], list)

        restored = BCLM1ModelConfig.from_dict(d)
        assert restored.vocab_size == original.vocab_size
        assert restored.local_attn_layers == original.local_attn_layers
        assert restored.global_attn_layers == original.global_attn_layers

    def test_from_dict_full_config_json(self):
        """Simulate parsing the 'model' key from an actual config.json."""
        model_dict = {
            "vocab_size": 32768,
            "tokenizer": "tokenmonster:english-32000-consistent-v1",
            "embed_dim": 384,
            "n_layers": 12,
            "max_seq_len": 16384,
            "dropout": 0.0,
            "attn_heads": 6,
            "attn_kv_heads": 2,
            "local_attn_layers": [1, 5, 7, 11],
            "global_attn_layers": [3, 9],
            "attn_window_size": 1024,
            "conv_kernel_size": 4,
            "osc_n_pairs": 1,
            "osc_n_real": 16,
            "osc_clamp_min_decay": 1e-05,
            "bigram_table_factor": 5,
        }
        cfg = BCLM1ModelConfig.from_dict(model_dict)
        assert cfg.n_layers == 12
        assert cfg.local_attn_layers == (1, 5, 7, 11)
        assert cfg.global_attn_layers == (3, 9)
        assert cfg.tokenizer == "tokenmonster:english-32000-consistent-v1"
        assert cfg.osc_clamp_min_decay == pytest.approx(1e-5)