import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import sys
from pathlib import Path

import spaces  # MUST come before any CUDA-touching import
import torch

sys.path.insert(0, str(Path(__file__).parent / "src"))

import numpy as np
import soundfile as sf
import torchaudio.functional as AF
import gradio as gr

from minimax_music3_latent_refiner import MiniMaxMusic3RefinerPipeline

MODEL_ID = "terminusresearch/minimax-music3-latent-refiner-v0.10"
MAX_SECONDS = 120.0  # keep ZeroGPU duration honest; model is quadratic in frames

print("Loading MiniMax Music 3 latent refiner pipeline (refiner + DAV + MERT + CLAP)...")
# Load on CPU: from_pretrained's safetensors/torch.load paths touch real CUDA ops that the
# ZeroGPU main process cannot satisfy. The .to("cuda") moves below are intercepted by the
# spaces hijack, which packs tensors for streaming into the GPU worker.
REFINER = MiniMaxMusic3RefinerPipeline.from_pretrained(MODEL_ID, device="cpu")
REFINER.refiner.to("cuda")
REFINER.audio_vae.to("cuda")
REFINER.mert.to("cuda")
REFINER.clap.to("cuda")
for _v in REFINER.normalization.values():
    _v.to("cuda")
print("Pipeline loaded.")


def _load_waveform(path: str):
    """Read any audio file as float32 [channels, samples] at its native rate."""
    if not path:
        raise gr.Error("Please upload an audio file or pick an example first.")
    data, sr = sf.read(path, dtype="float32", always_2d=True)
    if data.shape[1] > 2:
        data = data[:, :2]
    return torch.from_numpy(data.T.copy()), sr


def _trim(waveform: torch.Tensor, sr: int, max_seconds: float):
    limit = int(max_seconds * sr)
    if waveform.shape[-1] > limit:
        return waveform[..., :limit], True
    return waveform, False


@spaces.GPU(duration=120)
def refine(
    audio_path: str,
    steps: int = 32,
    window_seconds: float = 30.0,
    overlap_seconds: float = 2.0,
    direct: bool = False,
):
    """Restore a damaged music clip with the MiniMax Music 3 latent refiner.

    Args:
        audio_path: degraded/damaged music audio to restore.
        steps: deterministic Euler bridge sampling steps (release default 32).
        window_seconds: overlapping inference window in seconds (30 s matches training).
        overlap_seconds: cross-fade overlap between windows in latent frames.
        direct: process the whole clip as one dense sequence instead of windows.

    Returns:
        Tuple of (restored audio at 44.1 kHz stereo, status text).
    """
    if not audio_path:
        raise gr.Error("Please upload an audio file or pick an example first.")
    waveform, sr = _load_waveform(audio_path)
    waveform, trimmed = _trim(waveform, sr, MAX_SECONDS)
    steps = int(steps)
    if steps < 1:
        steps = 32
    ws = None if direct else float(window_seconds)
    result = REFINER(
        waveform,
        sr,
        steps=steps,
        window_seconds=ws,
        overlap_seconds=float(overlap_seconds),
    )
    out = result.audio.squeeze(0).T.numpy()
    out_path = Path("/tmp") / "refined.wav"
    sf.write(out_path, out, result.sample_rate, subtype="PCM_16")
    status = (
        f"Restored {out.shape[0] / result.sample_rate:.1f}s of audio "
        f"({steps} bridge steps, {'one dense sequence' if ws is None else f'{ws:.0f}s windows'})."
    )
    if trimmed:
        status += f" Input was trimmed to the {MAX_SECONDS:.0f}s demo limit."
    return str(out_path), status


def degrade(audio_path: str, lowpass_hz: float, noise_db: float, bits: float):
    """Apply a degradation chain (bandwidth cut + noise + bit depth + soft clip)
    to a clean clip, so you can hear what the refiner restores.

    Args:
        audio_path: clean music audio to damage.
        lowpass_hz: lowpass cutoff in Hz (bandwidth reduction).
        noise_db: additive white noise level in dBFS.
        bits: bit-depth reduction (bits per sample).

    Returns:
        Tuple of (damaged audio, status text).
    """
    if not audio_path:
        raise gr.Error("Upload clean audio to damage first.")
    import scipy.signal as sps

    data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
    if data.shape[1] > 2:
        data = data[:, :2]
    limit = int(MAX_SECONDS * sr)
    data = data[:limit]
    y = data.copy()
    nyq = sr / 2
    if lowpass_hz < nyq:
        y = sps.sosfilt(sps.butter(2, lowpass_hz / nyq, btype="low", output="sos"), y, axis=0)
    y = sps.sosfilt(sps.butter(2, 60.0 / nyq, btype="high", output="sos"), y, axis=0)
    if noise_db > -90:
        rng = np.random.default_rng(0)
        y = y + rng.normal(0.0, 1.0, y.shape).astype(np.float32) * (10 ** (noise_db / 20.0))
    if bits >= 2:
        levels = 2 ** int(bits)
        y = np.round(y * levels) / levels
    peak = max(float(np.abs(y).max()), 1e-6)
    y = np.tanh(y * (0.95 / peak) * 2.0) * 0.9
    out_path = str(Path("/tmp") / "damaged.wav")
    sf.write(out_path, np.clip(y, -1, 1), sr, subtype="PCM_16")
    return out_path, f"Damaged {y.shape[0] / sr:.1f}s — now press Restore to refine it."


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # 🎵 MiniMax Music 3 — Latent Refiner v0.10

        Restore damaged music while keeping the performance, timing, and arrangement intact.
        A 137M-parameter bridge transformer that works directly in MiniMax Music 3's continuous
        DAV latent space, conditioned on MERT frame features and CLAP audio embeddings.

        **Upload a degraded clip** (bandwidth-limited, noisy, quantized, clipped recordings), or
        use the *Damage* tab to apply the training-style degradation chain to your own clean audio,
        then press **Restore**.

        Model: [`terminusresearch/minimax-music3-latent-refiner-v0.10`](https://huggingface.co/terminusresearch/minimax-music3-latent-refiner-v0.10)
        """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Tab("Restore"):
            with gr.Row():
                with gr.Column():
                    audio_in = gr.Audio(
                        label="Damaged music (input)",
                        type="filepath",
                        sources=["upload", "microphone"],
                    )
                    restore_btn = gr.Button("Restore", variant="primary")
                    status = gr.Textbox(label="Status", interactive=False, lines=2)
                with gr.Column():
                    audio_out = gr.Audio(label="Restored music (output)", type="filepath")
            with gr.Accordion("Advanced settings", open=False):
                steps = gr.Slider(1, 64, value=32, step=1, label="Bridge steps")
                window_seconds = gr.Slider(
                    5.0, 60.0, value=30.0, step=5.0,
                    label="Window seconds (30 s matches training)",
                )
                overlap_seconds = gr.Slider(
                    0.0, 10.0, value=2.0, step=0.5, label="Window overlap (seconds)"
                )
                direct = gr.Checkbox(
                    value=False,
                    label="One dense sequence (no windows; not the quality baseline)",
                )
        with gr.Tab("Damage your own audio"):
            gr.Markdown(
                "Apply the refiner's training-style degradation chain (bandwidth reduction, additive "
                "noise, bit-depth reduction, soft clipping) to a clean clip, then restore it in the "
                "**Restore** tab to compare. Public-domain (FreePD/CC0) source clips."
            )
            clean_in = gr.Audio(label="Clean music (input)", type="filepath", sources=["upload", "microphone"])
            with gr.Row():
                lowpass_hz = gr.Slider(500, 20000, value=3200, step=100, label="Lowpass cutoff (Hz)")
                noise_db = gr.Slider(-90, -10, value=-34, step=1, label="Noise level (dBFS)")
                bits = gr.Slider(2, 16, value=6, step=1, label="Bit depth")
            damage_btn = gr.Button("Damage", variant="secondary")
            damaged_out = gr.Audio(label="Damaged music (output)", type="filepath")
            damage_status = gr.Textbox(label="Status", interactive=False, lines=1)

        gr.Examples(
            examples=[
                ["piano_strings.wav"],
                ["epic_orchestral.wav"],
            ],
            inputs=[audio_in],
            outputs=[audio_out, status],
            fn=refine,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (public-domain music with the training-style degradation applied)",
        )
        gr.Markdown(
            """
            #### Reference clean sources
            These are the undamaged originals of the two examples — listen side by side to judge the restoration.
            """
        )
        with gr.Row():
            gr.Audio(value="reference_piano_strings_clean.wav", label="Clean original — piano & strings", type="filepath")
            gr.Audio(value="reference_epic_orchestral_clean.wav", label="Clean original — epic orchestral", type="filepath")

    restore_btn.click(
        refine,
        inputs=[audio_in, steps, window_seconds, overlap_seconds, direct],
        outputs=[audio_out, status],
        api_name="restore",
    )
    damage_btn.click(
        degrade,
        inputs=[clean_in, lowpass_hz, noise_db, bits],
        outputs=[damaged_out, damage_status],
        api_name="damage",
    )

demo.launch(mcp_server=True)