import os
import tempfile
import shutil
from pathlib import Path

os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import gradio as gr
from pyharp import ModelCard, build_endpoint

MODES = ["music", "drum", "vocal", "vocal-contour", "chord", "beat"]

model_card = ModelCard(
    name="Omnizart",
    description="Automatic music transcription: given audio, produce MIDI or transcription for pitched instruments, drums, vocals, chords, and beat.",
    author="Yu-Te Wu et al. (MCT Lab)",
    tags=["transcription", "midi", "amt", "chord", "beat", "drum"],
)

_apps = {}

def get_omnizart_app(mode: str):
    if mode not in _apps:
        print(f"Loading omnizart {mode} module...", flush=True)
        if mode == "music":
            from omnizart.music import app as m; _apps[mode] = m
        elif mode == "drum":
            from omnizart.drum import app as m; _apps[mode] = m
        elif mode == "vocal":
            from omnizart.vocal import app as m; _apps[mode] = m
        elif mode == "vocal-contour":
            from omnizart.vocal_contour import app as m; _apps[mode] = m
        elif mode == "chord":
            from omnizart.chord import app as m; _apps[mode] = m
        elif mode == "beat":
            from omnizart.beat import app as m; _apps[mode] = m
        print(f"Loaded {mode}.", flush=True)
    return _apps[mode]

def process_fn(input_audio_path: str, mode: str) -> str:
    print(f"Transcribing with mode: {mode}...", flush=True)
    output_dir = Path(tempfile.mkdtemp())

    app = get_omnizart_app(mode)
    result = app.transcribe(input_audio_path, output=str(output_dir))

    print(f"Transcription result type: {type(result)}", flush=True)
    print(f"Transcription result: {result}", flush=True)

    out_path = tempfile.mktemp(suffix=".mid")

    # Handle both PrettyMIDI object and file path returns
    if hasattr(result, 'write'):
        # It's a PrettyMIDI object — write it directly
        result.write(out_path)
        print(f"Wrote PrettyMIDI to {out_path}", flush=True)
    elif isinstance(result, (str, Path)) and Path(result).exists():
        shutil.copy(result, out_path)
        print(f"Copied file to {out_path}", flush=True)
    else:
        # Search output dir for any generated files
        candidates = list(output_dir.glob("**/*.*"))
        print(f"Output dir contents: {candidates}", flush=True)
        if not candidates:
            raise RuntimeError(f"No output produced. Result was: {result}")
        shutil.copy(candidates[0], out_path)

    print("Done.", flush=True)
    return out_path

with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio"),
        gr.Dropdown(choices=MODES, value="music", label="Transcription Mode"),
    ]
    output_components = [
        gr.File(
            label="Transcription Output (MIDI or text)",
            file_types=[".mid", ".midi", ".txt", ".csv"],
        ).set_info("MIDI file for music/drum/vocal modes; text file for chord/beat modes."),
    ]
    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)