"""Text encoder for ACE-Step using Qwen3-Embedding-0.6B via mlx-lm.

Wraps the Qwen3 embedding model to produce bidirectional hidden states
for both text captions and lyrics. The same encoder serves both purposes;
lyrics simply use a larger max_length.

Output shapes consumed by the DiT:
    text_hidden_states:  [B, seq_len, 1024]
    text_attention_mask: [B, seq_len]
    lyric_hidden_states: [B, seq_len, 1024]
    lyric_attention_mask:[B, seq_len]
"""

from pathlib import Path
from typing import Optional, Tuple

import mlx.core as mx
import mlx.nn as nn
from mlx_lm import tokenizer_utils
from mlx_lm.utils import hf_repo_to_path, load_model

# Default model for text encoding
DEFAULT_TEXT_ENCODER = "Qwen/Qwen3-Embedding-0.6B"


def load_text_encoder(
    model_id: str = DEFAULT_TEXT_ENCODER,
) -> Tuple[nn.Module, object]:
    """Load the Qwen3-Embedding-0.6B model and tokenizer via mlx-lm.

    The embedding model stores weights without the ``model.`` prefix that
    mlx-lm expects (e.g. ``layers.0.*`` instead of ``model.layers.0.*``).
    We load with ``strict=False`` then re-load with the correct prefix.

    Args:
        model_id: HuggingFace model ID or local path. Defaults to
            ``Qwen/Qwen3-Embedding-0.6B``.

    Returns:
        (model, tokenizer) -- the mlx-lm ``Model`` and ``TokenizerWrapper``.
        ``model.model`` is the inner ``Qwen3Model`` that returns hidden states.
    """
    model_path = hf_repo_to_path(model_id)

    # Load with strict=False — the embedding model stores weights without
    # the 'model.' prefix that mlx-lm expects. We'll re-load with prefix.
    model, config = load_model(model_path, lazy=False, strict=False)

    # Re-load weights with 'model.' prefix
    weight_files = sorted(Path(model_path).glob("*.safetensors"))
    weights = {}
    for f in weight_files:
        weights.update(mx.load(str(f)))

    if weights:
        sample_key = next(iter(weights))
        if not sample_key.startswith("model."):
            prefixed = {f"model.{k}": v for k, v in weights.items()}
            model.load_weights(list(prefixed.items()))

    tokenizer = tokenizer_utils.load(model_path)
    return model, tokenizer


def _get_hidden_states(
    model: nn.Module,
    input_ids: mx.array,
) -> mx.array:
    """Run the transformer backbone with bidirectional attention and return
    the final-layer hidden states.

    This bypasses the ``lm_head`` projection and, critically, replaces the
    default causal mask with ``None`` so that every token can attend to every
    other token (bidirectional attention required for embedding models).

    Args:
        model: The outer ``Model`` object returned by ``mlx_lm.load``.
        input_ids: Integer token IDs, shape ``[B, seq_len]``.

    Returns:
        Hidden states of shape ``[B, seq_len, hidden_size]``.
    """
    backbone = model.model  # Qwen3Model

    h = backbone.embed_tokens(input_ids)

    # Bidirectional: pass mask=None so every position attends to all others.
    # The mlx-lm Qwen3Model normally builds a causal mask when cache is None,
    # so we call each layer directly with mask=None.
    for layer in backbone.layers:
        h = layer(h, mask=None, cache=None)

    h = backbone.norm(h)
    return h


def _tokenize(
    tokenizer,
    text: str,
    max_length: int,
) -> Tuple[mx.array, mx.array]:
    """Tokenize a single string with truncation and padding metadata.

    Returns:
        (input_ids, attention_mask) -- both as MLX arrays with a leading
        batch dimension of 1.
    """
    # The TokenizerWrapper delegates to the underlying HF tokenizer.
    encoded = tokenizer.encode(text)

    # Truncate if necessary (keep room for nothing extra -- the HF tokenizer
    # may or may not add special tokens depending on config).
    if len(encoded) > max_length:
        encoded = encoded[:max_length]

    input_ids = mx.array(encoded, dtype=mx.int32)[None, :]  # [1, seq_len]
    attention_mask = mx.ones(input_ids.shape, dtype=mx.float32)  # [1, seq_len]

    return input_ids, attention_mask


def encode_text(
    model: nn.Module,
    tokenizer,
    text: str,
    max_length: int = 512,
) -> Tuple[mx.array, mx.array]:
    """Encode a text caption into hidden states.

    Args:
        model: Model returned by :func:`load_text_encoder`.
        tokenizer: Tokenizer returned by :func:`load_text_encoder`.
        text: The text caption to encode.
        max_length: Maximum token length (default 512).

    Returns:
        (hidden_states, attention_mask)
            hidden_states:  ``[1, seq_len, 1024]``
            attention_mask: ``[1, seq_len]``
    """
    input_ids, attention_mask = _tokenize(tokenizer, text, max_length)
    hidden_states = _get_hidden_states(model, input_ids)
    return hidden_states, attention_mask


def encode_lyrics(
    model: nn.Module,
    tokenizer,
    lyrics: str,
    max_length: int = 4096,
) -> Tuple[mx.array, mx.array]:
    """Encode lyrics into hidden states.

    Identical to :func:`encode_text` but with a larger default
    ``max_length`` to accommodate full song lyrics.

    Args:
        model: Model returned by :func:`load_text_encoder`.
        tokenizer: Tokenizer returned by :func:`load_text_encoder`.
        lyrics: The lyrics string to encode.
        max_length: Maximum token length (default 4096).

    Returns:
        (hidden_states, attention_mask)
            hidden_states:  ``[1, seq_len, 1024]``
            attention_mask: ``[1, seq_len]``
    """
    input_ids, attention_mask = _tokenize(tokenizer, lyrics, max_length)
    hidden_states = _get_hidden_states(model, input_ids)
    return hidden_states, attention_mask
