import spaces  # imported first, as ZeroGPU requires

import sys
sys.stdout.reconfigure(line_buffering=True)

import tempfile
from pathlib import Path

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 pyharp.tags import Channels, SampleRate, Subcategory

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

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

model_card = ModelCard(
    name="Instruct-MusicGen",
    description=(
        "Edits a music recording to follow a text instruction, e.g. adding, "
        "removing, or isolating an instrument. The model is fixed-rate and "
        "mono, so the result is resampled back to the input's rate and copied "
        "across its channels: the output has the format it came in with, but "
        "no stereo image. The code is Apache-2.0, but the "
        "model loads Meta's MusicGen weights, which are CC-BY-NC-4.0, so "
        "results cannot be used commercially. Repository: "
        "https://github.com/ldzhangyx/instruct-MusicGen Paper: "
        "https://arxiv.org/abs/2405.18386"
    ),
    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=[Subcategory.EDITING, SampleRate(SAMPLE_RATE), Channels(1), 'instruction-tuning', 'domain:music'],
)

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

ckpt_path = None
ckpt_error = None

model = None


def download_checkpoint():
    """Fetch the finetuned adapter weights, 14.6 GB."""
    global ckpt_path, ckpt_error
    try:
        ckpt_path = hf_hub_download(repo_id="ldzhangyx/instruct-MusicGen", filename="finetuned.ckpt",
                                    revision="2c96cb10fc7abc69cd69a28d79397be7003080b9")
        print("Checkpoint downloaded.")
    except Exception as e:
        ckpt_error = str(e)
        print(f"Download error: {e}")


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ (rather than
# __main__) and starts with the server, so the checkpoint is fetched there alone, before
# process_fn can reserve a GPU. The server never runs process_fn, so it never needs it.
if __name__ != "__main__":
    download_checkpoint()


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

    Construction builds the backbone on CUDA, which ZeroGPU's emulation records only
    from the main thread, so this runs at module level in the worker (see below)
    rather than in a background thread.
    """
    global model
    # weights_only=False bc ckpt has non-tensor objects
    model = InstructMusicGenAdapterLitModule.load_from_checkpoint(
        ckpt_path, weights_only=False
    )
    model.eval()


# Built at module level in the worker, which runs process_fn, so that no request spends the
# caller's ZeroGPU allowance on loading. The server never runs process_fn, so it never holds the
# model.
if __name__ != "__main__" and ckpt_error is None:
    load_model()


def process_fn(input_audio_path: str, instruction: str) -> str:
    """Edit the input audio according to the instruction."""
    # Checked before the GPU is reserved, so a request that cannot run spends none of the
    # caller's ZeroGPU allowance
    if ckpt_error is not None:
        raise gr.Error(f"Checkpoint download failed: {ckpt_error}")
    if model is None:
        raise gr.Error("The model is still loading, please wait a moment and try again.")
    if not instruction or not instruction.strip():
        raise gr.Error("Instruction cannot be empty.")
    if not input_audio_path:
        raise gr.Error("Provide the music to edit.")
    try:
        duration = sf.info(input_audio_path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The audio file is empty.")
    if duration > MAX_INPUT_SECONDS:
        raise gr.Error(
            f"The music is {duration:.1f} s long. Trim it to {MAX_INPUT_SECONDS} seconds or less."
        )

    return edit(input_audio_path, instruction)


# The edit is generated as long as the input, in one pass the GPU holds whole, so the input is
# limited to the longest measured to fit
MAX_INPUT_SECONDS = 60


def gpu_seconds(input_audio_path: str, instruction: str) -> float:
    # Measured at about 24 s for 8 s, 48 s for 20 s and 74 s for 30 s, padded
    return 15 + 3 * sf.info(input_audio_path).duration


@spaces.GPU(duration=gpu_seconds)
@torch.inference_mode()
def edit(input_audio_path: str, instruction: str) -> str:
    """Adapted from generate_edited_audio in the original repo's src/inference.py."""
    # 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_channels = 1 if input_audio.ndim == 1 else input_audio.shape[1]
    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()
    )

    # The model is mono at 32 kHz, so the result is resampled back to the input's rate and copied
    # across its channels
    if sample_rate != SAMPLE_RATE:
        generated_audio = torchaudio.functional.resample(generated_audio, SAMPLE_RATE, sample_rate)
    generated_audio = generated_audio.unsqueeze(-1).repeat(1, input_channels).numpy()

    input_ext = Path(input_audio_path).suffix.lower()
    with tempfile.NamedTemporaryFile(suffix=input_ext if input_ext in LOSSLESS else ".wav",
                                     delete=False) as f:
        output_audio_path = f.name
    sf.write(output_audio_path, generated_audio, sample_rate)
    return output_audio_path


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "The music to edit, up to 60 seconds. The instruction decides what happens "
                "to it, so the arrangement should contain whatever you are asking to change."
            )),
            gr.Textbox(
                label="Instruction",
                value="Only Drums",
                info="What to change: add/only/no + an instrument, e.g. 'Only Drums', 'No Bass', "
                     "'Add Piano', as in the model's demo examples. Piano, bass, drums, and 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,
        )

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