"""YuE2-3B inference engine for ZeroGPU.

Wraps the official `yue2_infer` package shipped in the m-a-p/YuE2-3B repo: one AR/NAR backbone
plans a chord-annotated ABC score, writes 25 Hz semantic tokens, turns them into acoustic latents
by flow matching, and YuE2-Vae decodes 48 kHz stereo. Everything here runs at import time (model
load) except `generate_wav`, which must be called from inside a `@spaces.GPU` function.
"""

import re
import subprocess
import sys
import time

import spaces  # noqa: F401  must precede torch (ZeroGPU patches CUDA init)
import numpy as np
import torch

MODEL_ID = "m-a-p/YuE2-3B"
VAE_ID = "m-a-p/YuE2-Vae"
# The inference package lives in the model repo. Its metadata hard-pins huggingface-hub 0.x, which
# Gradio 6 cannot run with, so it goes in without dependencies (the real ones are in requirements.txt).
WHEEL_URL = f"https://huggingface.co/{MODEL_ID}/resolve/main/yue2_infer-0.1.5-py3-none-any.whl"

try:
    import yue2  # noqa: F401
except ImportError:
    subprocess.run([sys.executable, "-m", "pip", "install", "--no-deps", "-q", WHEEL_URL], check=True)

from yue2 import YuE2Pipeline  # noqa: E402
from yue2.modeling_vae import YuE2VAE  # noqa: E402
from yue2.modeling_yue2 import YuE2ForCausalLM  # noqa: E402
from yue2.protocol import CONTEXT, SongRequest, token_prefixes  # noqa: E402

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

# YuE2Pipeline.__init__ caps the process at a 24 GiB memory fraction, which forces a real CUDA init
# under ZeroGPU's startup emulation; skip only that cap.
_set_memory_fraction = torch.cuda.set_per_process_memory_fraction
torch.cuda.set_per_process_memory_fraction = lambda *args, **kwargs: None
try:
    PIPE = YuE2Pipeline.from_pretrained(MODEL_ID, vae=VAE_ID, device=DEVICE, progress=False)
finally:
    torch.cuda.set_per_process_memory_fraction = _set_memory_fraction

# Weights are loaded here rather than lazily by the pipeline: transformers 5 takes `dtype` (the
# package passes `torch_dtype`), and the VAE stays resident next to the backbone, whereas the
# package's decode() moves the 7 GB backbone to the CPU and back for every song.
PIPE._model = YuE2ForCausalLM.from_pretrained(
    PIPE.model_dir, local_files_only=True, dtype=torch.bfloat16
).eval().to(DEVICE)
VAE = YuE2VAE.from_pretrained(PIPE.vae_dir, decoder_only=True, device=DEVICE, local_files_only=True)

SAMPLE_RATE = int(VAE.config.sample_rate)
FRAME_RATE = SAMPLE_RATE / VAE.config.downsampling_ratio  # semantic tokens per second of audio
ABC_MAX_TOKENS = PIPE.generation_config.abc.max_tokens
ODE_STEPS = PIPE.generation_config.ode_steps
MAX_AUDIO_SECONDS = 240.0
MAX_STYLE_CHARS, MAX_LYRICS_CHARS = 1500, 10000
FADE_SECONDS = 3.0

_TAG_LINE = re.compile(r"^\s*\[[^\]]*\]\s*$")


def is_instrumental(lyrics: str) -> bool:
    """True when the lyrics hold only section tags, i.e. there are no words to sing."""
    return not any(line.strip() and not _TAG_LINE.match(line) for line in (lyrics or "").splitlines())


def _cot(lyrics: str) -> str:
    # Full planning writes a melody+chords score whose lead is a Vocal voice, which YuE2 then sings;
    # instrumentals skip the symbolic plan.
    return "off" if is_instrumental(lyrics) else "full"


def max_audio_seconds(duration: float) -> float:
    """Hard cap on the rendered length. YuE2 picks the song length from the lyrics (the composer
    sizes them to `duration`); the cap only bounds GPU time, with headroom for intros and outros."""
    return min(float(duration) * 1.3 + 15.0, MAX_AUDIO_SECONDS)


def _gpu_time(audio_seconds: float, abc_tokens: float) -> float:
    # Fitted on ZeroGPU with this same package (hugging-apps/yue2-instrumental-cot-lora): ~104 tok/s
    # while planning, ~0.265 s per audio-second for semantic decoding + VAE, ~0.0023 s per
    # audio-second per ODE step, ~11 s fixed (worker start + VAE).
    return 11.0 + abc_tokens / 104.0 + audio_seconds * (0.265 + 0.0023 * ODE_STEPS)


def estimate_gpu_seconds(duration: float) -> int:
    """ZeroGPU reservation: worst case (ABC plan at its token cap, song at its length cap) + 10%."""
    return int(min(_gpu_time(max_audio_seconds(duration), ABC_MAX_TOKENS) * 1.1, 300))


def eta_seconds(duration: float) -> int:
    """Typical wall time for a song of ~`duration` seconds (plans run 1.6k-2.7k ABC tokens)."""
    return int(_gpu_time(float(duration) * 1.1, 2200))


def validate(style: str, lyrics: str) -> None:
    """CPU-side input validation (raise ValueError with a user-facing message).

    Runs before the ZeroGPU call: exceptions raised inside the GPU worker lose their message.
    """
    style = (style or "").strip()
    lyrics = (lyrics or "").strip()
    if not style:
        raise ValueError("A style description is required.")
    if not lyrics:
        raise ValueError("Lyrics are required (section tags only, e.g. [Instrumental], make an instrumental).")
    if len(style) > MAX_STYLE_CHARS:
        raise ValueError(f"Keep the style under {MAX_STYLE_CHARS} characters.")
    if len(lyrics) > MAX_LYRICS_CHARS:
        raise ValueError(f"Keep the lyrics under {MAX_LYRICS_CHARS} characters.")
    prompt_tokens = len(token_prefixes(SongRequest(style=style, lyrics=lyrics, cot=_cot(lyrics)), PIPE.tokenizer))
    if prompt_tokens + ABC_MAX_TOKENS + 3 + int(MAX_AUDIO_SECONDS * FRAME_RATE) > CONTEXT:
        raise ValueError("The lyrics are too long for YuE2's context window; please shorten them.")


@torch.inference_mode()
def generate_wav(style: str, lyrics: str, duration: float, seed: int):
    """Run plan -> semantic tokens -> flow matching -> VAE. Must run inside @spaces.GPU.

    Returns (int16 stereo ndarray [samples, 2], sample_rate, audio_seconds, wall_seconds).
    """
    style, lyrics = (style or "").strip(), (lyrics or "").strip()
    start = time.time()
    plan = PIPE.plan(style=style, lyrics=lyrics, cot=_cot(lyrics), seed=int(seed))
    max_tokens = int(max_audio_seconds(duration) * FRAME_RATE)
    semantic = PIPE.generate_semantic(plan, sampling={"max_tokens": max_tokens})
    latents = PIPE.synthesize(semantic)
    audio = VAE.decode_tiled(
        torch.from_numpy(latents).T.contiguous().unsqueeze(0), core_frames=1024, halo_frames=16, output_device="cpu"
    )
    wav = audio[0].clamp(-1.0, 1.0).T.contiguous().numpy()
    if semantic.truncated:  # stopped at the length cap: fade out rather than cut mid-phrase
        n = min(len(wav), int(FADE_SECONDS * SAMPLE_RATE))
        wav[-n:] *= np.linspace(1.0, 0.0, n, dtype=np.float32)[:, None]
    print(
        f"[engine] cot={plan.request.cot} abc_tokens={len(plan.abc_ids)} (truncated={plan.truncated}) "
        f"semantic_tokens={len(semantic.tokens)}/{max_tokens} (truncated={semantic.truncated}) "
        f"abc={plan.timing.get('seconds', 0):.1f}s semantic={semantic.timing['seconds']:.1f}s",
        flush=True,
    )
    return (wav * 32767.0).astype(np.int16), SAMPLE_RATE, len(wav) / SAMPLE_RATE, time.time() - start