import spaces
import gradio as gr
import torch
from pyharp import ModelCard, build_endpoint, LabelList, AudioLabel, OutputLabel
from skey.key_detection import detect_key

# Check for GPU availability
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Create a ModelCard
model_card = ModelCard(
    name="S-KEY Detection",
    description="Major/Minor Key detection from audio using Self-supervised Key Detection.",
    author="Deezer Research",
    tags=["key detection", "mir"],
)

# Define the process function
@spaces.GPU
def process_fn(input_audio_path):
    try:
        ext = input_audio_path.split(".")[-1].lower()

        detected_keys = detect_key(
            input_audio_path,
            extension=ext,
            device=DEVICE,
            cli=False
        )

        print("Raw output:", detected_keys)

        if not detected_keys:
            raise ValueError("No key detected")

        detected_key = detected_keys[0]  # because function returns list[str]

        labels = LabelList()
        labels.labels.append(
            AudioLabel(
                t=0.0,
                duration=5.0,
                label=detected_key,
                amplitude=0.8,
                description="Detected by S-KEY",
                color=OutputLabel.rgb_color_to_int(255, 100, 100, 0.8),
            )
        )

        return input_audio_path, labels

    except Exception as e:
        print("Error:", e)
        raise e


# Build Gradio Endpoint
with gr.Blocks() as demo:
    input_audio = gr.Audio(label="Input Audio", type="filepath").harp_required(True)
    # output_text = gr.Textbox(label="Detected Key")
    output_audio = gr.Audio(label="Output Audio", type="filepath")
    output_labels = gr.JSON(label="Analysis Labels")

    build_endpoint(
        model_card=model_card,
        process_fn=process_fn,
        input_components=[input_audio],
        output_components=[output_audio, output_labels],
    )


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