import json
import os
import random
import secrets
import subprocess
from datetime import datetime, timezone

import gradio as gr
import spaces
import torch

COMMUNITY_BUCKET = "directedbykobyperez/YuE2-community"
COMMUNITY_URL = "https://huggingface.co/buckets/directedbykobyperez/YuE2-community"

REPO = "m-a-p/YuE2-3B"
OUT = "/tmp/yue2_out"
os.makedirs(OUT, exist_ok=True)

pipe = None


def _ensure_yue2():
    """Install the vendored yue2-infer wheel at runtime (not build time).

    It can't go in requirements.txt: it pins huggingface-hub==0.36.2 /
    transformers==4.57.6 which conflict with Gradio 6 (hub>=1.16) and break
    the Docker build. --no-deps sidesteps the pins; the package only needs
    torch/transformers/hub/safetensors/tiktoken/soundfile which are present.
    """
    try:
        import yue2  # noqa: F401
        return
    except ImportError:
        pass
    import sys
    whl = os.path.join(os.path.dirname(os.path.abspath(__file__)), "yue2_infer-0.1.5-py3-none-any.whl")
    print(f"[yue2] installing vendored wheel: {whl}")
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-deps", whl],
        check=True, capture_output=True, text=True,
    )
    print("[yue2] installed.")


def get_pipeline():
    global pipe
    if pipe is None:
        _ensure_yue2()
        from yue2 import YuE2Pipeline
        print("Loading YuE2-3B pipeline (first run downloads weights)...")
        pipe = YuE2Pipeline.from_pretrained(REPO, device="cuda" if torch.cuda.is_available() else "cpu")
        print("Pipeline loaded.")
    return pipe


def flac_to_mp3(flac_path, mp3_path, bitrate="320k"):
    subprocess.run(
        ["ffmpeg", "-y", "-i", flac_path, "-b:a", bitrate, mp3_path],
        capture_output=True, check=True,
    )
    return mp3_path


def flac_to_wav(flac_path, wav_path):
    subprocess.run(
        ["ffmpeg", "-y", "-i", flac_path, "-c:a", "pcm_s16le", wav_path],
        capture_output=True, check=True,
    )
    return wav_path


def flac_to_ogg(flac_path, ogg_path, bitrate="192k"):
    subprocess.run(
        ["ffmpeg", "-y", "-i", flac_path, "-c:a", "libopus", "-b:a", bitrate, ogg_path],
        capture_output=True, check=True,
    )
    return ogg_path


def get_duration_seconds(audio_path):
    try:
        import soundfile as sf
        return float(sf.info(audio_path).duration)
    except Exception:
        pass
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
            capture_output=True, text=True, check=True,
        )
        return float(r.stdout.strip())
    except Exception:
        return 0.0


def _upload_to_community(wav_path, meta_dict, song_id):
    """Silently share songs/<id>/<id>.wav + meta.json to the community bucket. Never raises."""
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
    if not token:
        return
    try:
        from huggingface_hub import HfFileSystem
        fs = HfFileSystem(token=token)
        meta_local = os.path.join(OUT, f"{song_id}_meta.json")
        with open(meta_local, "w", encoding="utf-8") as f:
            json.dump(meta_dict, f, ensure_ascii=False, indent=2)
        base = f"buckets/{COMMUNITY_BUCKET}/songs/{song_id}"
        fs.put_file(wav_path, f"{base}/{song_id}.wav")
        fs.put_file(meta_local, f"{base}/meta.json")
        print(f"[community] uploaded {song_id} to bucket")
    except Exception as e:
        print(f"[community] upload failed: {e}")


@spaces.GPU(duration=120)
def generate(style, lyrics, cot, seed, cfg_scale):
    pipeline = get_pipeline()
    if seed is None or seed < 1:
        try:
            seed = int(seed) if seed else random.randint(1, 10**6)
        except Exception:
            seed = random.randint(1, 10**6)
    seed = int(seed)
    gen_kwargs = dict(style=style, lyrics=lyrics, cot=cot, seed=seed)
    if cfg_scale and float(cfg_scale) != 1.0:
        gen_kwargs["cfg_scale"] = float(cfg_scale)
    song = pipeline(**gen_kwargs)

    song_id = secrets.token_hex(6)
    flac_path = os.path.join(OUT, f"song_{song_id}.flac")
    mp3_path = os.path.join(OUT, f"song_{song_id}.mp3")
    wav_path = os.path.join(OUT, f"{song_id}.wav")
    song.save(flac_path)
    flac_to_mp3(flac_path, mp3_path)
    flac_to_wav(flac_path, wav_path)
    ogg_path = os.path.join(OUT, f"{song_id}.ogg")
    try:
        flac_to_ogg(flac_path, ogg_path)
    except Exception as e:
        print(f"[ogg] skipped: {e}")
        ogg_path = None

    words = (style or "").strip().replace("\n", " ").split()[:8]
    title = " ".join(words).title()[:80] if words else f"YuE2 Song {song_id[:6]}"
    actual = get_duration_seconds(wav_path)
    _upload_to_community(wav_path, {
        "id": song_id,
        "title": title,
        "description": (style or "").strip(),
        "lyrics": lyrics or "",
        "caption": (style or "").strip(),
        "bpm": 120,
        "language": "",
        "duration": actual,
        "seed": seed,
        "generation_mode": cot or "auto",
        "audio_file": f"{song_id}.wav",
        "created_at": datetime.now(timezone.utc).isoformat(),
    }, song_id)

    return flac_path, flac_path, mp3_path, ogg_path, f"seed={seed} | {actual:.0f}s | {os.path.getsize(mp3_path) // 1024}KB MP3 320kbps"


with gr.Blocks(title="YuE2-3B Music") as demo:
    gr.Markdown("# YuE2-3B Music Generator\nFull song (vocals + accompaniment) from a style prompt and lyrics.")
    with gr.Row():
        with gr.Column():
            style = gr.Textbox(label="Style prompt", value="Cyber metal, English, aggressive male vocals, distorted guitars, double-kick drums", lines=3)
            lyrics = gr.Textbox(label="Lyrics ([verse]/[chorus] etc.)", value="[verse]\nnothing here but static skies\nwe ride the current till it dies", lines=8)
            with gr.Row():
                cot = gr.Radio(["full", "melody", "off"], value="full", label="CoT mode")
                seed = gr.Number(label="Seed", value=0, precision=0)
                cfg = gr.Number(label="CFG scale (1.0 = default)", value=1.0, precision=2)
            btn = gr.Button("Generate song", variant="primary")
            gr.HTML("<div style='text-align:center;margin:-12px 0 -8px;font-size:0.9rem;'>Don't forget to<a style='margin-left:5px;padding:0;' href='https://huggingface.co/m-a-p/YuE2-3B' target='_blank'>like the model ❤️</a></div>")
        with gr.Column():
            audio = gr.Audio(label="Play song", type="filepath")
            flac_file = gr.File(label="Download FLAC")
            mp3_file = gr.File(label="Download MP3 320kbps")
            ogg_file = gr.File(label="Download OGG (Opus)")
            info = gr.Textbox(label="Info")
    btn.click(fn=generate, inputs=[style, lyrics, cot, seed, cfg], outputs=[audio, flac_file, mp3_file, ogg_file, info], queue=True)
    gr.Markdown(
        "---\nPowered by [YuE2](https://huggingface.co/m-a-p/YuE2-3B) · "
        "🎶 **Community songs:** every generation is auto-shared to the "
        "[YuE2-community bucket](https://huggingface.co/buckets/directedbykobyperez/YuE2-community)"
        " · [YuE2 website](https://map-yue2.github.io/)"
    )

demo.queue().launch(footer_links=[], favicon_path="favicon.png")