"""MiniMax Music 3 inference engine for ZeroGPU.

Ported from the official MiniMaxAI/MiniMax-Music3 Space (diffusers ModularPipeline +
AoTI kernels compiled for the RTX Pro 6000 ZeroGPU hardware). Everything here runs at
import time (model load + AoTI binding) except `generate_wav`, which must be called from
inside a `@spaces.GPU` function.
"""

import copy as _copy
import time

import spaces  # must precede torch (ZeroGPU patches CUDA init)
import numpy as np
import torch
import torch.nn as _nn
from huggingface_hub import snapshot_download

from diffusers import ModularPipeline
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from transformers import StaticCache
from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM

MODEL_ID = "MiniMaxAI/MiniMax-Music3"
AOTI_ID = "diffusers-internal-dev/MiniMax-Music3-aoti"

PIPE = ModularPipeline.from_pretrained(MODEL_ID)
PIPE.load_components(dtype=torch.bfloat16)
PIPE.to("cuda")

SAMPLE_RATE = int(PIPE.sampling_rate)
FRAME_RATE = float(PIPE.frame_rate)
MAX_FRAMES = 9000


def _encode_prompt(caption, lyrics, device):
    # Mirrors the modular TextEncoderStep: the AR stage is driven manually below.
    import diffusers.modular_pipelines.minimax_music3.encoders as P

    text = (
        f"{P._IM_START}{P._CAPTION_START}{P._clean_caption(caption)}{P._CAPTION_END}"
        f"{P._LYRICS_START}{P._normalize_lyrics(lyrics)}{P._LYRICS_END}{P._IM_END}{P._AUDIO_START}"
    )
    input_ids = PIPE.tokenizer(text, return_tensors="pt")["input_ids"]
    if input_ids.shape[1] > P._MAX_PROMPT_TOKENS:
        raise ValueError(
            f"The assembled prompt has {input_ids.shape[1]} tokens; the maximum is {P._MAX_PROMPT_TOKENS}."
        )
    unconditional_ids = input_ids.clone()
    unconditional_ids[:, 1:-2] = P._AUDIO_CFG_TOKEN_ID
    return torch.cat((input_ids, unconditional_ids), dim=0).to(device)


# ── AoTI kernels ─────────────────────────────────────────────────────────────
# The transformer artifact is static over full 689-latent chunks; the once-per-song
# final short chunk falls back to eager.
_AOTI_DIR = snapshot_download(AOTI_ID)
_eager_transformer_forward = PIPE.transformer.forward
spaces.aoti_load_from_package_dir(PIPE.transformer, f"{_AOTI_DIR}/transformer")
_aoti_transformer_forward = PIPE.transformer.forward


def _guarded_transformer_forward(hidden_states, timestep, encoder_hidden_states, return_dict=True):
    if hidden_states.shape[-1] == 689:
        out = _aoti_transformer_forward(hidden_states, timestep, encoder_hidden_states)
        if not isinstance(out, Transformer2DModelOutput):
            out = Transformer2DModelOutput(sample=out[0] if isinstance(out, (tuple, list)) else out)
        return out
    return _eager_transformer_forward(hidden_states, timestep, encoder_hidden_states, return_dict=return_dict)


PIPE.transformer.forward = _guarded_transformer_forward
spaces.aoti_load_from_package_dir(PIPE.vocoder, f"{_AOTI_DIR}/vocoder")

# AoTI LM decode step, one artifact per StaticCache bucket; eager per-frame glue. Eager
# full-sequence prefill writes directly into each artifact's cache buffers (aliased
# StaticCache), matching eager exactly.
_LM = PIPE.language_model
_BUCKETS = [1024, 2048, 4096, 8192]
_STOP_CHECK_INTERVAL = 25

_lm_headless = _copy.copy(_LM)
_lm_headless._modules = dict(_LM._modules)  # nn.Module shallow copies share _modules
_lm_headless.lm_head = _nn.Identity()
_lm_headless.generation_config = _copy.deepcopy(_LM.generation_config)
_lm_headless.generation_config.cache_implementation = "static"

_LM_STEPS = {}
for _bucket in _BUCKETS:
    _exportable = TorchExportableModuleForDecoderOnlyLM(
        _lm_headless, batch_size=2, max_cache_len=_bucket, device="cuda"
    )
    for _m in _exportable.modules():
        _m._non_persistent_buffers_set.clear()
    spaces.aoti_load_from_package_dir(_exportable.model, f"{_AOTI_DIR}/lm_step_{_bucket}")
    _LM_STEPS[_bucket] = _exportable.model


def _aliased_cache(step_module, bucket):
    cache = StaticCache(max_cache_len=bucket, config=_LM.config.get_text_config())
    cache.early_initialization(
        2, _LM.config.num_key_value_heads, _LM.config.head_dim, _LM.dtype, torch.device("cuda")
    )
    for i, layer in enumerate(cache.layers):
        layer.keys = step_module.get_buffer(f"key_cache_{i}")
        layer.values = step_module.get_buffer(f"value_cache_{i}")
        layer.cumulative_length = step_module.get_buffer(f"cumulative_length_{i}")
        layer.keys.zero_()
        layer.values.zero_()
        layer.cumulative_length.zero_()
    return cache


def _hop_lm_cache(src_bucket, dst_bucket, used):
    src, dst = _LM_STEPS[src_bucket], _LM_STEPS[dst_bucket]
    for i in range(_LM.config.num_hidden_layers):
        dst.get_buffer(f"key_cache_{i}")[:, :, :used] = src.get_buffer(f"key_cache_{i}")[:, :, :used]
        dst.get_buffer(f"value_cache_{i}")[:, :, :used] = src.get_buffer(f"value_cache_{i}")[:, :, :used]
        dst.get_buffer(f"cumulative_length_{i}").copy_(src.get_buffer(f"cumulative_length_{i}"))


def _iter_frames_aoti(text_ids, max_frames, generator=None):
    import diffusers.modular_pipelines.minimax_music3.encoders as P

    prompt_len = text_ids.shape[1]
    bucket = _BUCKETS[0]
    while bucket < prompt_len + 16:
        bucket *= 2
    step = _LM_STEPS[bucket]
    cache = _aliased_cache(step, bucket)
    prompt_embeds = _LM.model.embed_tokens(text_ids)
    output = _LM.model(
        inputs_embeds=prompt_embeds,
        past_key_values=cache,
        cache_position=torch.arange(prompt_len, device="cuda"),
        use_cache=True,
    )
    last_hidden = output.last_hidden_state[:, -1]

    vocab_mask = torch.ones(_LM.config.vocab_size, dtype=torch.bool, device="cuda")
    vocab_mask[P._AUDIO_CODE_OFFSET : P._AUDIO_CODE_OFFSET + P._SEMANTIC_VOCAB_SIZE] = False
    vocab_mask[P._AUDIO_END_TOKEN_ID] = False

    emitted = 0
    position = prompt_len
    pending = []
    for frame_index in range(max_frames + 1):
        if position + 2 >= bucket:
            new_bucket = bucket * 2
            _hop_lm_cache(bucket, new_bucket, position)
            bucket = new_bucket
            step = _LM_STEPS[bucket]
        logits = _LM.lm_head(last_hidden).float()
        logits = logits.masked_fill(vocab_mask, -float("inf"))
        conditional, unconditional = logits[0:1], logits[1:2]
        guided = unconditional + (conditional - unconditional) * P._AR_CFG_SCALE
        threshold = torch.topk(conditional, P._AR_CFG_TOP_K, dim=-1).values[..., -1, None]
        guided = guided.masked_fill(conditional < threshold, -float("inf"))
        guided = guided.masked_fill(vocab_mask.unsqueeze(0), -float("inf"))
        sampled = P._sample_top_k(guided, generator)
        semantic_code = (sampled - P._AUDIO_CODE_OFFSET).clamp_min(0).repeat(2)
        frame_codes, depth_hidden = P._generate_depth_codes(PIPE, last_hidden, semantic_code, generator)
        frame_hidden = torch.cat((last_hidden[:1].clone(), depth_hidden), dim=-1) if frame_index > 0 else None
        pending.append((sampled, frame_hidden))
        if len(pending) >= _STOP_CHECK_INTERVAL or frame_index == max_frames:
            stop_flags = torch.cat([s == P._AUDIO_END_TOKEN_ID for s, _ in pending]).tolist()
            for flag, (_, fh) in zip(stop_flags, pending):
                if flag:
                    return
                if fh is not None:
                    emitted += 1
                    yield fh
                    if emitted >= max_frames:
                        return
            pending = []
        feedback = P._embed_audio_frame(PIPE, frame_codes)
        last_hidden = step(inputs_embeds=feedback, cache_position=torch.tensor([position], device="cuda"))[:, -1]
        position += 1
    for _, fh in pending:
        if fh is not None:
            yield fh


PIPE._iter_frames = _iter_frames_aoti


# ── Windowed flow-matching decode + vocoder ──────────────────────────────────
_CHUNK, _HOP, _HOP_SAMPLES = 200, 100, 86 * 512
_CROP_RIGHT_SAMPLES = (344 - 86) * 512


@torch.inference_mode()
def _decode_window(hidden_window, previous, generator, steps, guidance):
    previous_latent, previous_condition = previous
    condition = PIPE.condition_encoder(hidden_window)
    condition = condition.to(PIPE.transformer.dtype)
    latents = torch.randn(
        (1, PIPE.transformer.config.in_channels, condition.shape[1]),
        generator=generator, device="cuda", dtype=condition.dtype,
    )
    overlap, noise_prompt = 0, None
    if previous_latent is not None:
        overlap = min(previous_latent.shape[-1], latents.shape[-1])
        noise_prompt = latents[..., :overlap].clone()
        condition[:, :overlap] = previous_condition[:, :overlap]
    condition_input = torch.cat((condition, torch.zeros_like(condition)), dim=0)
    PIPE.scheduler.set_timesteps(sigmas=np.linspace(1.0, 1.0 / steps, steps), device="cuda")
    for timestep in PIPE.scheduler.timesteps:
        if overlap > 0:
            t = timestep.to(latents.dtype)
            latents[..., :overlap] = (1.0 - (1.0 - 1e-6) * t) * noise_prompt + t * previous_latent[..., :overlap]
        velocity = PIPE.transformer(
            latents.expand(2, -1, -1).contiguous(), timestep.expand(2).to(latents.dtype), condition_input
        ).sample
        velocity = velocity[1:2] + guidance * (velocity[0:1] - velocity[1:2])
        latents = PIPE.scheduler.step(velocity, timestep, latents).prev_sample
    if overlap > 0:
        latents[..., :overlap] = previous_latent[..., :overlap]
    overlap_start = max(0, latents.shape[-1] - 2 * 172)
    overlap_end = max(overlap_start, latents.shape[-1] - 172)
    carry = (latents[..., overlap_start:overlap_end], condition[:, overlap_start:overlap_end])
    waveform = PIPE.vocoder(latents.to(PIPE.vocoder.dtype)).float().clamp(-1.0, 1.0)[0]
    return waveform, carry


@torch.inference_mode()
def _stream_windows(text_ids, max_frames, ar_generator, dit_generator, steps, guidance):
    frames = []
    windows_done = 0
    carry = (None, None)
    for hidden in PIPE._iter_frames(text_ids, max_frames, ar_generator):
        frames.append(hidden)
        window_start = windows_done * _HOP
        if len(frames) > window_start + _CHUNK:
            window = torch.stack(frames[window_start : window_start + _CHUNK], dim=1)
            waveform, carry = _decode_window(window, carry, dit_generator, steps, guidance)
            left = 0 if windows_done == 0 else _HOP_SAMPLES
            windows_done += 1
            yield waveform[:, left : waveform.shape[-1] - _CROP_RIGHT_SAMPLES]
    if not frames:
        raise RuntimeError("The model generated zero audio frames; try different lyrics or a longer duration.")
    total = len(frames)
    window_starts = [0] if total <= _CHUNK else list(range(0, total - _HOP, _HOP))
    for w in range(windows_done, len(window_starts)):
        window_start = window_starts[w]
        window = torch.stack(frames[window_start : min(window_start + _CHUNK, total)], dim=1)
        waveform, carry = _decode_window(window, carry, dit_generator, steps, guidance)
        left = 0 if w == 0 else _HOP_SAMPLES
        right = _CROP_RIGHT_SAMPLES if w < len(window_starts) - 1 else 0
        yield waveform[:, left : waveform.shape[-1] - right]


def _to_int16(waveform):
    return (waveform.cpu().numpy().T * 32767.0).astype(np.int16)


def estimate_gpu_seconds(duration, steps=30):
    """Fitted on-Space (xlarge): wall = 0.75*dur + 0.20*dur*(steps/30) + ~15s cold-worker margin."""
    return min(int(float(duration) * (0.75 + 0.20 * float(steps) / 30.0) + 15), 600)


def validate(caption: str, lyrics: str) -> None:
    """CPU-side input validation (raise ValueError with a user-facing message).

    Runs before the ZeroGPU call: exceptions raised inside the GPU worker lose their message.
    """
    caption = (caption or "").strip()
    lyrics = (lyrics or "").strip()
    if not caption:
        raise ValueError("A music description (structured caption) is required.")
    if not lyrics:
        raise ValueError("Lyrics are required (use [instrumental] for songs without vocals).")
    _encode_prompt(caption, lyrics, "cpu")  # raises ValueError past the token limit


@torch.inference_mode()
def generate_wav(caption: str, lyrics: str, duration: float, seed: int, steps: int = 30, guidance: float = 1.7):
    """Run the full AR -> windowed DiT -> vocoder pipeline. Must run inside @spaces.GPU.

    Returns (int16 stereo ndarray [samples, 2], sample_rate, audio_seconds, wall_seconds).
    """
    caption = (caption or "").strip()
    lyrics = (lyrics or "").strip()
    steps, guidance = int(steps), float(guidance)
    text_ids = _encode_prompt(caption, lyrics, "cuda")
    max_frames = min(int(float(duration) * FRAME_RATE), MAX_FRAMES)
    ar_generator = torch.Generator("cuda").manual_seed(int(seed))
    dit_generator = torch.Generator("cuda").manual_seed(int(seed) + 1)

    start = time.time()
    chunks = list(_stream_windows(text_ids, max_frames, ar_generator, dit_generator, steps, guidance))
    full = _to_int16(torch.cat(chunks, dim=-1))
    audio_seconds = full.shape[0] / SAMPLE_RATE
    return full, SAMPLE_RATE, audio_seconds, time.time() - start