import sys

sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    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)
                func = args[0]
                return func


from pyharp import ModelCard, build_endpoint
from pyharp.labels import LabelList, AudioLabel

import gradio as gr
import torch

from beat_this.inference import File2Beats

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

# built on CPU at import time so a broken checkpoint fails fast in the logs;
# moved to DEVICE lazily on first request, since ZeroGPU only allows CUDA
# calls made inside an @spaces.GPU-decorated call.
model = File2Beats(checkpoint_path="final0.ckpt", device="cpu", dbn=False)
model_ready = DEVICE == "cpu"

model_card = ModelCard(
    name="Beat This!",
    description="Detects the beat and downbeat (start-of-bar) positions in a piece of music.",
    author="Francesco Foscarin, Jan Schlüter",
    tags=["beat tracking", "rhythm"],
)

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


@spaces.GPU
@torch.inference_mode()
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."""
    global model, model_ready
    if not model_ready:
        model = File2Beats(checkpoint_path="final0.ckpt", device=DEVICE, dbn=False)
        model_ready = True

    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


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Output Audio").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,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)