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

import gradio as gr
import torch
import librosa
import numpy as np
import soundfile as sf

# adapted from AI-Music-Detection-HFSpace/app.py: torchaudio 2.9+ routes
# torchaudio.load through load_with_torchcodec, which isn't installed here.
# Loading via librosa/soundfile instead sidesteps that dependency entirely.
import torchaudio


def _patched_load(filepath, *args, **kwargs):
    frame_offset = kwargs.pop("frame_offset", 0)
    num_frames = kwargs.pop("num_frames", -1)
    channels_first = kwargs.pop("channels_first", True)
    if len(args) >= 1:
        frame_offset = args[0]
    if len(args) >= 2:
        num_frames = args[1]

    data, sample_rate = librosa.load(str(filepath), sr=None, mono=False)
    if data.ndim == 1:
        data = data[np.newaxis, :]
    if frame_offset and int(frame_offset) > 0:
        data = data[:, int(frame_offset):]
    if num_frames and int(num_frames) > 0:
        data = data[:, : int(num_frames)]

    duration_s = data.shape[-1] / max(sample_rate, 1)
    if duration_s < 1.5:
        raise RuntimeError(
            f"Audio is too short ({duration_s:.2f} s). "
            "Please upload a clip of at least 2 seconds (a few seconds of music works best)."
        )
    waveform = torch.from_numpy(np.ascontiguousarray(data)).float()
    if not channels_first:
        waveform = waveform.transpose(0, 1).contiguous()
    return waveform, sample_rate


torchaudio.load = _patched_load


class _AudioInfo:
    __slots__ = ("sample_rate", "num_frames", "num_channels", "bits_per_sample", "encoding")


def _patched_info(filepath, *args, **kwargs):
    info = sf.info(str(filepath))
    out = _AudioInfo()
    out.sample_rate = info.samplerate
    out.num_frames = info.frames
    out.num_channels = info.channels
    out.bits_per_sample = 0
    out.encoding = info.format
    return out


torchaudio.info = _patched_info

import tempfile
import threading
from pathlib import Path

from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from inference import inference

CHECKPOINT_REPO = "teamup-tech/FST-AI-Music-Detection-checkpoints"
CHECKPOINT_DIR = Path("checkpoints")

checkpoints_loading = True
checkpoint_error = None


def download_checkpoints():
    """Fetch the Stage-1 and Stage-2 checkpoints so inference.py finds them
    already on disk instead of downloading ~1.3 GB on the first request."""
    global checkpoints_loading, checkpoint_error
    try:
        CHECKPOINT_DIR.mkdir(exist_ok=True)
        for filename in ("Stage-1.ckpt", "Stage-2.ckpt"):
            hf_hub_download(repo_id=CHECKPOINT_REPO, filename=filename, local_dir=str(CHECKPOINT_DIR))
        print("Checkpoints ready.")
    except Exception as e:
        checkpoint_error = str(e)
        print(f"Checkpoint download error: {e}")
    finally:
        checkpoints_loading = False


threading.Thread(target=download_checkpoints, daemon=True).start()


model_card = ModelCard(
    name="Fusion Segment Transformer",
    description=(
        "Detects AI-generated music by fusing MERT audio embeddings with a "
        "structural, downbeat-segmented view of the track. Submitted to ICASSP 2026."
    ),
    author="Yumin Kim, Seonghyeon Go",
    tags=["music", "classification", "ai-detection"],
)


@spaces.GPU
@torch.inference_mode()
def process_fn(audio_path: str) -> str:
    """Runs the two-stage detector on an uploaded track and writes the verdict to a text file."""
    if checkpoints_loading:
        raise gr.Error("Model checkpoints are still downloading, please wait a moment and try again.")
    if checkpoint_error:
        raise gr.Error(f"Failed to download checkpoints: {checkpoint_error}")

    result = inference(audio_path)
    prediction = result.get("prediction", "Unknown")
    confidence = result.get("confidence", "0.00")
    input_name = Path(audio_path).name

    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        f.write(f"{input_name}\n\nAI-Detection Results\nPrediction: {prediction} ({confidence}% confidence)\n")
        out_path = f.name

    return out_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
    ]
    output_components = [
        gr.File(type="filepath", label="Detection Result", file_types=[".txt"]).set_info(
            "Prediction (Real/Fake) and confidence, as a text file."
        ),
    ]

    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)