"""AceStep end-to-end music generation pipeline for MLX.

Wires together the text encoder (Qwen3-Embedding-0.6B via mlx-lm),
the DiT denoiser, the Oobleck VAE decoder, and the flow-matching
sampler into a single callable that goes from text prompt to WAV file.

Usage::

    pipeline = AceStepPipeline("./weights")
    pipeline.load_models()
    path = pipeline.generate(
        caption="upbeat jazz fusion, saxophone solo, 120 bpm",
        lyrics="[verse]\\nWalking down the avenue...",
        duration=30.0,
    )
"""

from pathlib import Path
from typing import Callable, Optional, Tuple

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

from .config import AceStepConfig, AceStepDiTConfig, AceStepVAEConfig, load_config
from .dit import AceStepConditionGenerationModel
from .text_encoder import encode_lyrics, encode_text, load_text_encoder
from .vae import AutoencoderOobleck
from ...core.audio_io import load_wav, save_wav
from ...core.sampler import prepare_noise, sample_loop, time_schedule


class AceStepPipeline:
    """End-to-end music generation pipeline.

    Holds references to the three neural components (DiT, VAE, text
    encoder) and orchestrates the full text-to-music flow:

    1. Encode caption and lyrics into conditioning tensors.
    2. Build context latents (src_latent zeros + chunk_mask ones).
    3. Generate initial noise and run the flow-matching ODE.
    4. Decode the denoised latent through the VAE.
    5. Save the resulting stereo audio as a WAV file.

    All heavy model weights are loaded lazily via :meth:`load_models`
    so that the constructor is cheap.
    """

    def __init__(self, weights_dir: str):
        """Store paths without loading anything.

        Args:
            weights_dir: Root directory containing converted MLX weights.
                Expected layout::

                    weights_dir/
                        dit/model.safetensors
                        dit/config.json
                        vae/model.safetensors
                        vae/config.json
        """
        self.weights_dir = Path(weights_dir)
        self.config: Optional[AceStepConfig] = None
        self.dit_model: Optional[AceStepConditionGenerationModel] = None
        self.vae_model: Optional[AutoencoderOobleck] = None
        self.text_model: Optional[nn.Module] = None
        self.tokenizer = None
        self.silence_latent: Optional[mx.array] = None
        self._loaded = False

    # ------------------------------------------------------------------
    # Weight loading
    # ------------------------------------------------------------------

    def load_models(self) -> None:
        """Load DiT, VAE, and text encoder into memory.

        DiT weights are read from ``weights_dir/dit/``, VAE weights
        from ``weights_dir/vae/``, and the text encoder is fetched
        via ``mlx-lm`` (Qwen3-Embedding-0.6B).  Weight tensors are
        evaluated eagerly so that any shape mismatch surfaces
        immediately.
        """
        dit_dir = self.weights_dir / "dit"
        vae_dir = self.weights_dir / "vae"

        # --- Config -------------------------------------------------
        dit_config_path = dit_dir / "config.json"
        vae_config_path = vae_dir / "config.json"

        self.config = load_config(
            dit_config_path=str(dit_config_path) if dit_config_path.exists() else None,
            vae_config_path=str(vae_config_path) if vae_config_path.exists() else None,
        )

        # --- DiT ----------------------------------------------------
        self.dit_model = AceStepConditionGenerationModel(self.config.dit)
        dit_weights_path = dit_dir / "model.safetensors"
        if dit_weights_path.exists():
            weights = mx.load(str(dit_weights_path))
            self.dit_model.load_weights(list(weights.items()))
            mx.eval(self.dit_model.parameters())

        # --- VAE ----------------------------------------------------
        self.vae_model = AutoencoderOobleck(self.config.vae)
        vae_weights_path = vae_dir / "model.safetensors"
        if vae_weights_path.exists():
            weights = mx.load(str(vae_weights_path))
            self.vae_model.load_weights(list(weights.items()))
            mx.eval(self.vae_model.parameters())

        # --- Silence latent -----------------------------------------
        silence_path = dit_dir / "silence_latent.safetensors"
        if silence_path.exists():
            sl = mx.load(str(silence_path))
            # Expected key: "silence_latent", shape [1, T, 64]
            key = next(iter(sl))
            self.silence_latent = sl[key]
            mx.eval(self.silence_latent)

        # --- Text encoder -------------------------------------------
        self.text_model, self.tokenizer = load_text_encoder(
            self.config.text_encoder_model
        )

        self._loaded = True

    # ------------------------------------------------------------------
    # Condition encoding
    # ------------------------------------------------------------------

    def encode_conditions(
        self,
        caption: str,
        lyrics: str = "",
        duration: float = 30.0,
    ) -> Tuple[mx.array, mx.array, mx.array]:
        """Encode text caption and lyrics into conditioning tensors.

        Steps:

        1. Run the text encoder on the caption to get
           ``text_hidden_states`` and ``text_mask``.
        2. Run the text encoder on the lyrics (or a single space when
           lyrics are empty) to get ``lyric_hidden_states`` and
           ``lyric_mask``.
        3. Build ``context_latents`` of shape
           ``[1, latent_frames, 128]`` where the first 64 channels are
           zeros (src_latent -- no cover/repaint source) and the last
           64 channels are ones (chunk_mask -- generate everything).
        4. Call ``model.prepare_condition`` to pack and project the
           conditioning sequence.

        Args:
            caption: Descriptive text prompt (genre, mood, tempo, etc.).
            lyrics: Song lyrics. Use ``""`` for instrumental.
            duration: Target duration in seconds.

        Returns:
            ``(encoder_hidden_states, encoder_mask, context_latents)``
        """
        self._ensure_loaded()

        # 1. Encode caption
        text_hidden_states, text_mask = encode_text(
            self.text_model, self.tokenizer, caption
        )
        mx.eval(text_hidden_states, text_mask)

        # 2. Encode lyrics (use a single space for instrumental)
        lyric_text = lyrics if lyrics else " "
        lyric_hidden_states, lyric_mask = encode_lyrics(
            self.text_model, self.tokenizer, lyric_text
        )
        mx.eval(lyric_hidden_states, lyric_mask)

        # 3. Build context latents: [1, T_latent, 128]
        #    For text-to-music: src_latent = silence_latent (NOT zeros)
        #    The silence_latent is a learned embedding that the model expects
        #    when no source audio is provided. Using zeros breaks output.
        latent_frames = self.config.duration_to_latent_frames(duration)
        acoustic_dim = self.config.dit.audio_acoustic_hidden_dim  # 64

        if self.silence_latent is not None:
            src_latent = self.silence_latent[:, :latent_frames, :]
        else:
            src_latent = mx.zeros((1, latent_frames, acoustic_dim))
        chunk_mask = mx.ones((1, latent_frames, acoustic_dim))
        context_latents = mx.concatenate([src_latent, chunk_mask], axis=-1)

        # 4. Pack conditions via the DiT encoder
        encoder_hidden_states, encoder_mask, context_latents = (
            self.dit_model.prepare_condition(
                text_hidden_states=text_hidden_states,
                text_attention_mask=text_mask,
                lyric_hidden_states=lyric_hidden_states,
                lyric_attention_mask=lyric_mask,
                context_latents=context_latents,
            )
        )
        mx.eval(encoder_hidden_states, encoder_mask, context_latents)

        return encoder_hidden_states, encoder_mask, context_latents

    # ------------------------------------------------------------------
    # Generation
    # ------------------------------------------------------------------

    def generate(
        self,
        caption: str,
        lyrics: str = "",
        duration: float = 30.0,
        steps: int = 8,
        cfg_scale: float = 1.0,
        shift: float = 3.0,
        seed: Optional[int] = None,
        output: str = "output.wav",
        callback: Optional[Callable[[int, int, mx.array], None]] = None,
    ) -> str:
        """Full text-to-music generation pipeline.

        Args:
            caption: Text prompt describing the music.
            lyrics: Lyrics for the song. ``""`` for instrumental.
            duration: Target duration in seconds (default 30).
            steps: Number of denoising steps (default 8 for turbo).
            cfg_scale: Classifier-free guidance scale. ``1.0`` disables
                guidance (turbo model has CFG baked in).
            shift: Timestep schedule shift (default 3.0 for turbo).
            seed: Random seed for reproducibility. ``None`` for random.
            output: Output WAV file path.
            callback: Optional ``(step, total, xt)`` progress callback.

        Returns:
            The output file path as a string.
        """
        self._ensure_loaded()

        # 1. Encode conditions
        encoder_hidden_states, encoder_mask, context_latents = self.encode_conditions(
            caption, lyrics, duration
        )

        # 2. Compute latent shape
        latent_frames = self.config.duration_to_latent_frames(duration)
        acoustic_dim = self.config.dit.audio_acoustic_hidden_dim
        latent_shape = (1, latent_frames, acoustic_dim)

        # 3. Prepare noise
        noise = prepare_noise(latent_shape, seed=seed)
        mx.eval(noise)

        # 4. Time schedule
        schedule = time_schedule(steps, shift=shift)

        # 5. Build denoise function wrapping the DiT (with cross-attention cache)
        from .cache import CrossAttentionCache

        cache = CrossAttentionCache()

        def denoise_fn(xt: mx.array, t: float) -> mx.array:
            return self.dit_model.denoise_step(
                xt=xt,
                t=t,
                encoder_hidden_states=encoder_hidden_states,
                context_latents=context_latents,
                cache=cache,
            )

        # 6. Run the denoising loop
        latent = sample_loop(
            denoise_fn=denoise_fn,
            noise=noise,
            schedule=schedule,
            method="ode",
            callback=callback,
        )
        mx.eval(latent)

        # 7. VAE decode: latent [1, T_latent, 64] -> audio [1, T_samples, 2]
        audio = self.vae_model.decode(latent)
        mx.eval(audio)

        # 8. Convert to numpy and save
        audio_np = np.array(audio[0])  # [T_samples, 2]
        save_wav(audio_np, output, sample_rate=self.config.sample_rate)

        return output

    # ------------------------------------------------------------------
    # Audio encoding (for repaint/cover)
    # ------------------------------------------------------------------

    def encode_audio(self, audio_path: str) -> mx.array:
        """Load audio and encode through VAE to latent space.

        Args:
            audio_path: Path to a WAV/FLAC/OGG audio file.

        Returns:
            Latent tensor of shape ``[1, T_latent, 64]``.
        """
        self._ensure_loaded()
        audio_np, sr = load_wav(audio_path)  # [samples, 2] @ 48kHz
        audio_mx = mx.array(audio_np)[None]  # [1, samples, 2]
        latent, _ = self.vae_model.encode(audio_mx, deterministic=True)
        mx.eval(latent)
        return latent

    # ------------------------------------------------------------------
    # Repaint generation
    # ------------------------------------------------------------------

    def generate_repaint(
        self,
        src_audio: str,
        caption: str,
        lyrics: str = "",
        repaint_start: float = 0.0,
        repaint_end: float = -1.0,
        steps: int = 8,
        cfg_scale: float = 1.0,
        shift: float = 3.0,
        seed: Optional[int] = None,
        output: str = "output.wav",
        callback: Optional[Callable[[int, int, mx.array], None]] = None,
    ) -> str:
        """Generate music by repainting a region of source audio.

        Encodes the source audio to latent space, then regenerates the
        specified time region using the DiT while preserving the rest.
        When ``repaint_start=0`` and ``repaint_end=-1`` (full duration),
        this acts as a full remix guided by the source audio's structure.

        Args:
            src_audio: Path to source audio file.
            caption: Text prompt for the generated content.
            lyrics: Lyrics for the song. ``""`` for instrumental.
            repaint_start: Start of region to regenerate (seconds).
            repaint_end: End of region (seconds). ``-1`` for full duration.
            steps: Number of denoising steps.
            cfg_scale: Classifier-free guidance scale.
            shift: Timestep schedule shift.
            seed: Random seed for reproducibility.
            output: Output WAV file path.
            callback: Optional ``(step, total, xt)`` progress callback.

        Returns:
            The output file path as a string.
        """
        self._ensure_loaded()

        # 1. Encode source audio to latent space
        target_latents = self.encode_audio(src_audio)  # [1, T, 64]
        latent_frames = target_latents.shape[1]
        acoustic_dim = self.config.dit.audio_acoustic_hidden_dim  # 64
        duration = self.config.latent_frames_to_duration(latent_frames)
        latent_rate = self.config.latent_rate  # 25 Hz

        # 2. Build chunk_mask: 1 = regenerate, 0 = keep original
        if repaint_end < 0:
            repaint_end = duration
        start_frame = int(repaint_start * latent_rate)
        end_frame = int(repaint_end * latent_rate)
        start_frame = max(0, min(start_frame, latent_frames))
        end_frame = max(start_frame, min(end_frame, latent_frames))

        mask_np = np.zeros((1, latent_frames, acoustic_dim), dtype=np.float32)
        mask_np[:, start_frame:end_frame, :] = 1.0
        chunk_mask = mx.array(mask_np)

        # 3. Build src_latents: source audio outside repaint region,
        #    silence_latent inside repaint region
        if self.silence_latent is not None:
            silence = self.silence_latent[:, :latent_frames, :]
        else:
            silence = mx.zeros_like(target_latents)
        # Where mask=1 (repaint region), use silence; where mask=0, use source
        mask_bool = chunk_mask[:, :, :1]  # [1, T, 1] — broadcast across channels
        src_latents = target_latents * (1.0 - mask_bool) + silence * mask_bool

        # 4. Concatenate: [src_latents, chunk_mask] -> [1, T, 128]
        context_latents = mx.concatenate([src_latents, chunk_mask], axis=-1)

        # 5. Encode text conditions
        text_hidden_states, text_mask = encode_text(
            self.text_model, self.tokenizer, caption
        )
        mx.eval(text_hidden_states, text_mask)

        lyric_text = lyrics if lyrics else " "
        lyric_hidden_states, lyric_mask = encode_lyrics(
            self.text_model, self.tokenizer, lyric_text
        )
        mx.eval(lyric_hidden_states, lyric_mask)

        # 6. Pack conditions via the DiT encoder
        encoder_hidden_states, encoder_mask, context_latents = (
            self.dit_model.prepare_condition(
                text_hidden_states=text_hidden_states,
                text_attention_mask=text_mask,
                lyric_hidden_states=lyric_hidden_states,
                lyric_attention_mask=lyric_mask,
                context_latents=context_latents,
            )
        )
        mx.eval(encoder_hidden_states, encoder_mask, context_latents)

        # 7. Prepare noise and schedule
        latent_shape = (1, latent_frames, acoustic_dim)
        noise = prepare_noise(latent_shape, seed=seed)
        mx.eval(noise)
        schedule = time_schedule(steps, shift=shift)

        # 8. Build denoise function
        from .cache import CrossAttentionCache

        cache = CrossAttentionCache()

        def denoise_fn(xt: mx.array, t: float) -> mx.array:
            return self.dit_model.denoise_step(
                xt=xt,
                t=t,
                encoder_hidden_states=encoder_hidden_states,
                context_latents=context_latents,
                cache=cache,
            )

        # 9. Run denoising loop
        latent = sample_loop(
            denoise_fn=denoise_fn,
            noise=noise,
            schedule=schedule,
            method="ode",
            callback=callback,
        )
        mx.eval(latent)

        # 10. VAE decode and save
        audio = self.vae_model.decode(latent)
        mx.eval(audio)
        audio_np = np.array(audio[0])  # [T_samples, 2]
        save_wav(audio_np, output, sample_rate=self.config.sample_rate)

        return output

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _ensure_loaded(self) -> None:
        """Raise if models have not been loaded yet."""
        if not self._loaded:
            raise RuntimeError("Models not loaded. Call pipeline.load_models() first.")
