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 os
import shutil
import tempfile
import threading
import time
import urllib.request
from pathlib import Path

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

import picogen2
from picogen2.mirtoolkit.beat_this import BeatThis
from picogen2.mirtoolkit.sheetsage import SheetSage

REPO_ROOT = Path(__file__).parent


def _download_with_progress(url: str, dest: Path, label: str, interval_s: float = 20.0):
    """Downloads url to dest, logging progress at most once per interval_s."""
    if dest.exists():
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = dest.with_name(dest.name + ".part")
    with urllib.request.urlopen(url) as response, open(tmp_path, "wb") as f:
        total = int(response.headers.get("Content-Length", 0))
        downloaded = 0
        last_log = time.monotonic()
        while chunk := response.read(1024 * 1024):
            f.write(chunk)
            downloaded += len(chunk)
            now = time.monotonic()
            if now - last_log >= interval_s:
                pct = 100 * downloaded / total if total else 0
                print(f"{label}: {downloaded / 1e9:.2f}/{total / 1e9:.2f}GB ({pct:.0f}%)")
                last_log = now
    tmp_path.rename(dest)
    print(f"{label}: done ({dest.stat().st_size / 1e9:.2f}GB)")


# SheetSage (this model's audio feature extractor) was trained on ~24s segments and is
# most accurate on short clips; longer songs also risk exceeding the GPU time budget below.
MAX_INPUT_SECONDS = 30.0

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

decoder = None
decoder_ready = False  # has the decoder been moved onto the GPU yet?
tokenizer = None
beat_detector = None
sheetsage_model = None

model_loading = True
model_error = None


def load_assets():
    """Downloads PiCoGen2's checkpoint and Jukebox's weights, stages SheetSage's vendored
    checkpoints, and builds the beat tracker -- all on CPU so the app can start serving
    while this runs."""
    global decoder, tokenizer, beat_detector, model_loading, model_error
    try:
        tokenizer = picogen2.Tokenizer()
        decoder = picogen2.PiCoGenDecoder.from_pretrained(device="cpu")

        # SheetSage's upstream S3 bucket (its own retrieve_asset download source) has
        # been dead for a while (see https://github.com/chrisdonahue/sheetsage/issues/44,
        # 45, 46). Its checkpoints are small (245MB total) and vendored directly in this
        # repo instead; stage them where sheetsage.assets.retrieve_asset expects to find
        # them so it treats them as already downloaded.
        sheetsage_cache = Path.home() / ".sheetsage" / "sheetsage" / "v0.2"
        sheetsage_cache.mkdir(parents=True, exist_ok=True)
        for item in (REPO_ROOT / "sheetsage" / "weights").iterdir():
            dest = sheetsage_cache / item.name
            if item.name.startswith(".") or dest.exists():
                continue
            if item.is_dir():
                shutil.copytree(item, dest)
            else:
                shutil.copy(item, dest)

        # Jukebox's own weights are hosted on OpenAI's CDN, unrelated to (and unaffected
        # by) SheetSage's dead bucket. Same default cache path jukebox's own downloader
        # uses (see jukebox/make_models.py:load_checkpoint), pre-fetched here so the
        # first real request doesn't have to wait on a 10GB download inside its GPU
        # time budget.
        jukebox_cache = Path(os.environ.get("JUKEBOX_CACHE_DIR", "~/.cache")).expanduser()
        _download_with_progress(
            "https://openaipublic.azureedge.net/jukebox/models/5b/vqvae.pth.tar",
            jukebox_cache / "jukebox" / "models" / "5b" / "vqvae.pth.tar",
            "jukebox vqvae",
        )
        _download_with_progress(
            "https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_2.pth.tar",
            jukebox_cache / "jukebox" / "models" / "5b" / "prior_level_2.pth.tar",
            "jukebox prior_level_2",
        )

        beat_detector = BeatThis(cuda=False)  # CPU is already fast enough (~1.5s/30s clip)
        print("Models loaded (CPU); Jukebox weights ready.")
    except Exception as e:
        model_error = str(e)
        print(f"Load error: {e}")
    finally:
        model_loading = False


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


model_card = ModelCard(
    name="PiCoGen2",
    description=(
        "Generates a piano cover from a short pop song clip: SheetSage extracts melody/"
        "harmony audio features, and a GPT-NeoX decoder turns them into piano notes."
    ),
    author="Chih-Pin Tan, Hsin Ai, Yi-Hsin Chang, Shuen-Huei Guan, Yi-Hsuan Yang",
    tags=["music generation", "piano cover", "midi"],
)


@spaces.GPU(duration=300)
@torch.inference_mode()
def process_fn(input_audio_path: str, temperature: float) -> str:
    """Detects beats, extracts SheetSage audio features, and generates a piano cover."""
    global decoder, sheetsage_model, decoder_ready

    if model_loading:
        raise gr.Error("Model is still loading, please wait a moment and try again.")
    if decoder is None:
        raise gr.Error(f"Model failed to load: {model_error}")
    if DEVICE != "cuda":
        raise gr.Error("This model requires a GPU; none is available.")

    duration = sf.info(input_audio_path).duration
    if duration > MAX_INPUT_SECONDS:
        raise gr.Error(
            f"Input is {duration:.1f}s long; please trim it to {MAX_INPUT_SECONDS:.0f}s "
            "or shorter."
        )

    if not decoder_ready:
        decoder = decoder.to(DEVICE)  # only safe here, inside @spaces.GPU
        decoder_ready = True
    if sheetsage_model is None:
        sheetsage_model = SheetSage()  # constructs Jukebox on the GPU internally

    beats, downbeats = beat_detector(input_audio_path)
    if len(downbeats) < 2:
        # SheetSage needs at least one full bar (2 downbeats) to run
        raise gr.Error("Input audio is too short; try a longer audio.")
    beat_information = {"beats": beats.tolist(), "downbeats": downbeats.tolist()}

    sheetsage_output = sheetsage_model(
        audio_path=input_audio_path, beat_information=beat_information
    )

    out_events = picogen2.decode(
        model=decoder,
        tokenizer=tokenizer,
        beat_information=beat_information,
        melody_last_embs=sheetsage_output["melody_last_hidden_state"],
        harmony_last_embs=sheetsage_output["harmony_last_hidden_state"],
        temperature=temperature,
        device=DEVICE,
    )

    with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as f:
        output_midi_path = f.name
    tokenizer.events_to_midi(out_events).dump(output_midi_path)
    return output_midi_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio")
        .set_info(f"Short song clip, up to {MAX_INPUT_SECONDS:.0f}s")
        .harp_required(True),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            step=0.1,
            value=1.0,
            label="Temperature",
            info="Sampling temperature for the piano decoder (default: 1.0, per repo config)",
        ),
    ]
    output_components = [
        gr.File(type="filepath", label="Piano Cover", file_types=[".mid", ".midi"]).set_info(
            "Generated piano cover, as a MIDI file."
        ),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch(pwa=True)