import spaces  # imported first, as ZeroGPU requires

import sys

sys.stdout.reconfigure(line_buffering=True)

from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint
from pyharp.labels import LabelList, AudioLabel
from pyharp.tags import Channels, SampleRate, Subcategory

import gradio as gr
import soundfile as sf
import torch

from beat_this.inference import File2Beats

model_card = ModelCard(
    name="Beat This!",
    description=(
        "Detects the beat and downbeat (start-of-bar) positions in a piece of "
        "music. The input is mixed to mono and resampled to 22.05 kHz internally. "
        "Repository: https://github.com/CPJKU/beat_this Paper: "
        "https://arxiv.org/abs/2407.21658"
    ),
    author="Francesco Foscarin, Jan Schlüter, Gerhard Widmer",
    tags=[Subcategory.MUSIC_ANALYSIS, SampleRate(22050), Channels(1)],
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Built and placed in the worker alone, which imports this file as __mp_main__ and runs process_fn,
# so a broken checkpoint fails in the logs when the Space starts and no request builds the model
# again. The checkpoint lives in a model repository rather than in this Space, pinned to a revision.
if __name__ != "__main__":
    model = File2Beats(
        checkpoint_path=hf_hub_download(
            repo_id="teamup-tech/Beat-This-weights",
            filename="final0.ckpt",
            revision="4afbe376535aaab71b76ae2ca2d7c16b08877321",
        ),
        device=DEVICE,
        dbn=False,
    )

BEAT_COLOR = AudioLabel.rgb_color_to_int(120, 170, 255)  # blue
DOWNBEAT_COLOR = AudioLabel.rgb_color_to_int(255, 130, 40)  # orange

# downbeats stay in the overhead row (no amplitude); beats get an amplitude
# near the top of the waveform, so they show up as a row just below downbeats
BEAT_AMPLITUDE = 0.9


def process_fn(input_audio_path: str) -> tuple[str, LabelList]:
    """Finds beat and downbeat times in the input audio and returns them as labeled markers on the audio."""
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not input_audio_path:
        raise gr.Error("Provide the music to track.")
    try:
        info = sf.info(input_audio_path)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if info.frames == 0:
        raise gr.Error("The audio file is empty.")

    return track_beats(input_audio_path)


@spaces.GPU
@torch.inference_mode()
def track_beats(input_audio_path: str) -> tuple[str, LabelList]:
    beats, downbeats = model(input_audio_path)
    # downbeats are a subset of beats with matching float values
    downbeat_times = set(downbeats.tolist())

    output_labels = LabelList()
    for t in beats:
        is_downbeat = float(t) in downbeat_times
        if is_downbeat:
            output_labels.append(
                AudioLabel(t=float(t), label="downbeat", color=DOWNBEAT_COLOR)
            )
        else:
            output_labels.append(
                AudioLabel(
                    t=float(t), label="beat", color=BEAT_COLOR, amplitude=BEAT_AMPLITUDE
                )
            )
    # HARP only attaches labels to an output track, so pass the audio through unchanged
    return input_audio_path, output_labels


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "Music to track. Beats and downbeats come back as labels over the "
                "audio."
            )),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Audio with Beats").set_info(
                "Input audio, unchanged, with detected beats and downbeats marked."
            ),
            gr.JSON(label="Beats").set_info(
                "Detected beat and downbeat times, labeled \"beat\" or \"downbeat\"."
            ),
        ]

        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    demo.queue().launch(show_error=True, pwa=True)