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 sys
sys.stdout.reconfigure(line_buffering=True)

import tempfile
import threading

import gradio as gr
import soundfile as sf
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from src.constants import CODEBOOK_SIZE, SAMPLE_RATE
from src.models.instructmusicgenadapter_module import InstructMusicGenAdapterLitModule

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

ckpt_path = None
ckpt_ready = False
ckpt_error = None

model = None
model_ready = False
model_lock = threading.Lock()


def download_checkpoint():
    """Fetch the finetuned adapter weights in the background so the server can
    start immediately instead of blocking on a 14.6GB download."""
    global ckpt_path, ckpt_ready, ckpt_error
    try:
        ckpt_path = hf_hub_download(repo_id="ldzhangyx/instruct-MusicGen", filename="finetuned.ckpt")
        print("Checkpoint downloaded.")
    except Exception as e:
        ckpt_error = str(e)
        print(f"Download error: {e}")
    finally:
        ckpt_ready = True


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


model_card = ModelCard(
    name="Instruct-MusicGen",
    description="Edits a music recording to follow a text instruction, e.g. adding, "
                 "removing, or isolating an instrument.",
    author="Yixiao Zhang, Yukara Ikemiya, Woosung Choi, Naoki Murata, "
           "Marco A. Martinez-Ramirez, Liwei Lin, Gus Xia, Wei-Hsiang Liao, "
           "Yuki Mitsufuji, Simon Dixon",
    tags=["music editing", "instruction tuning"],
)


def load_model():
    """Build the model and load the finetuned weights.

    Must run inside @spaces.GPU. Construction itself builds the backbone on CUDA,
    so it can't happen in a background thread on ZeroGPU.
    """
    global model
    # weights_only=False bc ckpt has non-tensor objects
    model = InstructMusicGenAdapterLitModule.load_from_checkpoint(
        ckpt_path, weights_only=False
    )
    model.eval()


@spaces.GPU
@torch.inference_mode()
def process_fn(input_audio_path: str, instruction: str) -> str:
    """Edit the input audio according to the instruction.

    Adapted from generate_edited_audio in the original repo's src/inference.py.
    """
    global model, model_ready

    if not ckpt_ready:
        raise gr.Error("Checkpoint is still downloading, please wait a moment and try again.")
    if ckpt_error is not None:
        raise gr.Error(f"Checkpoint download failed: {ckpt_error}")

    with model_lock:
        if not model_ready:
            load_model()
            model_ready = True

    if not instruction or not instruction.strip():
        raise gr.Error("Instruction cannot be empty.")

    # model was trained on "Music piece. Instruct: <desired_edit>." 
    # only <desired_edit> is variable, so rest is hardcoded
    instruction = f"Music piece. Instruct: {instruction.strip().rstrip('.')}."

    input_audio, sample_rate = sf.read(input_audio_path)
    input_audio_tensor = torch.tensor(input_audio).float()
    if input_audio_tensor.ndim == 2:
        # downmix to mono -- the original inference.py assumes a mono waveform!!
        input_audio_tensor = input_audio_tensor.mean(dim=-1)
    input_audio_tensor = input_audio_tensor.unsqueeze(0).unsqueeze(0)
    if sample_rate != SAMPLE_RATE:
        # the original inference.py never resamples, but the compression model
        # expects SAMPLE_RATE (32kHz); a DAW export at 44.1/48kHz needs converting first
        input_audio_tensor = torchaudio.functional.resample(input_audio_tensor, sample_rate, SAMPLE_RATE)
    input_audio_tensor = input_audio_tensor.to(DEVICE)

    instruction_list = [instruction]

    with torch.autocast("cuda", dtype=torch.float16):
        description, cond_code = model.model.musicgen._prepare_tokens_and_attributes(
            instruction_list, input_audio_tensor
        )

    cond_code = torch.cat(
        [cond_code, torch.ones_like(cond_code[:, :, 0:1]) * CODEBOOK_SIZE], dim=-1
    )

    with torch.autocast("cuda", dtype=torch.float16):
        audio_values = model.model.generate(
            text_description=instruction_list,
            condition_audio_code=cond_code,
            num_samples=1,
        )

    generated_audio = (
        model.model.musicgen.compression_model.decode(audio_values, None)
        .squeeze()
        .float()  # decode runs under autocast, so this is still float16, but soundfile needs float32
        .cpu()
        .detach()
        .numpy()
    )

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        output_audio_path = f.name
    sf.write(output_audio_path, generated_audio, SAMPLE_RATE)
    return output_audio_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        gr.Textbox(
            label="Instruction",
            value="Only Drums",
            info="What to change: add/only/no + an instrument, e.g. 'Only Drums', 'No Bass', "
                 "'Add Piano' (per the model's demo examples; piano/bass/drums/guitar work best)",
        ),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Edited Audio").set_info(
            "Edited version of the input audio reflecting the instruction."
        ),
    ]

    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)