from __future__ import annotations

from pathlib import Path
from tempfile import TemporaryDirectory

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

from runtime import EXAMPLE, render

model_card = ModelCard(
    name="DExter",
    description="Render an expressive piano MIDI performance from a MusicXML score.",
    author="Huan Zhang and collaborators",
    tags=["piano", "performance-rendering", "musicxml", "midi-generation"],
)


@spaces.GPU(duration=75)
def process_fn(xml_path: str | None, use_example: bool, seed: int) -> str:
    if not xml_path and not use_example:
        raise gr.Error("Upload a MusicXML score or select the example.")
    source = xml_path or str(EXAMPLE)
    with TemporaryDirectory(prefix="dexter-") as tmp:
        output = Path(tmp) / "performance.mid"
        try:
            render(source, output, int(seed))
        except Exception as exc:
            raise gr.Error(f"DExter inference failed: {exc}") from exc
        return output_components[0].move_resource_to_block_cache(str(output))


with gr.Blocks(title="DExter", delete_cache=(3600, 86400)) as demo:
    inputs = [
        gr.File(type="filepath", file_types=[".xml", ".musicxml"], label="Piano MusicXML score").set_info(
            "Uncompressed MusicXML, up to 2 MB and 400 notes. Upload takes priority over the example."
        ),
        gr.Checkbox(value=True, label="Use bundled Schubert example if no file is uploaded"),
        gr.Number(value=13, precision=0, minimum=0, maximum=2147483647, label="Random seed"),
    ]
    output_components = [gr.File(label="Expressive MIDI performance", file_types=[".mid"])]
    build_endpoint(
        model_card=model_card,
        input_components=inputs,
        output_components=output_components,
        process_fn=process_fn,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=4).launch(
        show_error=True, pwa=True, max_file_size="2mb"
    )