"""Tests for the AceStepPipeline orchestration layer.

These are shape-only / structural tests that use tiny configs and mock
weight loading. No real model weights are required.
"""

import tempfile
from pathlib import Path
from typing import Optional
from unittest.mock import MagicMock, patch

import mlx.core as mx
import numpy as np
import pytest

from mlx_music.models.acestep.config import (
    AceStepConfig,
    AceStepDiTConfig,
    AceStepVAEConfig,
)
from mlx_music.models.acestep.dit import AceStepConditionGenerationModel
from mlx_music.models.acestep.pipeline import AceStepPipeline
from mlx_music.models.acestep.vae import AutoencoderOobleck


# ---------------------------------------------------------------------------
# Tiny configs for fast unit tests
# ---------------------------------------------------------------------------


def _small_dit_config() -> AceStepDiTConfig:
    return AceStepDiTConfig(
        hidden_size=64,
        intermediate_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=16,
        num_lyric_encoder_hidden_layers=1,
        num_timbre_encoder_hidden_layers=1,
        in_channels=24,  # 8 + 8 + 8
        audio_acoustic_hidden_dim=8,
        patch_size=2,
        text_hidden_dim=32,
        timbre_hidden_dim=8,
        sliding_window=4,
        layer_types=["sliding_attention", "full_attention"],
    )


def _small_vae_config() -> AceStepVAEConfig:
    # decoder_input_channels must match DiT audio_acoustic_hidden_dim (8)
    # encoder_hidden_size must be 2 * decoder_input_channels (16)
    return AceStepVAEConfig(
        encoder_hidden_size=16,
        downsampling_ratios=[2, 4],
        channel_multiples=[1, 2],
        decoder_channels=8,
        decoder_input_channels=8,
        audio_channels=2,
        sampling_rate=48000,
    )


def _small_config() -> AceStepConfig:
    return AceStepConfig(
        dit=_small_dit_config(),
        vae=_small_vae_config(),
        latent_rate=25,
        sample_rate=48000,
    )


# ---------------------------------------------------------------------------
# Helper: build a pipeline with tiny models (no disk weights)
# ---------------------------------------------------------------------------


def _make_loaded_pipeline() -> AceStepPipeline:
    """Build a pipeline with small random-weight models.

    Bypasses disk I/O entirely -- constructs models directly from tiny
    configs and marks the pipeline as loaded.
    """
    pipeline = AceStepPipeline("/nonexistent")
    pipeline.config = _small_config()
    pipeline.dit_model = AceStepConditionGenerationModel(pipeline.config.dit)
    pipeline.vae_model = AutoencoderOobleck(pipeline.config.vae)
    # Mock the text encoder and tokenizer
    pipeline.text_model = MagicMock()
    pipeline.tokenizer = MagicMock()
    pipeline._loaded = True
    return pipeline


# ---------------------------------------------------------------------------
# Constructor tests
# ---------------------------------------------------------------------------


class TestPipelineInit:
    def test_constructor_stores_path(self):
        p = AceStepPipeline("/some/path")
        assert p.weights_dir == Path("/some/path")

    def test_constructor_does_not_load(self):
        p = AceStepPipeline("/some/path")
        assert not p._loaded
        assert p.dit_model is None
        assert p.vae_model is None
        assert p.text_model is None
        assert p.tokenizer is None

    def test_not_loaded_raises(self):
        p = AceStepPipeline("/some/path")
        with pytest.raises(RuntimeError, match="Models not loaded"):
            p.encode_conditions("test caption")

    def test_not_loaded_generate_raises(self):
        p = AceStepPipeline("/some/path")
        with pytest.raises(RuntimeError, match="Models not loaded"):
            p.generate("test caption")


# ---------------------------------------------------------------------------
# Condition encoding tests (mocked text encoder)
# ---------------------------------------------------------------------------


class TestEncodeConditions:
    """Test condition encoding with mocked text encoder outputs."""

    def _patch_text_encoder(self, pipeline: AceStepPipeline):
        """Patch encode_text and encode_lyrics to return shaped tensors."""
        text_dim = pipeline.config.dit.text_hidden_dim

        def fake_encode_text(model, tokenizer, text, max_length=512):
            seq_len = min(len(text.split()) + 2, max_length)
            h = mx.random.normal((1, seq_len, text_dim))
            m = mx.ones((1, seq_len))
            return h, m

        def fake_encode_lyrics(model, tokenizer, lyrics, max_length=4096):
            seq_len = min(len(lyrics.split()) + 2, max_length)
            h = mx.random.normal((1, seq_len, text_dim))
            m = mx.ones((1, seq_len))
            return h, m

        return fake_encode_text, fake_encode_lyrics

    def test_returns_three_tensors(self):
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with (
            patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
            patch("mlx_music.models.acestep.pipeline.encode_lyrics", fake_lyrics),
        ):
            enc_h, enc_m, ctx = pipeline.encode_conditions(
                "jazz fusion", "la la la", duration=2.0
            )
            mx.eval(enc_h, enc_m, ctx)

        assert enc_h.ndim == 3
        assert enc_m.ndim == 2
        assert ctx.ndim == 3

    def test_context_latent_shape(self):
        pipeline = _make_loaded_pipeline()
        cfg = pipeline.config
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        duration = 2.0
        expected_frames = cfg.duration_to_latent_frames(duration)
        context_channels = (
            2 * cfg.dit.audio_acoustic_hidden_dim
        )  # 128 for full, 16 for small

        with (
            patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
            patch("mlx_music.models.acestep.pipeline.encode_lyrics", fake_lyrics),
        ):
            _, _, ctx = pipeline.encode_conditions("test", duration=duration)
            mx.eval(ctx)

        assert ctx.shape == (1, expected_frames, context_channels)

    def test_encoder_hidden_states_dim(self):
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with (
            patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
            patch("mlx_music.models.acestep.pipeline.encode_lyrics", fake_lyrics),
        ):
            enc_h, _, _ = pipeline.encode_conditions("pop ballad")
            mx.eval(enc_h)

        # The packed condition hidden dim should be hidden_size
        assert enc_h.shape[2] == pipeline.config.dit.hidden_size

    def test_empty_lyrics_uses_space(self):
        """Empty lyrics should still produce valid conditioning."""
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with (
            patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
            patch("mlx_music.models.acestep.pipeline.encode_lyrics", fake_lyrics),
        ):
            enc_h, enc_m, ctx = pipeline.encode_conditions(
                "ambient drone", lyrics="", duration=1.0
            )
            mx.eval(enc_h, enc_m, ctx)

        assert enc_h.shape[0] == 1
        assert not mx.any(mx.isnan(enc_h)).item()


# ---------------------------------------------------------------------------
# Full generate tests (mocked text encoder + small models)
# ---------------------------------------------------------------------------


class TestGenerate:
    """Integration tests for the full generation pipeline with tiny models."""

    def _patch_text_encoder(self, pipeline: AceStepPipeline):
        text_dim = pipeline.config.dit.text_hidden_dim
        # Use deterministic values so conditioning is identical across calls
        _text_h = mx.ones((1, 5, text_dim)) * 0.1
        _lyric_h = mx.ones((1, 3, text_dim)) * 0.2

        def fake_encode_text(model, tokenizer, text, max_length=512):
            return _text_h, mx.ones((1, 5))

        def fake_encode_lyrics(model, tokenizer, lyrics, max_length=4096):
            return _lyric_h, mx.ones((1, 3))

        return fake_encode_text, fake_encode_lyrics

    def test_returns_output_path(self):
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with tempfile.TemporaryDirectory() as tmpdir:
            out_path = str(Path(tmpdir) / "test.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                result = pipeline.generate(
                    caption="test",
                    duration=1.0,
                    steps=2,
                    seed=42,
                    output=out_path,
                )

            assert result == out_path
            assert Path(out_path).exists()

    def test_output_is_wav(self):
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with tempfile.TemporaryDirectory() as tmpdir:
            out_path = str(Path(tmpdir) / "music.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                pipeline.generate(
                    caption="rock", duration=1.0, steps=2, seed=0, output=out_path
                )

            # Check file is non-empty
            assert Path(out_path).stat().st_size > 0

    def test_seeded_reproducibility(self):
        """Same seed should produce identical audio output."""
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with tempfile.TemporaryDirectory() as tmpdir:
            path_a = str(Path(tmpdir) / "a.wav")
            path_b = str(Path(tmpdir) / "b.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                pipeline.generate(
                    caption="test", duration=1.0, steps=2, seed=42, output=path_a
                )
                pipeline.generate(
                    caption="test", duration=1.0, steps=2, seed=42, output=path_b
                )

            import soundfile as sf

            audio_a, sr_a = sf.read(path_a)
            audio_b, sr_b = sf.read(path_b)
            assert sr_a == sr_b
            np.testing.assert_array_equal(audio_a, audio_b)

    def test_different_seeds_differ(self):
        """Different seeds should produce different audio."""
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with tempfile.TemporaryDirectory() as tmpdir:
            path_a = str(Path(tmpdir) / "a.wav")
            path_b = str(Path(tmpdir) / "b.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                pipeline.generate(
                    caption="test", duration=1.0, steps=2, seed=1, output=path_a
                )
                pipeline.generate(
                    caption="test", duration=1.0, steps=2, seed=2, output=path_b
                )

            import soundfile as sf

            audio_a, _ = sf.read(path_a)
            audio_b, _ = sf.read(path_b)
            assert not np.array_equal(audio_a, audio_b)

    def test_callback_is_called(self):
        """Progress callback should fire once per denoising step."""
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)
        call_log = []

        def cb(step, total, xt):
            call_log.append((step, total))

        with tempfile.TemporaryDirectory() as tmpdir:
            out_path = str(Path(tmpdir) / "cb.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                pipeline.generate(
                    caption="test",
                    duration=1.0,
                    steps=4,
                    seed=0,
                    output=out_path,
                    callback=cb,
                )

        assert len(call_log) == 4
        assert call_log[0][1] == 4  # total steps
        assert call_log[-1][0] == 3  # last step index

    def test_no_nan_in_output(self):
        """Generated audio should not contain NaN."""
        pipeline = _make_loaded_pipeline()
        fake_text, fake_lyrics = self._patch_text_encoder(pipeline)

        with tempfile.TemporaryDirectory() as tmpdir:
            out_path = str(Path(tmpdir) / "nan_check.wav")

            with (
                patch("mlx_music.models.acestep.pipeline.encode_text", fake_text),
                patch(
                    "mlx_music.models.acestep.pipeline.encode_lyrics",
                    fake_lyrics,
                ),
            ):
                pipeline.generate(
                    caption="test", duration=1.0, steps=2, seed=99, output=out_path
                )

            import soundfile as sf

            audio, _ = sf.read(out_path)
            assert not np.any(np.isnan(audio)), "Output audio contains NaN"


# ---------------------------------------------------------------------------
# Dimension algebra tests
# ---------------------------------------------------------------------------


class TestDimensionAlgebra:
    """Verify the key dimensional relationships in the pipeline."""

    def test_latent_rate_matches_config(self):
        cfg = _small_config()
        # 25 Hz latent rate: 1 second = 25 frames
        assert cfg.duration_to_latent_frames(1.0) == 25
        assert cfg.duration_to_latent_frames(30.0) == 750

    def test_context_channels(self):
        """Context latents should be 2 * acoustic_dim."""
        cfg = _small_config()
        acoustic_dim = cfg.dit.audio_acoustic_hidden_dim
        context_channels = cfg.dit.in_channels - acoustic_dim
        # in_channels = acoustic + src_latent + chunk_mask = 3 * acoustic_dim
        assert context_channels == 2 * acoustic_dim

    def test_vae_hop_length(self):
        """Default VAE hop length is product of downsampling ratios."""
        cfg = AceStepVAEConfig()
        # 2 * 4 * 4 * 6 * 10 = 1920
        assert cfg.hop_length == 1920

    def test_small_vae_hop_length(self):
        cfg = _small_vae_config()
        # 2 * 4 = 8
        assert cfg.hop_length == 8

    def test_latent_dim_matches_acoustic_dim_full(self):
        """Full config: VAE latent_dim should equal DiT acoustic_dim."""
        full_dit = AceStepDiTConfig()
        full_vae = AceStepVAEConfig()
        assert full_vae.decoder_input_channels == full_dit.audio_acoustic_hidden_dim

    def test_in_channels_decomposition_full(self):
        """Full config: in_channels = 3 * acoustic_dim = 192."""
        full_dit = AceStepDiTConfig()
        assert full_dit.in_channels == 192
        assert full_dit.audio_acoustic_hidden_dim == 64
        assert full_dit.in_channels == 3 * full_dit.audio_acoustic_hidden_dim


# ---------------------------------------------------------------------------
# Load models (structural test with mocked disk)
# ---------------------------------------------------------------------------


class TestPipelineCacheIntegration:
    """Verify the denoise_fn closure uses cache correctly."""

    def test_denoise_fn_uses_cache(self):
        """Cache should be populated after first call and reused after."""
        from mlx_music.models.acestep.cache import CrossAttentionCache

        cfg = _small_dit_config()
        model = AceStepConditionGenerationModel(cfg)
        mx.eval(model.parameters())

        T = 4
        enc_len = 6
        encoder_hs = mx.random.normal((1, enc_len, cfg.hidden_size))
        context = mx.random.normal(
            (1, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )
        mx.eval(encoder_hs, context)

        cache = CrossAttentionCache()

        def denoise_fn(xt, t):
            return model.denoise_step(xt, t, encoder_hs, context, cache=cache)

        xt = mx.random.normal((1, T, cfg.audio_acoustic_hidden_dim))
        mx.eval(xt)

        # Step 0: should populate cache
        _ = denoise_fn(xt, 1.0)
        mx.eval(_)
        assert len(cache) == cfg.num_hidden_layers

        # Step 1: should reuse cache (same len, no new layers)
        _ = denoise_fn(xt, 0.9)
        mx.eval(_)
        assert len(cache) == cfg.num_hidden_layers


class TestLoadModels:
    """Test the load_models method with a real temp directory."""

    def test_load_from_empty_dir_creates_models(self):
        """load_models should create model instances even with no weights."""
        with tempfile.TemporaryDirectory() as tmpdir:
            dit_dir = Path(tmpdir) / "dit"
            vae_dir = Path(tmpdir) / "vae"
            dit_dir.mkdir()
            vae_dir.mkdir()

            pipeline = AceStepPipeline(tmpdir)

            # Mock the text encoder loading (requires network)
            with patch(
                "mlx_music.models.acestep.pipeline.load_text_encoder"
            ) as mock_load:
                mock_load.return_value = (MagicMock(), MagicMock())
                pipeline.load_models()

            assert pipeline._loaded
            assert pipeline.dit_model is not None
            assert pipeline.vae_model is not None
            assert pipeline.config is not None

    def test_config_defaults_when_no_json(self):
        """Without config.json, defaults should be used."""
        with tempfile.TemporaryDirectory() as tmpdir:
            dit_dir = Path(tmpdir) / "dit"
            vae_dir = Path(tmpdir) / "vae"
            dit_dir.mkdir()
            vae_dir.mkdir()

            pipeline = AceStepPipeline(tmpdir)

            with patch(
                "mlx_music.models.acestep.pipeline.load_text_encoder"
            ) as mock_load:
                mock_load.return_value = (MagicMock(), MagicMock())
                pipeline.load_models()

            # Should use default config values
            assert pipeline.config.dit.hidden_size == 2048
            assert pipeline.config.vae.hop_length == 1920
            assert pipeline.config.latent_rate == 25
