import unittest

DEFAULT_LFM_CONFIG = {
    "model_type": "lfm2",
    "vocab_size": 65536,
    "hidden_size": 512,
    "num_hidden_layers": 4,
    "num_attention_heads": 8,
    "num_key_value_heads": 4,
    "max_position_embeddings": 32768,
    "norm_eps": 1e-5,
    "conv_bias": False,
    "conv_L_cache": 3,
    "block_dim": 512,
    "block_ff_dim": 1536,
    "block_multiple_of": 256,
    "block_ffn_dim_multiplier": None,
    "block_auto_adjust_ff_dim": True,
    "rope_theta": 10000.0,
    # layer_types: one per layer - "conv" or "full_attention"
    "layer_types": ["conv", "full_attention", "conv", "full_attention"],
}


def get_test_config():
    """Create a test config dict with all required fields."""
    return {
        "model_type": "lfm_audio",
        "sample_rate": 24000,
        "codebooks": 8,
        "audio_vocab_size": 2049,
        "preprocessor": {"sample_rate": 16000, "features": 128},
        "encoder": {"d_model": 512, "n_layers": 2},  # Small for testing
        "depthformer": {"layers": 2, "dim": 256},  # Small for testing
        "lfm": DEFAULT_LFM_CONFIG,
    }


class TestPreprocessorConfig(unittest.TestCase):

    def test_defaults(self):
        """Test PreprocessorConfig default values."""
        from mlx_audio.sts.models.lfm_audio.config import PreprocessorConfig

        config = PreprocessorConfig()
        self.assertEqual(config.sample_rate, 16000)
        self.assertEqual(config.features, 128)
        self.assertEqual(config.n_fft, 512)
        self.assertEqual(config.hop_length, 160)
        self.assertEqual(config.win_length, 400)


class TestConformerEncoderConfig(unittest.TestCase):

    def test_defaults(self):
        from mlx_audio.sts.models.lfm_audio.config import ConformerEncoderConfig

        config = ConformerEncoderConfig()
        self.assertEqual(config.feat_in, 128)
        self.assertEqual(config.d_model, 512)
        self.assertEqual(config.n_layers, 17)
        self.assertEqual(config.n_heads, 8)


class TestDepthformerConfig(unittest.TestCase):

    def test_defaults(self):
        from mlx_audio.sts.models.lfm_audio.config import DepthformerConfig

        config = DepthformerConfig()
        self.assertEqual(config.layers, 6)
        self.assertEqual(config.dim, 1024)
        self.assertEqual(config.num_heads, 32)
        self.assertEqual(config.num_kv_heads, 8)


class TestLFM2AudioConfig(unittest.TestCase):

    def test_from_dict(self):
        from mlx_audio.sts.models.lfm_audio.config import LFM2AudioConfig

        config_dict = get_test_config()
        config = LFM2AudioConfig.from_dict(config_dict)

        self.assertEqual(config.model_type, "lfm_audio")
        self.assertEqual(config.sample_rate, 24000)
        self.assertEqual(config.codebooks, 8)
        self.assertEqual(config.audio_vocab_size, 2049)
        self.assertIsNotNone(config.preprocessor)
        self.assertIsNotNone(config.encoder)
        self.assertIsNotNone(config.depthformer)
        self.assertIsNotNone(config.lfm)


class TestLFM2AudioModelOutput(unittest.TestCase):

    def test_call_returns_text_and_audio_logits(self):
        import mlx.core as mx

        from mlx_audio.sts.models.lfm_audio.config import LFM2AudioConfig
        from mlx_audio.sts.models.lfm_audio.model import LFM2AudioModel

        config = LFM2AudioConfig.from_dict(get_test_config())
        model = LFM2AudioModel(config)
        model.eval()

        # Create dummy text input
        batch_size = 1
        seq_len = 5
        text_tokens = mx.zeros((batch_size, seq_len), dtype=mx.int32)

        # Call model
        text_logits, audio_logits = model(text_tokens=text_tokens)
        mx.eval(text_logits)

        # Check text_logits is an array with correct shape
        self.assertIsInstance(text_logits, mx.array)
        self.assertEqual(text_logits.shape[0], batch_size)
        self.assertEqual(text_logits.shape[1], seq_len)

        # Check audio_logits is a list of 8 arrays (one per codebook)
        self.assertIsInstance(audio_logits, list)
        self.assertEqual(len(audio_logits), config.codebooks)

        for logit in audio_logits:
            mx.eval(logit)
            self.assertIsInstance(logit, mx.array)
            self.assertEqual(logit.shape[0], batch_size)
            self.assertEqual(logit.shape[1], seq_len)
            self.assertEqual(logit.shape[2], config.audio_vocab_size)


class TestLFM2AudioModelDtype(unittest.TestCase):

    def test_model_components_exist(self):
        from mlx_audio.sts.models.lfm_audio.config import LFM2AudioConfig
        from mlx_audio.sts.models.lfm_audio.model import LFM2AudioModel

        config = LFM2AudioConfig.from_dict(get_test_config())
        model = LFM2AudioModel(config)

        # Check components exist
        self.assertIsNotNone(model.audio_encoder)
        self.assertIsNotNone(model.audio_adapter)
        self.assertIsNotNone(model.lfm)
        self.assertIsNotNone(model.audio_embedding)
        self.assertIsNotNone(model.audio_head)
        self.assertIsNotNone(model.depth_embeddings)
        self.assertIsNotNone(model.depth_linear)

        # Check depth_embeddings count matches codebooks
        self.assertEqual(len(model.depth_embeddings), config.codebooks)

    def test_sample_rate_property(self):
        from mlx_audio.sts.models.lfm_audio.config import LFM2AudioConfig
        from mlx_audio.sts.models.lfm_audio.model import LFM2AudioModel

        config = LFM2AudioConfig.from_dict(get_test_config())
        model = LFM2AudioModel(config)

        self.assertEqual(model.sample_rate, 24000)


class TestLFMModality(unittest.TestCase):

    def test_modality_values(self):
        from mlx_audio.sts.models.lfm_audio.model import LFMModality

        self.assertEqual(LFMModality.TEXT, 1)
        self.assertEqual(LFMModality.AUDIO_IN, 2)
        self.assertEqual(LFMModality.AUDIO_OUT, 3)


class TestSpecialTokens(unittest.TestCase):

    def test_token_values(self):
        from mlx_audio.sts.models.lfm_audio.model import (
            AUDIO_EOS_TOKEN,
            AUDIO_START_TOKEN,
            IM_END_TOKEN,
            TEXT_END_TOKEN,
        )

        self.assertEqual(AUDIO_START_TOKEN, 128)
        self.assertEqual(IM_END_TOKEN, 7)
        self.assertEqual(TEXT_END_TOKEN, 130)
        self.assertEqual(AUDIO_EOS_TOKEN, 2048)


class TestGenerationConfig(unittest.TestCase):

    def test_defaults(self):
        from mlx_audio.sts.models.lfm_audio.model import GenerationConfig

        config = GenerationConfig()
        self.assertEqual(config.max_new_tokens, 512)
        self.assertEqual(config.temperature, 1.0)
        self.assertEqual(config.top_k, 50)
        self.assertEqual(config.audio_temperature, 1.0)
        self.assertEqual(config.audio_top_k, 4)


if __name__ == "__main__":
    unittest.main()
