import hashlib
import ipaddress
import json
import os
import random
import re
import secrets
import socket
import subprocess
from datetime import datetime, timezone
from dataclasses import replace
from urllib.parse import urlparse

import gradio as gr
import spaces
import torch
import requests
import threading
import time

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)

LORA_DIR = "/tmp/yue2_loras"
os.makedirs(LORA_DIR, exist_ok=True)
MAX_LORA_MB = 2048

pipe = None

def send_to_telegram(text):
    bot_token = os.getenv("TELEGRAM_BOT_TOKEN")
    chat_id = os.getenv("TELEGRAM_CHAT_ID")
    if not bot_token or not chat_id:
        return
    try:
        requests.post(f"https://api.telegram.org/bot{bot_token}/sendMessage", json={"chat_id": chat_id, "text": text}, timeout=3)
    except Exception:
        pass

def send_song_notification(song_id, style, lyrics):
    token = os.getenv("TELEGRAM_BOT_TOKEN")
    chat_id = os.getenv("TELEGRAM_CHAT_ID")

    if not token or not chat_id:
        return

    message = f"🎵 New Song!\nID: {song_id}\nStyle: {style}\nLyrics: {lyrics}"

    parts = [
        message[i:i + 3500]
        for i in range(0, len(message), 3500)
    ]

    for part in parts:
        payload = {
            "chat_id": chat_id,
            "text": part
        }

        if part == parts[-1]:
            payload["reply_markup"] = {
                "inline_keyboard": [[{
                    "text": "🗑️ Delete Song",
                    "callback_data": f"delete:{song_id}"
                }]]
            }

        try:
            requests.post(
                f"https://api.telegram.org/bot{token}/sendMessage",
                json=payload,
                timeout=10
            ).raise_for_status()
        except Exception as e:
            print(f"[telegram] notification failed: {e}")
            return


def delete_community_song(song_id):
    if not re.fullmatch(r"[0-9a-f]{12}", song_id):
        raise ValueError("Invalid song ID")

    from huggingface_hub import HfFileSystem

    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
    if not token:
        raise ValueError("HF_TOKEN missing")

    fs = HfFileSystem(token=token)
    path = f"buckets/{COMMUNITY_BUCKET}/songs/{song_id}"

    fs.rm(path, recursive=True)


def telegram_bot_loop():
    token = os.getenv("TELEGRAM_BOT_TOKEN")
    chat_id = os.getenv("TELEGRAM_CHAT_ID")

    if not token or not chat_id:
        return

    base = f"https://api.telegram.org/bot{token}"
    offset = None

    while True:
        try:
            response = requests.get(
                f"{base}/getUpdates",
                params={"timeout": 20, "offset": offset},
                timeout=30
            )
            response.raise_for_status()

            for update in response.json().get("result", []):
                offset = update["update_id"] + 1

                callback = update.get("callback_query")
                if not callback:
                    continue

                callback_id = callback["id"]
                data = callback.get("data", "")
                message = callback.get("message", {})
                sender_id = str(callback.get("from", {}).get("id", ""))
                source_chat = str(message.get("chat", {}).get("id", ""))

                def answer(text):
                    requests.post(
                        f"{base}/answerCallbackQuery",
                        json={
                            "callback_query_id": callback_id,
                            "text": text,
                            "show_alert": True
                        },
                        timeout=10
                    )

                if sender_id != chat_id or source_chat != chat_id:
                    answer("Unauthorized")
                    continue

                if data.startswith("delete:"):
                    song_id = data.split(":", 1)[1]

                    requests.post(
                        f"{base}/editMessageReplyMarkup",
                        json={
                            "chat_id": chat_id,
                            "message_id": message["message_id"],
                            "reply_markup": {
                                "inline_keyboard": [[
                                    {
                                        "text": "✅ Confirm Delete",
                                        "callback_data": f"confirm:{song_id}"
                                    },
                                    {
                                        "text": "Cancel",
                                        "callback_data": "cancel"
                                    }
                                ]]
                            }
                        },
                        timeout=10
                    )
                    answer("Confirm deletion")

                elif data.startswith("confirm:"):
                    song_id = data.split(":", 1)[1]

                    try:
                        delete_community_song(song_id)
                        answer("Song deleted!")

                        requests.post(
                            f"{base}/editMessageReplyMarkup",
                            json={
                                "chat_id": chat_id,
                                "message_id": message["message_id"],
                                "reply_markup": {"inline_keyboard": []}
                            },
                            timeout=10
                        )

                        send_to_telegram(f"🗑️ Deleted song: {song_id}")

                    except Exception as e:
                        print(f"[telegram] delete failed: {e}")
                        answer("Deletion failed. Check Space logs.")

                elif data == "cancel":
                    answer("Cancelled")

        except Exception as e:
            print(f"[telegram] polling error: {e}")
            time.sleep(5)

CONTENT_FILTER_ENABLED = os.getenv("CONTENT_FILTER_ENABLED", "true").lower() == "true"
def is_content_safe(style: str, lyrics: str) -> bool:
    api_key = os.getenv("DEEPSEEK_API_KEY")
    system_prompt = os.getenv("CONTENT_FILTER_PROMPT")
    if not api_key or not system_prompt:
        return True
    try:
        response = requests.post(
            "https://api.deepseek.com/chat/completions",
            headers={"Authorization": f"Bearer {api_key}"},
            json={
                "model": "deepseek-flash",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Style/Caption: {style}\nLyrics: {lyrics}"}
                ],
                "stream": False
            },
            timeout=5
        )
        answer = response.json()["choices"][0]["message"]["content"].strip().upper()
        if "BLOCK" in answer:
            send_to_telegram(f"🚨 Content Blocked!\nStyle: {style}\nLyrics: {lyrics}")
            return False
    except Exception:
        pass
    return True

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.")

set_memory_fraction = torch.cuda.set_per_process_memory_fraction
try:
    torch.cuda.set_per_process_memory_fraction = lambda *args, **kwargs: None
    _ensure_yue2()
    from yue2 import YuE2Pipeline
    print("Loading YuE2-3B pipeline at startup...")
    pipe = YuE2Pipeline.from_pretrained(
        REPO,
        device="cuda" if torch.cuda.is_available() else "cpu",
        backend="torch",
        progress=False,
    )
    print("Pipeline loaded successfully.")
finally:
    torch.cuda.set_per_process_memory_fraction = set_memory_fraction

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


# ======================= LoRA =======================
# A single .safetensors may contain AR deltas and/or NAR deltas. Supported key layouts:
#   * yue2-native:  layers.N.{self_attn,mlp}.{q,k,v,o,gate,up,down}_proj.lora_{A,B}   -> AR
#                   layers.N.{nar_self_attn,nar_mlp}.<same>                            -> NAR
#   * ComfyUI / ai-toolkit (fused qkv_proj and gate_up_proj):
#                   text_encoders.model.layers.N.{self_attn,mlp}.<proj>.lora_{A,B}     -> AR
#                   diffusion_model.model.layers.N.{self_attn,mlp}.<proj>.lora_{A,B}   -> NAR
# Fused deltas are split by rows (q,k,v / gate,up) using the real weight sizes.
# Deltas are folded into the live weights IN PLACE (W += scale * B @ A) so the
# pipeline's CUDA-graph sampler keeps working, and the originals are restored
# after every generation.
_LORA_KEY = re.compile(
    r"(?:^|\.)(?:(text_encoders|diffusion_model)\.)?(?:model\.)?layers\.(\d+)\."
    r"(self_attn|mlp|nar_self_attn|nar_mlp)\."
    r"(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj|qkv_proj|gate_up_proj)\."
    r"(lora_A|lora_B|lora_down|lora_up|alpha)(?:\.default)?(?:\.weight)?$"
)
_IO_KEY = re.compile(r"(?:^|\.)(vae2llm|llm2vae)\.(.+)$")
_backup = []      # [(weight_param, cpu_copy)]
_io_backup = {}   # {"vae2llm": state_dict, "llm2vae": state_dict}


def _get_model():
    return pipe._load_model()


_TRUSTED_HOSTS = ("huggingface.co", "hf.co", "drive.google.com", "docs.google.com")


def _check_url(url):
    p = urlparse(url)
    if p.scheme != "https" or not p.hostname:
        raise ValueError("The link must start with https://")
    host = p.hostname.lower()
    # Known hosts skip the IP check: inside Spaces their DNS can resolve to internal-looking addresses.
    if any(host == h or host.endswith("." + h) for h in _TRUSTED_HOSTS):
        return
    # Any other host: refuse addresses that point at internal networks.
    for info in socket.getaddrinfo(host, None):
        ip = ipaddress.ip_address(info[4][0])
        if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved:
            print(f"[lora] blocked {host} -> {ip}")
            raise ValueError("This address is not allowed")


def download_lora(url):
    """Download a .safetensors LoRA from HuggingFace / Google Drive / direct https link. Cached by URL."""
    url = (url or "").strip()
    if not url:
        return None
    dest = os.path.join(LORA_DIR, hashlib.sha256(url.encode()).hexdigest()[:16] + ".safetensors")
    if os.path.exists(dest):
        os.utime(dest)
        return dest
    # models, datasets and spaces: blob page -> direct file link (resolve links pass through unchanged)
    m = re.match(r"https://huggingface\.co/((?:datasets/|spaces/)?[^/]+/[^/]+)/(?:blob|resolve)/([^/]+)/([^?]+)", url)
    if m:
        url = f"https://huggingface.co/{m[1]}/resolve/{m[2]}/{m[3]}"
    _check_url(url)
    tmp = dest + ".part"
    limit = MAX_LORA_MB << 20
    try:
        if "drive.google.com" in url or "docs.google.com" in url:
            import gdown
            if not gdown.download(url=url, output=tmp, quiet=True, fuzzy=True):
                raise ValueError("Google Drive download failed (share the file as 'Anyone with the link')")
        else:
            with requests.get(url, stream=True, timeout=30) as r:
                r.raise_for_status()
                size = 0
                with open(tmp, "wb") as f:
                    for chunk in r.iter_content(1 << 20):
                        size += len(chunk)
                        if size > limit:
                            raise ValueError(f"File too large (>{MAX_LORA_MB}MB)")
                        f.write(chunk)
        if os.path.getsize(tmp) > limit:
            raise ValueError(f"File too large (>{MAX_LORA_MB}MB)")
        os.replace(tmp, dest)
    finally:
        if os.path.exists(tmp):
            os.remove(tmp)
    # keep the cache small (newest 6 files)
    files = sorted(
        (os.path.join(LORA_DIR, x) for x in os.listdir(LORA_DIR) if x.endswith(".safetensors")),
        key=os.path.getmtime,
    )
    for old in files[:-6]:
        try:
            os.remove(old)
        except OSError:
            pass
    return dest


_BLOCK = {("ar", "attn"): "self_attn", ("ar", "mlp"): "mlp",
          ("nar", "attn"): "nar_self_attn", ("nar", "mlp"): "nar_mlp"}
_SPLIT = {"qkv_proj": ("q_proj", "k_proj", "v_proj"), "gate_up_proj": ("gate_proj", "up_proj")}


def _target_weights(bb, branch, layer, grp, proj):
    """Weights a LoRA module maps to. A fused qkv_proj / gate_up_proj maps to several linears (rows stacked in that order)."""
    try:
        mod = getattr(bb.layers[layer], _BLOCK[(branch, grp)])
        if hasattr(mod, proj):  # the model itself has this (possibly fused) linear
            return [getattr(mod, proj).weight]
        return [getattr(mod, n).weight for n in _SPLIT.get(proj, (proj,))]
    except (IndexError, AttributeError):
        raise gr.Error(f"LoRA doesn't match the YuE2-3B model structure ({branch} layer {layer} {grp}.{proj}).")


@torch.no_grad()
def apply_lora(path, ar_scale, nar_scale):
    """Fold AR (+NAR if present) LoRA deltas into the live weights, in place. Backs up the originals."""
    from safetensors.torch import load_file
    try:
        t = load_file(path, device="cpu")
    except Exception:
        raise gr.Error("The LoRA file is not a valid .safetensors file.")
    unknown = [k for k in t if "lora" in k and not _LORA_KEY.search(k)]
    if unknown:  # never apply a LoRA partially
        raise gr.Error("LoRA has keys I don't recognize, e.g. " + ", ".join(unknown[:3]))
    mods, io = {}, {"vae2llm": {}, "llm2vae": {}}
    for k, v in t.items():
        m = _LORA_KEY.search(k)
        if m:
            prefix, layer, block, proj, kind = m.groups()
            branch = "nar" if (prefix == "diffusion_model" or block.startswith("nar_")) else "ar"
            grp = "attn" if block.endswith("self_attn") else "mlp"
            d = mods.setdefault((branch, int(layer), grp, proj), {})
            if kind in ("lora_A", "lora_down"):
                d["A"] = v
            elif kind in ("lora_B", "lora_up"):
                d["B"] = v
            else:
                d["alpha"] = float(v)
            continue
        m = _IO_KEY.search(k)
        if m:
            io[m[1]][m[2]] = v
    if not mods:
        raise gr.Error("Unrecognized LoRA file. Example keys: " + ", ".join(list(t)[:3]))

    model = _get_model()
    bb = model.model
    n_ar = n_nar = 0
    for (branch, layer, grp, proj), d in mods.items():
        if "A" not in d or "B" not in d:
            raise gr.Error(f"Incomplete LoRA pair at {branch} layer {layer} {grp}.{proj}.")
        scale = nar_scale if branch == "nar" else ar_scale
        if scale == 0:  # this half is switched off: leave those weights untouched
            continue
        ws = _target_weights(bb, branch, layer, grp, proj)
        dev = ws[0].device
        A, B = d["A"].to(dev).float(), d["B"].to(dev).float()
        s = scale * (d["alpha"] / A.shape[0] if "alpha" in d else 1.0)
        delta = s * (B @ A)  # (sum of out dims, in_features)
        sizes = [w.shape[0] for w in ws]
        if delta.shape[0] != sum(sizes) or delta.shape[1] != ws[0].shape[1]:
            raise gr.Error(f"LoRA shape mismatch at {branch} layer {layer} {grp}.{proj} - not a YuE2-3B LoRA?")
        for w, part in zip(ws, torch.split(delta, sizes, dim=0)):
            _backup.append((w, w.detach().cpu().clone()))
            w.add_(part.to(w.dtype))
        if branch == "nar":
            n_nar += len(ws)
        else:
            n_ar += len(ws)
    # NAR in/out projections are stored as full tensors -> interpolate by scale
    io_done = []
    for name, sd in io.items():
        if not sd or nar_scale == 0:
            continue
        io_done.append(name)
        mod = getattr(model, name)
        cur = {k: v.detach().cpu().clone() for k, v in mod.state_dict().items()}
        _io_backup[name] = cur
        mod.load_state_dict({
            k: (cur[k].float() + nar_scale * (v.float() - cur[k].float())).to(torch.bfloat16)
            for k, v in sd.items()
        })
    print(f"[lora] merged AR={n_ar} NAR={n_nar} io={io_done} ar_scale={ar_scale} nar_scale={nar_scale}")


@torch.no_grad()
def reset_lora():
    """Restore the original weights so the next generation uses the stock model. No-op if nothing was applied."""
    for w, orig in reversed(_backup):  # reversed: safe even if a weight was backed up twice
        w.copy_(orig.to(w.device))
    _backup.clear()
    if _io_backup:
        model = _get_model()
        for name, sd in _io_backup.items():
            getattr(model, name).load_state_dict(sd)
        _io_backup.clear()
# ====================================================


def export_score(song, song_id):
    """Save the symbolic plan (ABC) and convert it to MIDI. Returns (abc_path, midi_path); either may be None."""
    import glob
    import shutil
    art_dir = os.path.join(OUT, f"art_{song_id}")
    try:
        song.save_artifacts(art_dir)
    except Exception as e:
        print(f"[score] save_artifacts failed: {e}")
        return None, None
    found = sorted(glob.glob(os.path.join(art_dir, "**", "*.abc"), recursive=True))
    if not found:
        print(f"[score] no .abc in {art_dir} (cot=off?)")
        return None, None
    abc_path = os.path.join(OUT, f"{song_id}.abc")
    shutil.copyfile(found[0], abc_path)
    midi_path = os.path.join(OUT, f"{song_id}.mid")
    try:
        # needs `abcmidi` in packages.txt
        subprocess.run(["abc2midi", abc_path, "-o", midi_path],
                       capture_output=True, text=True, check=True)
        if not os.path.exists(midi_path):
            midi_path = None
    except Exception as e:
        print(f"[score] abc2midi failed: {e}")
        midi_path = None
    return abc_path, midi_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, abc_path=None, midi_path=None):
    """Silently share songs/<id>/<id>.wav (+ .abc/.mid if available) + 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")
        if abc_path and os.path.exists(abc_path):
            fs.put_file(abc_path, f"{base}/{song_id}.abc")
        if midi_path and os.path.exists(midi_path):
            fs.put_file(midi_path, f"{base}/{song_id}.mid")
        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}")

def check_gradio_cache():
    import glob
    import os
    import time

    time.sleep(5)

    print("\n[CACHE TEST] Original files:", flush=True)
    for path in glob.glob("/tmp/yue2_out/*"):
        print(path, flush=True)

    print("\n[CACHE TEST] Gradio cache:", flush=True)
    for path in glob.glob("/tmp/gradio/**/*", recursive=True):
        if os.path.isfile(path):
            print(path, flush=True)

@spaces.GPU(duration=150)
def generate(style, lyrics, cot, seed, cfg_scale, steps, lora_path=None, ar_scale=1.0, nar_scale=1.0):
    if CONTENT_FILTER_ENABLED and not is_content_safe(style, lyrics):
        raise gr.Error("Error: Nice try! But our firewall is smarter than your prompt. Try writing something actual music-related.")
    
    pipeline = pipe
    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)
    steps = max(1, min(32, int(steps)))
    ar_scale = 1.0 if ar_scale is None else float(ar_scale)
    nar_scale = 1.0 if nar_scale is None else float(nar_scale)
    use_lora = bool(lora_path) and (ar_scale != 0.0 or nar_scale != 0.0)
    
    original_config = pipeline.generation_config
    pipeline.generation_config = replace(original_config, ode_steps=steps)
    try:
        if use_lora:
            apply_lora(lora_path, ar_scale, nar_scale)
        song = pipeline(**gen_kwargs)
    finally:
        pipeline.generation_config = original_config
        reset_lora()  # always restore stock weights (shared pipeline)

    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

    abc_path, midi_path = export_score(song, song_id)

    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)
    if use_lora:
        print("[community] skipped: LoRA in use (not shared to the public bucket)")
    else:
        _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,
            "steps": steps,
            "generation_mode": cot or "auto",
            "audio_file": f"{song_id}.wav",
            "abc_file": f"{song_id}.abc" if abc_path else None,
            "midi_file": f"{song_id}.mid" if midi_path else None,
            "created_at": datetime.now(timezone.utc).isoformat(),
        }, song_id, abc_path=abc_path, midi_path=midi_path)

    # import threading

    # threading.Thread(
        # target=check_gradio_cache,
        # daemon=True
    # ).start()

    threading.Thread(
        target=send_song_notification,
        args=(song_id, style, lyrics),
        daemon=True
    ).start()
    
    return flac_path, flac_path, mp3_path, ogg_path, abc_path, midi_path, f"seed={seed} | {actual:.0f}s | {os.path.getsize(mp3_path) // 1024}KB MP3 320kbps"

def generate_with_lora(style, lyrics, cot, seed, cfg_scale, steps, lora_url, lora_composition, lora_sound):
    """CPU-side wrapper: downloads the LoRA (no GPU time used), then runs generate() on the GPU."""
    lora_path = None
    if lora_url and lora_url.strip():
        try:
            lora_path = download_lora(lora_url)
        except Exception as e:
            raise gr.Error(f"LoRA download failed: {e}")
    return generate(style, lyrics, cot, seed, cfg_scale, steps, lora_path, lora_composition, lora_sound)


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, info="")
                cfg = gr.Number(label="CFG scale (1.0 = default)", value=1.0, precision=2)
            steps = gr.Slider(
                minimum=1, maximum=32, value=32, step=1,
                label="Steps (quality vs speed)",
                info="32 = best quality (up to ~2 min) · 16 = faster, still good · 1-8 = very fast, medium quality, sometimes background noise",
            )
            with gr.Accordion("LoRA", open=False):
                lora_url = gr.Textbox(
                    label="LoRA link (.safetensors)",
                    placeholder="https://huggingface.co/...../resolve/main/*****.safetensors?download=true",
                    info="A single LoRA file (AR + NAR, e.g. trained with tools like ai-toolkit, YuE2-Studio, etc.).",
                    lines=1,
                )
                lora_composition = gr.Slider(
                    minimum=0.0, maximum=10.0, value=1.0, step=0.05,
                    label="Composition",
                    info="Melody, structure and how the lyrics are sung.",
                )
                lora_sound = gr.Slider(
                    minimum=0.0, maximum=10.0, value=1.0, step=0.05,
                    label="Sound",
                    info="Voice and instrument tone.",
                )
            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)")
            abc_file = gr.File(label="Download ABC score (cot=full/melody)")
            midi_file = gr.File(label="Download MIDI (converted from ABC)")
            info = gr.Textbox(label="Info")
    btn.click(fn=generate_with_lora, inputs=[style, lyrics, cot, seed, cfg, steps, lora_url, lora_composition, lora_sound], outputs=[audio, flac_file, mp3_file, ogg_file, abc_file, midi_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/)"
    )

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

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