from __future__ import annotations

import json
import tempfile
import uuid
from pathlib import Path

import gradio as gr
import soundfile as sf

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

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

from merit_runtime import compare_audio
from pyharp import ModelCard, build_endpoint


MIN_AUDIO_SECONDS = 5
MAX_AUDIO_SECONDS = 10
OUTPUT_ROOT = Path(tempfile.gettempdir()) / "merit_outputs"

model_card = ModelCard(
    name="MERIT",
    description=(
        "Compare two music clips using independent melody, rhythm, "
        "and timbre similarity representations."
    ),
    author="AMAAI Lab",
    tags=[
        "music-information-retrieval",
        "music-similarity",
        "melody",
        "rhythm",
        "timbre",
    ],
)


def _validate_audio(path: str | None, label: str) -> str:
    if not path:
        raise gr.Error(f"Please upload {label}.")

    try:
        duration = sf.info(path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read {label}: {exc}") from exc

    if duration <= 0:
        raise gr.Error(f"{label} is empty.")
    if duration < MIN_AUDIO_SECONDS:
        raise gr.Error(
            f"{label} must be at least {MIN_AUDIO_SECONDS} seconds long. "
            f"Received {duration:.1f} seconds."
        )
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(
            f"{label} must be no longer than {MAX_AUDIO_SECONDS} seconds. "
            f"Received {duration:.1f} seconds."
        )
    return path


@spaces.GPU(duration=120)
def process_fn(
    reference_audio: str | None,
    comparison_audio: str | None,
) -> str:
    reference_audio = _validate_audio(reference_audio, "Reference Audio")
    comparison_audio = _validate_audio(comparison_audio, "Comparison Audio")

    try:
        scores = compare_audio(reference_audio, comparison_audio)
    except Exception as exc:
        raise gr.Error(f"MERIT inference failed: {exc}") from exc

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "merit_similarity.json"
    output_path.write_text(
        json.dumps(
            {
                "model": "MERIT",
                "score_type": "cosine_similarity",
                "score_range": [-1.0, 1.0],
                "scores": scores,
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    return str(output_path)


with gr.Blocks(title="MERIT Music Similarity") as demo:
    input_components = [
        gr.Audio(
            type="filepath",
            label="Reference Audio",
        )
        .harp_required(True)
        .set_info("First music clip, 5 to 10 seconds long."),
        gr.Audio(
            type="filepath",
            label="Comparison Audio",
        )
        .harp_required(True)
        .set_info("Second music clip, 5 to 10 seconds long."),
    ]
    output_components = [
        gr.File(
            type="filepath",
            file_types=[".json"],
            label="Similarity Results",
        ).set_info("Melody, rhythm, and timbre cosine similarity scores."),
    ]
    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,
    )