#!/usr/bin/env python3
"""Gradio/pyharp frontend for the Python 3.10 environment."""

from __future__ import annotations

import json
import os
import subprocess
from typing import Any

import gradio as gr
from pyharp import AudioLabel, LabelList, OutputLabel
from pyharp.core import ModelCard, build_endpoint

BACKEND_PYTHON = os.environ.get("BACKEND_PYTHON", "/opt/backend-py39/bin/python")
BACKEND_SCRIPT = os.environ.get("BACKEND_SCRIPT", "/app/backend_worker.py")


def call_backend(audio_path: str, timeout_s: float = 120.0) -> list[dict[str, float]]:
    completed = subprocess.run(
        [BACKEND_PYTHON, BACKEND_SCRIPT, audio_path],
        capture_output=True,
        check=False,
        text=True,
        timeout=timeout_s,
    )

    try:
        response: dict[str, Any] = json.loads(completed.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            "BeatNet backend returned invalid JSON. "
            f"stdout={completed.stdout!r} stderr={completed.stderr!r}"
        ) from exc

    if not response.get("ok"):
        error = response.get("error") or completed.stderr or "BeatNet backend failed"
        raise RuntimeError(error)
    return response["beats"]


def process_fn(inp_audio: str):
    beats = call_backend(inp_audio)

    output_labels = LabelList()
    for beat in beats:
        beat_value = beat["beat"]
        output_labels.labels.append(
            AudioLabel(
                t=beat["time"],
                label=f"{beat_value}",
                description=f"Beat: {beat_value}",
                color=OutputLabel.rgb_color_to_int(0, 164, 235),
                amplitude=1.0 if beat_value == 1.0 else 0.0,
            )
        )
    return inp_audio, output_labels


def build_app() -> gr.Blocks:
    model_card = ModelCard(
        name="BeatNet Beat Detection",
        description="Beat detection for audio.",
        author="Mojtaba Heydari, Frank Cwitkowitz, Zhiyao Duan",
        tags=["beat detection"],
    )

    with gr.Blocks() as app:
        gr.Markdown("## BeatNet Beat Detection")
        input_audio = gr.Audio(label="Input Audio", type="filepath", sources=["upload", "microphone"])
        output_wav = gr.Audio(type="filepath", label="Input Audio")
        output_label = gr.JSON(label="Beat Labels")

        build_endpoint(
            model_card=model_card,
            input_components=[input_audio],
            output_components=[output_wav, output_label],
            process_fn=process_fn,
        )

    return app


if __name__ == "__main__":
    build_app().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))