"""Tests for the audio encoder and speech modules."""

import numpy as np
import pytest
import torch

from neon.model.audio import AudioConfig, AudioEncoder, SpeechResponseHead, PersonaPlexSpeaker


class TestAudioConfig:
    """Test audio configuration."""

    def test_default_config(self):
        config = AudioConfig()
        assert config.enabled is True
        assert config.encoder_type == "whisper"
        assert config.sample_rate == 16000
        assert config.max_audio_seconds == 15.0
        assert config.personaplex_voice == "NATM1"

    def test_disabled_config(self):
        config = AudioConfig(enabled=False)
        assert config.enabled is False

    def test_mel_config(self):
        config = AudioConfig(encoder_type="mel")
        assert config.encoder_type == "mel"


class TestAudioEncoder:
    """Test the audio encoder module (without loading Whisper)."""

    def test_projection_shape(self):
        """Projection layer should produce correct shape."""
        config = AudioConfig(audio_hidden_size=512)
        encoder = AudioEncoder(config, backbone_hidden_size=2048)
        # Test projection only
        mock_features = torch.randn(4, 512)
        projected = encoder.projection(mock_features)
        assert projected.shape == (4, 2048)

    def test_custom_projection_dim(self):
        """Custom projection dim should override backbone hidden size."""
        config = AudioConfig(audio_hidden_size=256, projection_dim=1024)
        encoder = AudioEncoder(config, backbone_hidden_size=2048)
        mock_features = torch.randn(2, 256)
        projected = encoder.projection(mock_features)
        assert projected.shape == (2, 1024)
        assert encoder.output_dim == 1024

    def test_output_dim_property(self):
        """output_dim should match projection output."""
        config = AudioConfig(audio_hidden_size=512)
        encoder = AudioEncoder(config, backbone_hidden_size=3584)
        assert encoder.output_dim == 3584


class TestSpeechResponseHead:
    """Test the speech response prediction head."""

    def test_output_shapes(self):
        """Head should produce response type and token logits."""
        head = SpeechResponseHead(hidden_size=256, vocab_size=1000, max_response_tokens=50)
        features = torch.randn(4, 256)
        output = head(features)
        assert "response_type" in output
        assert "token_logits" in output
        assert output["response_type"].shape == (4, 5)  # 5 types
        assert output["token_logits"].shape == (4, 1000)

    def test_response_type_classes(self):
        """Should have 5 response types: narrate, confirm, warn, ask, silent."""
        head = SpeechResponseHead(hidden_size=128)
        features = torch.randn(1, 128)
        output = head(features)
        assert output["response_type"].shape[-1] == 5


class TestPersonaPlexSpeaker:
    """Test PersonaPlex speaker (no actual TTS calls)."""

    def test_creation(self):
        config = AudioConfig(personaplex_voice="NATF1")
        speaker = PersonaPlexSpeaker(config)
        assert speaker.voice == "NATF1"

    def test_speak_disabled(self):
        """speak() should return None when disabled."""
        config = AudioConfig(enable_speech_output=False)
        speaker = PersonaPlexSpeaker(config)
        result = speaker.speak("Hello")
        assert result is None

    def test_speak_empty_text(self):
        """speak() should return None for empty text."""
        config = AudioConfig()
        speaker = PersonaPlexSpeaker(config)
        result = speaker.speak("")
        assert result is None

    def test_response_templates(self):
        """Should have templates for each response type."""
        assert "narrate" in PersonaPlexSpeaker.RESPONSE_TEMPLATES
        assert "confirm" in PersonaPlexSpeaker.RESPONSE_TEMPLATES
        assert "warn" in PersonaPlexSpeaker.RESPONSE_TEMPLATES
        assert "ask" in PersonaPlexSpeaker.RESPONSE_TEMPLATES


class TestNeonVLAAudioIntegration:
    """Test that NeonVLA correctly integrates audio when configured."""

    def test_model_without_audio(self):
        """Model without audio config should work as before."""
        from neon.model.neon_vla import NeonConfig, NeonVLA

        config = NeonConfig(control_mode="arms_only", use_proprioception=False)
        model = NeonVLA(config)
        assert model.audio_encoder is None
        assert model.speech_head is None

    def test_model_with_audio_config(self):
        """Model with audio config should build encoder."""
        from neon.model.neon_vla import NeonConfig, NeonVLA

        config = NeonConfig(
            control_mode="arms_only",
            use_proprioception=False,
            audio={"enabled": True, "encoder_type": "whisper", "enable_speech_output": True},
        )
        model = NeonVLA(config)
        assert model.audio_encoder is not None
        assert model.speech_head is not None
        assert model.speaker is not None

    def test_model_audio_disabled(self):
        """Model with audio disabled should not build encoder."""
        from neon.model.neon_vla import NeonConfig, NeonVLA

        config = NeonConfig(
            control_mode="arms_only",
            audio={"enabled": False},
        )
        model = NeonVLA(config)
        assert model.audio_encoder is None

    def test_fusion_dim_with_audio(self):
        """Fusion layer input dim should include audio features."""
        from neon.model.neon_vla import NeonConfig, NeonVLA

        # Without audio
        config_no_audio = NeonConfig(control_mode="arms_only", use_proprioception=False)
        model_no_audio = NeonVLA(config_no_audio)
        fusion_in_no_audio = model_no_audio.fusion[0].in_features

        # With audio
        config_audio = NeonConfig(
            control_mode="arms_only",
            use_proprioception=False,
            audio={"enabled": True, "encoder_type": "whisper"},
        )
        model_audio = NeonVLA(config_audio)
        fusion_in_audio = model_audio.fusion[0].in_features

        # Audio model should have larger fusion input
        assert fusion_in_audio > fusion_in_no_audio
        assert fusion_in_audio == fusion_in_no_audio + model_audio.audio_encoder.output_dim
