import subprocess
import sys
import tempfile
import textwrap
import warnings
from typing import Literal
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import requests
import soundfile as sf
from pydantic import ValidationError

from mistral_common.audio import hertz_to_mel, mel_filter_bank
from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio
from mistral_common.protocol.transcription.request import TranscriptionRequest
from mistral_common.tokens.tokenizers.audio import Audio, AudioConfig, AudioSpectrogramConfig


def _make_dummy_base64() -> str:
    """Create a minimal valid base64-encoded WAV for testing."""
    audio = Audio(audio_array=np.zeros(1600, dtype=np.float32), sampling_rate=16000, format="wav")
    return audio.to_base64("wav")


def sin_wave(sampling_rate: int, duration: float) -> np.ndarray:
    return np.sin(np.ones([int(duration * sampling_rate)]))


def test_audio_config_rejects_sampling_rate_below_frame_rate() -> None:
    encoding_config = AudioSpectrogramConfig(num_mel_bins=80, hop_length=160, window_size=400)
    with pytest.raises(AssertionError):
        AudioConfig(sampling_rate=1, frame_rate=2.0, encoding_config=encoding_config)


def test_audio_resample() -> None:
    sampling_rate = 44_000
    original_array = sin_wave(sampling_rate, 1)

    audio = Audio(
        audio_array=original_array,
        sampling_rate=sampling_rate,
        format="wav",
    )

    audio.resample(sampling_rate)
    assert np.allclose(audio.audio_array, original_array, atol=1e-5)

    audio.resample(sampling_rate // 2)
    assert audio.sampling_rate == sampling_rate // 2
    assert len(audio.audio_array) == len(original_array) // 2
    # fmt: off
    expected_resampled_array = [0.622, 0.907, 0.802, 0.870, 0.819, 0.860, 0.826, 0.854, 0.831, 0.850]
    # fmt: on
    assert np.allclose(audio.audio_array[:10], expected_resampled_array, atol=1e-3)


def test_from_file() -> None:
    sampling_rate = 44100
    original_array = sin_wave(sampling_rate, 1)

    # Test with a local path
    with tempfile.NamedTemporaryFile(suffix=".wav") as tmp:
        with sf.SoundFile(tmp.name, "w", samplerate=sampling_rate, channels=1) as f:
            f.write(original_array)
        audio_local = Audio.from_file(tmp.name)
    assert isinstance(audio_local, Audio)
    assert np.allclose(audio_local.audio_array, original_array, atol=1e-5)
    assert audio_local.sampling_rate == sampling_rate


def test_from_url() -> None:
    # Test with an invalid URL
    with patch("mistral_common.tokens.tokenizers.audio._requests_lib.get") as mock_get:
        mock_get.side_effect = requests.RequestException("connection failed")
        with pytest.raises(
            ValueError,
            match="Failed to download audio from URL: https://example.com/invalid_audio.wav",
        ):
            Audio.from_url("https://example.com/invalid_audio.wav")

    # Test with an invalid content
    with patch("mistral_common.tokens.tokenizers.audio._requests_lib.get") as mock_get:
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.content = b"<html>this is not audio</html>"
        mock_get.return_value = mock_response
        with pytest.raises(
            ValueError,
            match=r"Failed to create Audio instance from URL: https://example\.com \.",
        ):
            Audio.from_url("https://example.com")

    # Test valid URL
    sampling_rate = 16_000
    audio_array = sin_wave(sampling_rate, 1)
    with tempfile.NamedTemporaryFile(suffix=".wav") as tmp:
        with sf.SoundFile(tmp.name, "w", samplerate=sampling_rate, channels=1) as f:
            f.write(audio_array)
        with open(tmp.name, "rb") as f:
            wav_bytes = f.read()

    with patch("mistral_common.tokens.tokenizers.audio._requests_lib.get") as mock_get:
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.content = wav_bytes
        mock_get.return_value = mock_response
        audio_url = Audio.from_url("https://fake-audio-url.com/sample.wav", strict=False)
        assert isinstance(audio_url, Audio)


@pytest.mark.parametrize("prefix", [True, False])
def test_audio_base64(prefix: bool) -> None:
    sampling_rate = 16_000
    original_array = sin_wave(sampling_rate, 3.3)

    LOSSY = ["mp3"]
    LOSSLESS = ["wav"]

    for format in LOSSY + LOSSLESS:
        audio = Audio(
            audio_array=original_array,
            sampling_rate=sampling_rate,
            format=format,
        )

        base64_str = audio.to_base64(format, prefix)
        new_audio = Audio.from_base64(base64_str)

        if prefix:
            assert base64_str.startswith(f"data:audio/{format};base64,")

        assert audio.sampling_rate == new_audio.sampling_rate
        if format in LOSSLESS:
            assert np.allclose(audio.audio_array, new_audio.audio_array, atol=1e-5)
        elif format in LOSSY:

            def rmse(a: np.ndarray, b: np.ndarray) -> float:
                return float(np.sqrt(np.mean((a - b) ** 2)))

            assert rmse(audio.audio_array, new_audio.audio_array) < 5e-3
        else:
            raise ValueError(f"Unknown format {format}")


@pytest.mark.parametrize(
    "freq, expected_mel",
    [
        (100.0, 1.5),  # Linear region
        (1000.0, 15.0),  # Boundary of linear and log regions
        (4000.0, 15.0 + np.log(4000.0 / 1000.0) * 27.0 / np.log(6.4)),  # Log region
        (10000.0, 15.0 + np.log(10000.0 / 1000.0) * 27.0 / np.log(6.4)),  # Log region
    ],
)
def test_hertz_to_mel_single_value(freq: float, expected_mel: float) -> None:
    mel = hertz_to_mel(freq)
    assert isinstance(mel, float)
    assert abs(mel - expected_mel) < 1e-6


def test_hertz_to_mel_array() -> None:
    # Test with an array of frequency values
    freq_array = np.array([100.0, 1000.0, 4000.0])
    expected_mel_array = np.array(
        [
            3.0 * 100.0 / 200.0,  # Linear region
            15.0,  # min_log_mel
            15.0 + np.log(4000.0 / 1000.0) * 27.0 / np.log(6.4),  # Log region
        ]
    )
    np.testing.assert_array_equal(hertz_to_mel(freq_array), expected_mel_array)


def test_mel_filter_bank() -> None:
    # Test mel_filter_bank function
    num_frequency_bins = 256
    num_mel_bins = 20
    min_frequency = 0
    max_frequency = 8000
    sampling_rate = 16000

    mel_filters = mel_filter_bank(num_frequency_bins, num_mel_bins, min_frequency, max_frequency, sampling_rate)

    # Check the shape of the output
    assert mel_filters.shape == (num_frequency_bins, num_mel_bins)

    # Check that all filters sum to 1 (approximately)
    assert np.allclose(mel_filters.sum(axis=0), 0.032, atol=0.005)

    # Check that there are no negative values
    assert (mel_filters >= 0).all()

    # integration test
    expected_array = np.array(
        [
            0.03160475,
            0.03200728,
            0.03184327,
            0.03176877,
            0.03208179,
            0.03160917,
            0.0320748,
            0.03180101,
            0.03186299,
            0.03190694,
            0.0319008,
            0.03182839,
            0.03191571,
            0.03184625,
            0.03188972,
            0.03185817,
            0.0318928,
            0.03186744,
            0.03187269,
            0.0318759,
        ]
    )
    diff = np.abs(mel_filters.sum(0) - expected_array)
    assert (diff < 1e-5).all()


class TestDeprecationWarnings:
    """Each test runs in a fresh subprocess to avoid module-level _warned flags."""

    @staticmethod
    def _run_code(code: str) -> subprocess.CompletedProcess[str]:
        return subprocess.run(
            [sys.executable, "-W", "all", "-c", textwrap.dedent(code)],
            capture_output=True,
            text=True,
            timeout=30,
        )

    def test_rawaudio_instantiation_warns_once(self) -> None:
        result = self._run_code("""\
            import warnings
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter("always")
                from mistral_common.protocol.instruct.chunk import RawAudio
                _ = RawAudio(data="abc", format="wav")
                _ = RawAudio(data="xyz", format="wav")
            dep = [x for x in w if issubclass(x.category, DeprecationWarning) and "RawAudio" in str(x.message)]
            assert len(dep) == 1, f"Expected 1 DeprecationWarning, got {len(dep)}: {dep}"
            print("OK")
        """)
        assert result.returncode == 0, result.stderr
        assert "OK" in result.stdout

    def test_audio_from_raw_audio_warns(self) -> None:
        result = self._run_code("""\
            import warnings
            import numpy as np
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter("always")
                from mistral_common.tokens.tokenizers.audio import Audio
                from mistral_common.protocol.instruct.chunk import RawAudio
                audio_obj = Audio(audio_array=np.zeros(1600, dtype=np.float32), sampling_rate=16000, format="wav")
                b64 = audio_obj.to_base64("wav")
                ra = RawAudio(data=b64, format="wav")
                result = Audio.from_raw_audio(ra)
            dep = [x for x in w if issubclass(x.category, DeprecationWarning) and "from_raw_audio" in str(x.message)]
            assert len(dep) == 1, f"Expected 1 from_raw_audio warning, got {len(dep)}: {dep}"
            assert hasattr(result, "audio_array")
            assert result.sampling_rate > 0
            print("OK")
        """)
        assert result.returncode == 0, result.stderr
        assert "OK" in result.stdout

    @pytest.mark.parametrize("deprecated_name", ["Audio", "AudioFormat", "EXPECTED_FORMAT_VALUES"])
    def test_audio_import_from_old_location_warns(
        self, deprecated_name: Literal["Audio", "AudioFormat", "EXPECTED_FORMAT_VALUES"]
    ) -> None:
        code = (
            "import warnings\n"
            "import numpy as np\n"
            "with warnings.catch_warnings(record=True) as w:\n"
            '    warnings.simplefilter("always")\n'
            f"    from mistral_common.audio import {deprecated_name}\n"
            "dep = [\n"
            "    x for x in w\n"
            "    if issubclass(x.category, DeprecationWarning)\n"
            '    and "mistral_common.audio" in str(x.message)\n'
            "]\n"
            'assert len(dep) == 1, f"Expected 1 import warning, got {len(dep)}: {dep}"\n'
            "# Verify the class is usable\n"
        )
        if deprecated_name == "Audio":
            code += (
                'a = Audio(audio_array=np.zeros(1600, dtype=np.float32), sampling_rate=16000, format="wav")\n'
                "assert a.sampling_rate == 16000\n"
            )
        elif deprecated_name == "AudioFormat":
            code += "assert isinstance(AudioFormat.MP3, AudioFormat)\n"
        else:
            code += 'assert "mp3" in EXPECTED_FORMAT_VALUES\n'
        code += 'print("OK")\n'

        result = self._run_code(code)
        assert result.returncode == 0, result.stderr
        assert "OK" in result.stdout

    def test_rawaudio_import_from_chunk_still_works(self) -> None:
        result = self._run_code("""\
            import warnings
            with warnings.catch_warnings(record=True):
                warnings.simplefilter("always")
                from mistral_common.protocol.instruct.chunk import RawAudio
                ra_str = RawAudio(data="hello", format="wav")
                assert ra_str.data == "hello"
                ra_bytes = RawAudio(data=b"world", format="wav")
                assert ra_bytes.data == b"world"
            print("OK")
        """)
        assert result.returncode == 0, result.stderr
        assert "OK" in result.stdout


class TestRawAudioBackwardCompat:
    def test_rawaudio_still_works(self) -> None:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            ra = RawAudio(data="some_data", format="wav")
        assert ra.data == "some_data"

    def test_rawaudio_format_accepted_silently(self) -> None:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            ra = RawAudio(data="data", format="wav")
        assert ra.format == "wav"

    def test_audiochunk_with_rawaudio_backward_compat(self) -> None:
        b64 = _make_dummy_base64()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            ra = RawAudio(data=b64, format="wav")
            chunk = AudioChunk(input_audio=ra)  # type: ignore[arg-type]
        assert chunk.input_audio == b64

    def test_audiochunk_with_rawaudio_dict_backward_compat(self) -> None:
        b64 = _make_dummy_base64()
        chunk = AudioChunk.model_validate({"input_audio": {"data": b64, "format": "wav"}})
        assert chunk.input_audio == b64

    def test_audiochunk_with_str_new_api(self) -> None:
        chunk = AudioChunk(input_audio="base64data")
        assert chunk.input_audio == "base64data"

    def test_audiochunk_with_bytes_new_api(self) -> None:
        chunk = AudioChunk(input_audio=b"raw_bytes")
        assert chunk.input_audio == b"raw_bytes"

    @pytest.mark.parametrize("empty_audio", ["", "   ", b""])
    def test_audiochunk_empty_input_raises_validation_error(self, empty_audio: str | bytes) -> None:
        with pytest.raises(ValidationError, match="should not be empty"):
            AudioChunk(input_audio=empty_audio)


def test_transcription_request_with_rawaudio_backward_compat() -> None:
    b64 = _make_dummy_base64()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", DeprecationWarning)
        ra = RawAudio(data=b64, format="wav")
        req = TranscriptionRequest(audio=ra, language=None, target_streaming_delay_ms=None)  # type: ignore[arg-type]
    assert isinstance(req.audio, str)
    assert req.audio == b64
