from __future__ import annotations

import uuid
from pathlib import Path
from tempfile import gettempdir

import gradio as gr

try:
    import spaces
except ImportError:
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=120):
                self.func = func

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

from pyharp import ModelCard, build_endpoint

from midi_gpt_runtime import rewrite_bars

OUTPUT_ROOT = Path(gettempdir()) / "midi_gpt_outputs"

model_card = ModelCard(
    name="MIDI-GPT",
    description=(
        "Rewrite selected bars in a multitrack MIDI arrangement while "
        "preserving the rest of the composition."
    ),
    author="Metacreation Lab",
    tags=[
        "music-generation",
        "symbolic-music",
        "midi",
        "music-editing",
        "multitrack",
    ],
)


@spaces.GPU(duration=120)
def process_fn(
    midi_path: str | None,
    track_index: float,
    start_bar: float,
    bar_count: str,
    temperature: float,
    top_p: float,
    seed: float,
    polyphony_limit: float,
) -> tuple[str, dict]:
    if not midi_path:
        raise gr.Error("Please upload a MIDI file.")

    source = Path(midi_path)
    if source.suffix.lower() not in {".mid", ".midi"}:
        raise gr.Error("The input must be a .mid or .midi file.")

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_path = output_dir / f"{source.stem}-rewritten.mid"

    try:
        details = rewrite_bars(
            input_path=source,
            output_path=output_path,
            track_index=int(track_index),
            start_bar=int(start_bar),
            bar_count=int(bar_count),
            temperature=float(temperature),
            top_p=float(top_p),
            seed=int(seed),
            polyphony_limit=int(polyphony_limit),
        )
    except Exception as exc:
        raise gr.Error(f"MIDI-GPT inference failed: {exc}") from exc

    return str(output_path), details


with gr.Blocks(title="MIDI-GPT") as demo:
    input_components = [
        gr.File(
            type="filepath",
            file_types=[".mid", ".midi"],
            label="Input MIDI",
        )
        .harp_required(True)
        .set_info("A multitrack MIDI file containing at least four bars."),
        gr.Number(
            value=0,
            minimum=0,
            precision=0,
            label="Track Index",
            info="Zero-based index of the track to rewrite.",
        ),
        gr.Number(
            value=0,
            minimum=0,
            precision=0,
            label="Start Bar",
            info="Zero-based index of the first bar to rewrite.",
        ),
        gr.Dropdown(
            choices=["1", "2", "4"],
            value="2",
            label="Number of Bars",
            info="Number of consecutive bars to rewrite.",
        ),
        gr.Slider(
            minimum=0.5,
            maximum=1.5,
            step=0.1,
            value=1.0,
            label="Temperature",
            info="Higher values produce more varied results.",
        ),
        gr.Slider(
            minimum=0.5,
            maximum=1.0,
            step=0.05,
            value=0.95,
            label="Top-p",
            info="Nucleus sampling threshold.",
        ),
        gr.Number(
            value=0,
            minimum=0,
            maximum=2_147_483_647,
            precision=0,
            label="Seed",
            info="Use the same seed and controls to reproduce a result.",
        ),
        gr.Slider(
            minimum=0,
            maximum=12,
            step=1,
            value=0,
            label="Maximum Polyphony",
            info="Maximum simultaneous notes; 0 leaves it unrestricted.",
        ),
    ]
    output_components = [
        gr.File(
            type="filepath",
            file_types=[".mid", ".midi"],
            label="Rewritten MIDI",
        ).set_info("Editable multitrack MIDI with the selected bars replaced."),
        gr.JSON(label="Generation Details"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True,
        pwa=True,
    )