import os
import subprocess
import tempfile
import soundfile as sf
import numpy as np
import gradio as gr
from pyharp import ModelCard, build_endpoint

synthesis_generator = None
expression_generator = None

INSTRUMENTS = [
    "violin", "viola", "cello", "double bass", "flute", "oboe",
    "clarinet", "saxophone", "bassoon", "trumpet", "horn", "trombone", "tuba"
]

def get_model():
    global synthesis_generator, expression_generator
    if synthesis_generator is None:
        print("Downloading MIDI-DDSP model weights...", flush=True)
        subprocess.run(["midi_ddsp_download_model_weights"], check=True)
        print("Loading MIDI-DDSP model...", flush=True)
        from midi_ddsp import load_pretrained_model
        synthesis_generator, expression_generator = load_pretrained_model()
        print("Model loaded.", flush=True)
    return synthesis_generator, expression_generator

model_card = ModelCard(
    name="MIDI-DDSP",
    description="Synthesize MIDI files into expressive audio using DDSP. Supports 13 orchestral instruments with realistic performance rendering.",
    author="Yusong Wu, Ethan Manilow, Yi Deng, et al. (Google Magenta)",
    tags=["midi", "synthesis", "ddsp", "performance-rendering", "orchestral"],
)

def process_fn(input_midi_path: str, instrument: str) -> str:
    print(f"Synthesizing {instrument}...", flush=True)
    syn_gen, exp_gen = get_model()

    from midi_ddsp.utils.midi_synthesis_utils import synthesize_mono_midi
    from midi_ddsp.data_handling.instrument_name_utils import INST_NAME_TO_ID_DICT

    instrument_id = INST_NAME_TO_ID_DICT[instrument]
    print(f"instrument_id: {instrument_id}", flush=True)

    output_dir = tempfile.mkdtemp()
    midi_audio, midi_control_params, midi_synth_params, conditioning_df = synthesize_mono_midi(
        syn_gen, exp_gen, input_midi_path, instrument_id, output_dir=None
    )

    audio = midi_audio.numpy()
    if audio.ndim > 1:
        audio = audio.squeeze()

    output_path = tempfile.mktemp(suffix=".wav")
    sf.write(output_path, audio, 16000)
    print("Done.", flush=True)
    return output_path

with gr.Blocks() as demo:
    input_components = [
        gr.File(
            type="filepath",
            label="Input MIDI File",
            file_types=[".mid", ".midi"],
        ).harp_required(True),
        gr.Dropdown(
            choices=INSTRUMENTS,
            value="violin",
            label="Instrument",
        ),
    ]
    output_components = [
        gr.Audio(
            type="filepath",
            label="Synthesized Audio",
        ).set_info("Expressive audio rendered from MIDI using DDSP."),
    ]
    app = build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

print("Launching Gradio...", flush=True)
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)