import sys

sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func

import tempfile
from pathlib import Path

import gradio as gr
import torch
from pyharp import ModelCard, build_endpoint

from muscriptor import TranscriptionModel
from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES, resolve_instrument_names

model_card = ModelCard(
    name="MuScriptor",
    description=(
        "Multi-instrument audio-to-MIDI transcription, trained on 170k songs "
        "from classical music to heavy metal."
    ),
    author="Simon Rouard, Michael Krause, Axel Roebel, Carl-Johann Simon-Gabriel, Alexandre Défossez (Kyutai x Mirelo)",
    tags=["transcription", "midi", "multi-instrument"],
)

_VALID_INSTRUMENTS = ", ".join(MT3_FULL_PLUS_GROUP_NAMES)

_models: dict[str, TranscriptionModel] = {}


def _get_model(variant: str) -> TranscriptionModel:
    """Load and cache a TranscriptionModel for a given size, one per variant.

    small/medium run on CPU; large runs on GPU. Built directly on its target
    device inside this GPU-decorated call, per ZeroGPU rules.
    """
    if variant not in _models:
        device = "cuda" if variant == "large" else "cpu"
        _models[variant] = TranscriptionModel.load_model(variant, device=device)
    return _models[variant]


def _resolve_instruments(text: str) -> tuple[list[str] | None, str]:
    """Resolve the comma-separated instruments box into exact group names.

    Names that don't resolve are dropped rather than raised: HARP's client
    shows generic error message, so an error here would not reach the user. 
    Second value returns what happened as a .txt output
    """
    tokens = [t for t in text.split(",") if t.strip()]
    if not tokens:
        return None, "No instrument restriction requested."

    resolved: list[str] = []
    problems: list[str] = []
    for token in tokens:
        try:
            resolved.extend(resolve_instrument_names([token]))
        except ValueError as e:
            problems.append(str(e))

    if not problems:
        return resolved, f"Matched: {', '.join(resolved)}."

    note = "; ".join(problems)
    if resolved:
        note = f"Matched: {', '.join(resolved)}. Ignored: {note}"
    else:
        note = f"No valid instrument names found, transcribing unrestricted. {note}"
    return (resolved or None), note


@spaces.GPU
@torch.inference_mode()
def process_fn(
    input_audio_path: str,
    variant: str,
    instruments_text: str,
    use_sampling: bool,
    temperature: float,
) -> tuple[str, str]:
    """Transcribe the input audio to MIDI, plus a note on the Instruments field."""
    model = _get_model(variant)
    instruments, instrument_note = _resolve_instruments(instruments_text)

    midi_bytes = model.transcribe_to_midi(
        input_audio_path,
        use_sampling=use_sampling,
        temperature=temperature,
        instruments=instruments,
    )

    with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as f:
        f.write(midi_bytes)
        output_midi_path = f.name

    notes_text = f"{Path(input_audio_path).name}\n\nInstrument Matching\n{instrument_note}\n"
    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        f.write(notes_text)
        output_notes_path = f.name

    return output_midi_path, output_notes_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        gr.Dropdown(
            choices=["small", "medium", "large"],
            value="medium",
            label="Model Size",
            info="small = fastest, least accurate. medium = balanced (default, per repo). large = most accurate, slower.",
        ),
        gr.Textbox(
            value="",
            label="Instruments",
            info=f"Separate names with commas. Leave blank to let the model detect instruments on its own. Valid names: {_VALID_INSTRUMENTS}.",
        ),
        gr.Checkbox(
            label="Use Sampling",
            info="Temperature sampling instead of greedy decoding (default: False, per repo).",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            step=0.1,
            value=1.0,
            label="Temperature",
            info="Only used when Use Sampling is on. Higher values = note choices are more varied and unpredictable. Lower values = values are closer to model's most confident guesses. (default: 1.0, per repo).",
        ),
    ]
    output_components = [
        gr.File(type="filepath", file_types=[".mid", ".midi"], label="Output MIDI").set_info(
            "Transcribed MIDI notes."
        ),
        gr.File(type="filepath", file_types=[".txt"], label="Instrument Matching").set_info(
            "Which requested instrument names were matched or ignored."
        ),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)