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 pyharp.tags import Subcategory

from runtime import EXAMPLE, render

model_card = ModelCard(
    name="DExter",
    description=(
        "Render an expressive piano MIDI performance from a MusicXML score. "
        "Repository: https://github.com/anusfoil/DExter Paper: "
        "https://arxiv.org/abs/2406.14850"
    ),
    author="Huan Zhang, Shreyan Chowdhury, Carlos Eduardo Cancino-Chacón, et al.",
    tags=[Subcategory.PERFORMANCE_RENDERING, '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))


if __name__ == "__main__":
    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",
            info=(
                "Runs on a bundled Schubert excerpt when you have not uploaded a "
                "score."
            )),
            gr.Number(value=13, precision=0, minimum=0, maximum=2147483647, label="Random seed",
            info=(
                "Changes the performance the model samples. The same seed and score "
                "give the same performance."
            )),
        ]
        output_components = [gr.File(label="Expressive MIDI performance", file_types=[".mid"]).set_info((
            "The score played with timing, dynamics and pedalling the model "
            "chose."
        ))]
        build_endpoint(
            model_card=model_card,
            input_components=inputs,
            output_components=output_components,
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1, max_size=4).launch(
        share=True, show_error=True, pwa=True, max_file_size="2mb"
    )