import spaces  # imported first, as ZeroGPU requires

import os
import pickle
import torch
import torch.nn as nn
import gradio as gr
from transformers import T5Tokenizer
from huggingface_hub import hf_hub_download
from time import time_ns
from uuid import uuid4

from transformer_model import Transformer
from pyharp.core import ModelCard, build_endpoint
from pyharp.tags import Subcategory

# HARP Model Card
model_card = ModelCard(
    name="text2midi",
    description=(
        "Generates a MIDI sequence from a text description of the music. Repository: "
        "https://github.com/amaai-lab/text2midi Paper: https://arxiv.org/abs/2412.16526"
    ),
    author="Keshav Bhandari, Abhinaba Roy, Kyra Wang, Geeta Puri, Simon Colton, Dorien Herremans",
    tags=[Subcategory.HOLISTIC, 'domain:music']
)

# Model/artifacts from HF Hub
REPO_ID = "amaai-lab/text2midi"
REVISION = "9bc91479c4a9b08a9a28a7dcfb3f1a56ee2ded36"


# Helpers
def _unique_path(ext: str) -> str:
    """Create a unique file path in /tmp to avoid naming collisions."""
    return os.path.join("/tmp", f"t2m_{time_ns()}_{uuid4().hex[:8]}{ext}")


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = None
r_tokenizer = None
tokenizer = None


def load_model():
    """Builds the model on the CPU, then places it on DEVICE."""
    global model, r_tokenizer, tokenizer

    # Load REMI vocab/tokenizer (pickle dict used by the provided model)
    with open(hf_hub_download(repo_id=REPO_ID, filename="vocab_remi.pkl", revision=REVISION), "rb") as f:
        r_tokenizer = pickle.load(f)

    vocab_size = len(r_tokenizer)
    model = Transformer(
        vocab_size,    # vocab size
        768,           # d_model
        8,             # nhead
        2048,          # dim_feedforward
        18,            # nlayers
        1024,          # max_seq_len
        False,         # use_rotary
        8,             # rotary_dim
        device="cpu"   # device
    )
    model_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin", revision=REVISION)
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model.eval()
    model.to(DEVICE)

    tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base", revision="7bcac572ce56db69c1ea7c8af255c5d7c9672fc2")


# Built in the worker, which runs process_fn, on its main thread, where ZeroGPU's emulation records
# the CUDA placement, so that no request spends the caller's allowance on building it. The server
# never runs process_fn.
if __name__ != "__main__":
    load_model()


# Core Text -> MIDI
def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) -> str:
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)

    input_ids = nn.utils.rnn.pad_sequence(inputs.input_ids, batch_first=True, padding_value=0).to(DEVICE)
    attention_mask = nn.utils.rnn.pad_sequence(inputs.attention_mask, batch_first=True, padding_value=0).to(DEVICE)

    with torch.no_grad():
        output = model.generate(input_ids, attention_mask, max_len=max_len, temperature=temperature)

    output_list = output[0].tolist()
    generated_midi = r_tokenizer.decode(output_list)

    midi_path = _unique_path(".mid")
    generated_midi.dump_midi(midi_path)
    return midi_path

def gpu_seconds(prompt: str, temperature: float, max_length: int) -> float:
    # Measured at about 42 s for the longest, 1500 tokens, padded
    return 10 + 0.05 * int(max_length)


def process_fn(prompt: str, temperature: float, max_length: int):
    """
    Generate a MIDI from a text description.

    Args:
        prompt (str): What the music should be.
        temperature (float): How freely the model departs from its likeliest
            continuation.
        max_length (int): Ceiling on the number of tokens generated.

    Returns:
        midi_path (str): The generated MIDI.
    """

    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not (prompt or "").strip():
        raise gr.Error("Describe the music to generate.")
    return compose(prompt, temperature, max_length)


@spaces.GPU(duration=gpu_seconds)
def compose(prompt: str, temperature: float, max_length: int) -> str:
    try:
        return generate_midi(prompt, float(temperature), int(max_length))
    except Exception as error:
        raise gr.Error(str(error))

# Gradio + HARP UI
if __name__ == "__main__":
    with gr.Blocks() as demo:
        gr.Markdown("## 🎶 text2midi")

        # Inputs
        prompt_in = gr.Textbox(
            label="Describe Your Music",
            info="Type a short phrase like 'calm piano with flowing arpeggios' ",
            ).harp_required(True)
        temperature_in = gr.Slider(minimum=0.8, maximum=1.1, value=0.9, step=0.1, label="Creativity", info=(
            "Adjusts how much freedom the model takes while composing.\n"
            "Lower = safer and more predictable (structured), "
            "Higher = more varied and expressive."
        ), interactive=True)

        maxlen_in = gr.Slider(minimum=500, maximum=1500, step=100, value=500, label="Composition Length", info=(
            "Determines how long the generated piece is in musical tokens.\n"
            "Higher values produce longer phrases (roughly more measures of music).")
        )

        # Outputs
        midi_out = gr.File(label="Generated MIDI", file_types=[".mid", ".midi"], type="filepath").set_info((
            "The generated notes, as a MIDI file."
        ))

        # Build HARP endpoint
        _ = build_endpoint(
            model_card=model_card,
            input_components=[prompt_in, temperature_in, maxlen_in],
            output_components=[midi_out],
            process_fn=process_fn
        )


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