import spaces  # MUST come before any torch / CUDA-touching import
import torch
import gradio as gr
import tempfile
import os
import time
import sys
import io
import json
import re
import base64
from pathlib import Path

from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

from muscriptor.models.lm import LMModel, TorchAutocast
from muscriptor.modules.conditioners import (
    MelSpectrogramConditioner,
    ClassConditioner,
    ConditioningProvider,
)
from muscriptor.tokenizer.mt3 import MT3Tokenizer, MT3_FULL_PLUS_GROUP_NAMES
from muscriptor.transcription_model import (
    TranscriptionModel,
    _resolve_source,
    _resolve_config,
    _remap_single_codebook_keys,
    _build_model,
)

_SAMPLE_RATE = 16000


def load_model_zerogpu():
    """Load the MuScriptor model in a ZeroGPU-compatible way.

    On ZeroGPU, safetensors.load_file(device="cuda") fails because there's no
    real GPU at module scope. However, .to("cuda") is intercepted by the
    spaces hijack. So we:
      1. Build the model with device="cuda" (conditioners store this as
         self.device and use it at runtime to route tensors).
      2. Load safetensors weights on CPU.
      3. Load_state_dict into the model (weights land on CPU because the
         model's tensors are still fake-CUDA via the hijack).
      4. Call model.to("cuda") — ZeroGPU intercepts this and packs weights
         to disk, streaming them into VRAM on the first @spaces.GPU call.
    """
    device = torch.device("cuda")  # ZeroGPU intercepts this

    source = _resolve_source("large")
    weights_path = Path(hf_hub_download(
        repo_id="MuScriptor/muscriptor-large",
        filename="model.safetensors",
    ))
    cfg = _resolve_config(source, weights_path)

    # Build the model with device="cuda" — conditioners store this and use it
    # at runtime to route tensors. ZeroGPU intercepts .to("cuda") calls inside.
    model = _build_model(device, cfg)
    model.eval()

    # Load weights on CPU (safetensors can't load to fake CUDA)
    state_dict = load_file(str(weights_path), device="cpu")
    state_dict = _remap_single_codebook_keys(state_dict)

    # Load state dict — the model's parameters are fake-CUDA (ZeroGPU),
    # but load_state_dict copies CPU data into them, which is fine.
    model.load_state_dict(state_dict)

    # Move everything to "cuda" — ZeroGPU intercepts this and packs to disk
    model.to("cuda")

    tokenizer = MT3Tokenizer(
        instrument_vocabulary="MT3_FULL_PLUS",
        max_shift_steps=1001,
    )

    return TranscriptionModel(model=model, tokenizer=tokenizer, device=device)


print("[muscriptor-space] Loading model...", file=sys.stderr, flush=True)
t0 = time.perf_counter()
model = load_model_zerogpu()
print(f"[muscriptor-space] Model loaded in {time.perf_counter() - t0:.1f}s", file=sys.stderr, flush=True)

# Build the instrument choices list (sorted by group ID for stable ordering)
INSTRUMENT_CHOICES = sorted(MT3_FULL_PLUS_GROUP_NAMES.keys(), key=lambda k: MT3_FULL_PLUS_GROUP_NAMES[k])

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
midi-player { width: 100%; display: block; margin-bottom: 8px; }
midi-visualizer { width: 100%; display: block; overflow: auto; }
"""

# The html-midi-player web components (<midi-player> / <midi-visualizer>) rely
# on Magenta.js + Tone.js. Following Gradio's official custom-HTML-components
# guide (https://gradio.app/guides/custom-HTML-components), the library is
# loaded via the gr.HTML `head` parameter instead of embedding <script> tags in
# the component value — Gradio sanitizes value HTML and strips <script> tags,
# which is why the raw-HTML approach never mounted the web components.
_MIDI_PLAYER_HEAD = (
    '<script src="https://cdn.jsdelivr.net/combine/'
    "npm/tone@14.7.58,"
    "npm/@magenta/music@1.23.1/es6/core.js,"
    'npm/html-midi-player@1.5.0"></script>'
)

# Static markup for the player/visualizer. js_on_load wires the transcribed
# MIDI (a base64 data URI passed as the component value) into the <midi-player>
# and links it to the <midi-visualizer> so both mount and render.
_MIDI_PLAYER_TEMPLATE = (
    '<div class="midi-player-container">'
    '<midi-player sound-font visualizer=".midi-visualizer"></midi-player>'
    '<midi-visualizer class="midi-visualizer" type="piano-roll"></midi-visualizer>'
    "</div>"
)

# Runs once when the component loads. `element` is the component's DOM node,
# `props.value` is the base64 data URI (or "" when there's nothing yet).
# `watch('value', ...)` re-runs the sync every time the server pushes a new
# value (i.e. after each transcription), so the web components actually update.
_MIDI_PLAYER_JS = """
function syncPlayer() {
    const player = element.querySelector('midi-player');
    const container = element.querySelector('.midi-player-container');
    if (!player || !container) { return; }
    if (props.value) {
        container.style.display = 'block';
        player.src = props.value;
    } else {
        container.style.display = 'none';
        player.removeAttribute('src');
    }
}
syncPlayer();
watch('value', syncPlayer);
"""


def _midi_player_src(midi_bytes: bytes | None) -> str:
    """Return the transcribed MIDI as a base64 data URI for the player.

    The MIDI data is embedded directly as a base64 data URI so it needs no
    separate file-serving route. The <midi-player> / <midi-visualizer> web
    components consume this via the gr.HTML custom component's js_on_load.
    """
    if not midi_bytes:
        return ""
    b64 = base64.b64encode(midi_bytes).decode("ascii")
    return f"data:audio/midi;base64,{b64}"


@spaces.GPU(duration=120)
def transcribe_audio(
    audio_path: str,
    instruments: list[str] | None,
    use_sampling: bool,
    temperature: float,
    progress=gr.Progress(track_tqdm=True),
):
    """Transcribe an audio recording into a downloadable MIDI file.

    Upload any music recording (multi-instrument, any genre) and MuScriptor
    will convert it into a MIDI file with per-note onset, offset, pitch, and
    instrument information.

    Args:
        audio_path: Path to the uploaded audio file (wav, mp3, flac, etc.).
        instruments: Optional list of instrument group names to condition the
            transcription (improves coherence when you know which instruments
            are present). Leave empty for automatic (unconditioned) transcription.
        use_sampling: If True, use stochastic sampling instead of greedy decoding.
        temperature: Sampling temperature (only used when use_sampling is True).
    """
    if audio_path is None:
        return None, "Please upload an audio file first.", ""

    t0 = time.perf_counter()

    # Run the transcription — returns MIDI bytes
    try:
        midi_bytes = model.transcribe_to_midi(
            audio_path,
            instruments=instruments if instruments else None,
            use_sampling=use_sampling,
            temperature=temperature,
            sliding_hop=1.0
        )
    except Exception as e:
        return None, f"Transcription failed: {e}", ""

    elapsed = time.perf_counter() - t0

    # Write to a temporary file for download
    tmp = tempfile.NamedTemporaryFile(suffix=".mid", delete=False)
    tmp.write(midi_bytes)
    tmp.close()

    # Also collect note statistics for the summary
    note_count = 0
    instrument_set = set()
    try:
        from mido import MidiFile
        midi = MidiFile(tmp.name)
        for track in midi.tracks:
            for msg in track:
                if msg.type == "note_on" and msg.velocity > 0:
                    note_count += 1
                    instrument_set.add(msg.program if not msg.is_meta else 0)
    except Exception:
        pass

    summary = (
        f"Transcription complete in {elapsed:.1f}s. "
        f"Found {note_count} notes across {len(instrument_set)} instrument program(s). "
        f"Play it back below or download the MIDI file."
    )

    player_src = _midi_player_src(midi_bytes)

    return tmp.name, summary, player_src


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # MuScriptor — Music Transcription (Audio → MIDI)

            Upload a music recording and get a downloadable MIDI file with transcribed notes.
            MuScriptor is a ~1.3B parameter multi-instrument automatic music transcription model
            developed by [Mirelo](https://www.mirelo.ai/) x [Kyutai](https://kyutai.org/).

            [Model card](https://huggingface.co/MuScriptor/muscriptor-large) · [Code](https://github.com/muscriptor/muscriptor) · [Audio samples](https://muscriptor.github.io)
            """
        )

        with gr.Row():
            with gr.Column():
                audio_input = gr.Audio(
                    label="Upload or record audio",
                    type="filepath",
                    sources=["upload", "microphone"],
                )

                with gr.Accordion("Advanced settings", open=False):
                    instrument_checkbox = gr.CheckboxGroup(
                        choices=INSTRUMENT_CHOICES,
                        value=[],
                        label="Instrument conditioning (optional)",
                        info="Select instruments present in the audio to improve transcription accuracy. Leave empty for automatic detection.",
                    )
                    use_sampling = gr.Checkbox(
                        label="Use sampling (stochastic decoding)",
                        value=False,
                        info="If enabled, uses temperature-based sampling instead of greedy decoding.",
                    )
                    temperature = gr.Slider(
                        label="Temperature",
                        minimum=0.1,
                        maximum=2.0,
                        value=1.0,
                        step=0.1,
                        info="Sampling temperature (only used when sampling is enabled).",
                    )

                transcribe_btn = gr.Button("Transcribe", variant="primary")

            with gr.Column():
                midi_player = gr.HTML(
                    value="",
                    label="MIDI playback & visualization",
                    head=_MIDI_PLAYER_HEAD,
                    html_template=_MIDI_PLAYER_TEMPLATE,
                    js_on_load=_MIDI_PLAYER_JS,
                )
                midi_output = gr.File(label="Download MIDI file")

        summary_output = gr.Textbox(label="Summary", interactive=False, visible=False)

    transcribe_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input, instrument_checkbox, use_sampling, temperature],
        outputs=[midi_output, summary_output, midi_player],
        api_name="transcribe",
    )

    gr.Examples(
        examples=[
            ["example_medicine.mp3", [], False, 1.0],
            ["example_organic_flow.mp3", [], False, 1.0],
            ["example_water_afro_pop.mp3", [], False, 1.0],
        ],
        inputs=[audio_input, instrument_checkbox, use_sampling, temperature],
        outputs=[midi_output, summary_output, midi_player],
        fn=transcribe_audio,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)