"""Tests for models/acestep/text_encoder.py

All tests use mocks so they run without downloading the Qwen3-Embedding-0.6B
model weights.
"""

from unittest.mock import MagicMock, patch

import mlx.core as mx
import mlx.nn as nn

from mlx_music.models.acestep.text_encoder import (
    DEFAULT_TEXT_ENCODER,
    _get_hidden_states,
    _tokenize,
    encode_lyrics,
    encode_text,
    load_text_encoder,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

HIDDEN_SIZE = 1024


def _make_mock_model(hidden_size: int = HIDDEN_SIZE) -> MagicMock:
    """Build a mock that mimics the mlx-lm Model / Qwen3Model structure.

    The mock's backbone (model.model) has:
        - embed_tokens: returns random embeddings
        - layers: a list with one mock layer (identity-ish)
        - norm: identity
    """
    # Inner backbone (Qwen3Model)
    backbone = MagicMock()

    def _embed(input_ids):
        B, S = input_ids.shape
        return mx.zeros((B, S, hidden_size))

    backbone.embed_tokens = _embed

    # Single mock transformer layer: just returns its input unchanged.
    mock_layer = MagicMock(side_effect=lambda h, mask=None, cache=None: h)
    backbone.layers = [mock_layer]

    # Final norm: identity
    backbone.norm = MagicMock(side_effect=lambda h: h)

    # Outer Model wrapper
    model = MagicMock()
    model.model = backbone
    return model


def _make_mock_tokenizer(token_ids: list[int] | None = None) -> MagicMock:
    """Build a mock tokenizer whose encode() returns *token_ids*."""
    tok = MagicMock()
    tok.encode.return_value = token_ids if token_ids is not None else [1, 2, 3]
    return tok


# ---------------------------------------------------------------------------
# Tests: load_text_encoder
# ---------------------------------------------------------------------------


class TestLoadTextEncoder:
    def test_load_calls_mlx_lm(self):
        """load_text_encoder delegates to mlx_lm utilities."""
        sentinel_model = MagicMock()
        sentinel_tok = MagicMock()
        mock_path = MagicMock()

        with (
            patch(
                "mlx_music.models.acestep.text_encoder.hf_repo_to_path",
                return_value=mock_path,
            ) as mock_hf,
            patch(
                "mlx_music.models.acestep.text_encoder.load_model",
                return_value=(sentinel_model, {}),
            ) as mock_load_model,
            patch(
                "mlx_music.models.acestep.text_encoder.tokenizer_utils",
            ) as mock_tok_utils,
            patch(
                "mlx_music.models.acestep.text_encoder.mx",
            ),
        ):
            mock_tok_utils.load.return_value = sentinel_tok
            # Mock Path.glob for weight file check
            mock_path.glob = MagicMock(return_value=[])
            model, tok = load_text_encoder()
            mock_hf.assert_called_once_with(DEFAULT_TEXT_ENCODER)
            assert model is sentinel_model

    def test_load_custom_model_id(self):
        """A custom model_id is forwarded."""
        mock_path = MagicMock()
        mock_path.glob = MagicMock(return_value=[])

        with (
            patch(
                "mlx_music.models.acestep.text_encoder.hf_repo_to_path",
                return_value=mock_path,
            ) as mock_hf,
            patch(
                "mlx_music.models.acestep.text_encoder.load_model",
                return_value=(MagicMock(), {}),
            ),
            patch(
                "mlx_music.models.acestep.text_encoder.tokenizer_utils",
            ),
            patch(
                "mlx_music.models.acestep.text_encoder.mx",
            ),
        ):
            load_text_encoder("my-org/my-model")
            mock_hf.assert_called_once_with("my-org/my-model")


# ---------------------------------------------------------------------------
# Tests: _tokenize
# ---------------------------------------------------------------------------


class TestTokenize:
    def test_basic_shapes(self):
        tok = _make_mock_tokenizer([10, 20, 30, 40])
        input_ids, attn_mask = _tokenize(tok, "hello world", max_length=512)

        assert input_ids.shape == (1, 4)
        assert attn_mask.shape == (1, 4)

    def test_attention_mask_all_ones(self):
        """Without padding, the attention mask should be all 1s."""
        tok = _make_mock_tokenizer([1, 2, 3])
        _, attn_mask = _tokenize(tok, "hi", max_length=512)
        assert mx.array_equal(attn_mask, mx.ones((1, 3)))

    def test_truncation(self):
        """Tokens beyond max_length are truncated."""
        long_ids = list(range(600))
        tok = _make_mock_tokenizer(long_ids)
        input_ids, attn_mask = _tokenize(tok, "long text", max_length=512)

        assert input_ids.shape == (1, 512)
        assert attn_mask.shape == (1, 512)
        # Verify the kept tokens are the first 512
        assert input_ids[0, 0].item() == 0
        assert input_ids[0, 511].item() == 511

    def test_short_input_no_truncation(self):
        tok = _make_mock_tokenizer([5, 6])
        input_ids, _ = _tokenize(tok, "ok", max_length=512)
        assert input_ids.shape == (1, 2)

    def test_dtype(self):
        tok = _make_mock_tokenizer([1])
        input_ids, attn_mask = _tokenize(tok, "x", max_length=10)
        assert input_ids.dtype == mx.int32
        assert attn_mask.dtype == mx.float32


# ---------------------------------------------------------------------------
# Tests: _get_hidden_states
# ---------------------------------------------------------------------------


class TestGetHiddenStates:
    def test_output_shape(self):
        model = _make_mock_model()
        input_ids = mx.array([[1, 2, 3]], dtype=mx.int32)
        h = _get_hidden_states(model, input_ids)
        assert h.shape == (1, 3, HIDDEN_SIZE)

    def test_batch_dimension(self):
        model = _make_mock_model()
        input_ids = mx.array([[1, 2], [3, 4]], dtype=mx.int32)
        h = _get_hidden_states(model, input_ids)
        assert h.shape == (2, 2, HIDDEN_SIZE)

    def test_bidirectional_mask(self):
        """The transformer layers must be called with mask=None for
        bidirectional attention (not causal)."""
        model = _make_mock_model()
        input_ids = mx.array([[1, 2, 3]], dtype=mx.int32)
        _get_hidden_states(model, input_ids)

        layer = model.model.layers[0]
        layer.assert_called_once()
        _, kwargs = layer.call_args
        assert kwargs.get("mask") is None
        assert kwargs.get("cache") is None

    def test_norm_is_applied(self):
        """The backbone's final norm must be called on the output."""
        model = _make_mock_model()
        input_ids = mx.array([[1]], dtype=mx.int32)
        _get_hidden_states(model, input_ids)
        model.model.norm.assert_called_once()


# ---------------------------------------------------------------------------
# Tests: encode_text
# ---------------------------------------------------------------------------


class TestEncodeText:
    def test_output_shapes(self):
        model = _make_mock_model()
        tok = _make_mock_tokenizer([1, 2, 3, 4, 5])
        hidden, mask = encode_text(model, tok, "a pop song about rain")

        assert hidden.shape == (1, 5, HIDDEN_SIZE)
        assert mask.shape == (1, 5)

    def test_default_max_length(self):
        """encode_text should use max_length=512 by default."""
        long_ids = list(range(1000))
        model = _make_mock_model()
        tok = _make_mock_tokenizer(long_ids)
        hidden, mask = encode_text(model, tok, "very long caption")

        assert hidden.shape == (1, 512, HIDDEN_SIZE)
        assert mask.shape == (1, 512)

    def test_custom_max_length(self):
        ids = list(range(200))
        model = _make_mock_model()
        tok = _make_mock_tokenizer(ids)
        hidden, mask = encode_text(model, tok, "caption", max_length=77)

        assert hidden.shape == (1, 77, HIDDEN_SIZE)

    def test_mask_values(self):
        model = _make_mock_model()
        tok = _make_mock_tokenizer([10, 20, 30])
        _, mask = encode_text(model, tok, "test")
        # All real tokens -> mask is all 1s
        assert mx.all(mask == 1.0).item()


# ---------------------------------------------------------------------------
# Tests: encode_lyrics
# ---------------------------------------------------------------------------


class TestEncodeLyrics:
    def test_output_shapes(self):
        model = _make_mock_model()
        tok = _make_mock_tokenizer([1, 2, 3])
        hidden, mask = encode_lyrics(model, tok, "verse 1\nchorus\nverse 2")

        assert hidden.shape == (1, 3, HIDDEN_SIZE)
        assert mask.shape == (1, 3)

    def test_default_max_length_is_4096(self):
        """encode_lyrics should handle up to 4096 tokens by default."""
        long_ids = list(range(5000))
        model = _make_mock_model()
        tok = _make_mock_tokenizer(long_ids)
        hidden, mask = encode_lyrics(model, tok, "long lyrics")

        assert hidden.shape == (1, 4096, HIDDEN_SIZE)
        assert mask.shape == (1, 4096)

    def test_short_lyrics(self):
        model = _make_mock_model()
        tok = _make_mock_tokenizer([7])
        hidden, mask = encode_lyrics(model, tok, "la")
        assert hidden.shape == (1, 1, HIDDEN_SIZE)

    def test_custom_max_length(self):
        ids = list(range(200))
        model = _make_mock_model()
        tok = _make_mock_tokenizer(ids)
        hidden, mask = encode_lyrics(model, tok, "lyrics", max_length=123)
        assert hidden.shape == (1, 123, HIDDEN_SIZE)

    def test_hidden_dim_matches_config(self):
        """Hidden dim must be 1024 to match AceStepDiTConfig.text_hidden_dim."""
        from mlx_music.models.acestep.config import AceStepDiTConfig

        assert AceStepDiTConfig().text_hidden_dim == HIDDEN_SIZE
