"""LLM composer: plain-English song description -> YuE2 inputs.

Produces a title, short display tags, a YuE2 style prompt (comma-separated descriptors: language,
genre, singer, instruments, mood, tempo, as in the official examples) and section-tagged lyrics.
Runs on the HF Inference Router (no GPU).
"""

import json
import os
import re
import time

from openai import OpenAI

ROUTER_URL = "https://router.huggingface.co/v1"

# (model, timeout_s, extra_body). Measured per compose (~600 output tokens): gpt-oss-120b on Groq
# with low reasoning ~2s; DeepSeek-V4-Flash with thinking disabled ~7s via :fastest (Fireworks),
# 14-35s on DeepInfra. Thinking/reasoning must be off: reasoning tokens otherwise dominate latency
# (6-27s). Under load any provider can 429 transiently, so the chain is walked twice.
_NO_THINK = {"thinking": {"type": "disabled"}}
COMPOSER_MODELS = (
    ("openai/gpt-oss-120b:groq", 30, {"reasoning_effort": "low"}),
    ("deepseek-ai/DeepSeek-V4-Flash-0731:fastest", 45, _NO_THINK),
    ("deepseek-ai/DeepSeek-V4-Flash-0731:deepinfra", 75, _NO_THINK),
    ("deepseek-ai/DeepSeek-V4-Flash-0731:novita", 100, _NO_THINK),
)
VISUAL_MODELS = (
    ("openai/gpt-oss-120b:groq", 30, {"reasoning_effort": "low"}),
    ("deepseek-ai/DeepSeek-V4-Flash-0731:fastest", 45, _NO_THINK),
)

# YuE2's native section labels (the ABC planner marks the same sections in its score).
SECTIONS = ("Intro", "Verse", "Pre-Chorus", "Chorus", "Bridge", "Interlude", "Instrumental", "Outro")
_SECTION_ALIASES = {
    "pre chorus": "Pre-Chorus", "prechorus": "Pre-Chorus", "post-chorus": "Chorus", "post chorus": "Chorus",
    "hook": "Chorus", "refrain": "Chorus", "break": "Interlude", "solo": "Interlude", "inst": "Instrumental",
}
_SECTION_NAMES = {s.lower(): s for s in SECTIONS} | _SECTION_ALIASES
NO_VOCALS = "no vocals, no singing, no choir, no spoken words"

_STYLE_RULES = f"""style: ONE line of 8-14 comma-separated musical descriptors, like the official YuE2 examples: start with the lyrics language ("English", "Mandarin Chinese", ...), then genre and subgenre, the singer (gender, timbre, delivery, e.g. "breathy female voice", "gritty male baritone"), 3-5 key instruments each with a short character ("rounded electric bass", "gated snare"), mood and energy arc, production character, and the tempo as "<number> BPM". Example: "English, warm piano pop, expressive female voice, acoustic piano, rounded bass and light drums, lyrical memorable melody, unhurried phrasing, 88 BPM". Plain musical words only: no sentences, quotes, title or lyric lines. Never contradict an explicit user constraint (vocal gender, tempo, required or excluded instruments). For an instrumental, start with "Instrumental" instead of a language, name the instrument carrying the melody, and end with "{NO_VOCALS}"."""

_LYRICS_RULES = """lyrics: singable lyrics using ONLY these section tags: [Intro] [Verse] [Pre-Chorus] [Chorus] [Bridge] [Interlude] [Outro]. Each tag ALONE on its own line with the section's lines below it, and one blank line between sections. [Intro] and [Interlude] are usually left empty (instrumental passages). Size the structure to the duration: <=30s: one [Verse] + one [Chorus]; ~60s: [Intro] [Verse] [Chorus] [Verse] [Chorus]; ~120s: [Intro] [Verse] [Pre-Chorus] [Chorus] [Verse] [Pre-Chorus] [Chorus] [Bridge] [Chorus] [Outro]; >=180s: that full structure plus an [Interlude] and a final repeated [Chorus]. 4 lines per sung section, short rhythmic lines of 5-9 words, roughly 12-16 sung words per 10 seconds; repeat the chorus words each time it returns. Musical instructions (tempo, instruments, dynamics) never belong in the lyrics. Write in the language the user asks for (default: English; YuE2 is strongest in English and Mandarin Chinese). If the song is instrumental, the lyrics are ONLY the section tags, with no words."""

COMPOSER_SYSTEM = f"""You write inputs for YuE2, an open music generation model that sings complete songs from a style prompt and section-tagged lyrics.
Given a song description and a target duration, produce:
1. title: a short, catchy song title (2-5 words, no quotes).
2. tags: 3-5 short comma-separated style tags for a music feed card, e.g. "synth-pop, female vocals, 120 bpm, euphoric".
3. {_STYLE_RULES}
4. {_LYRICS_RULES}
Unless the user explicitly asks for an instrumental, the song HAS a singer: the style must describe that singer and the lyrics must have words to sing.
Answer with ONLY a JSON object with keys: title, tags, style, lyrics. Inside the lyrics string, separate lines with JSON newline escapes, never with the two literal characters backslash and n."""

VISUAL_SYSTEM = (
    "Reply with exactly ONE concrete visual noun (a physical object, animal, or natural element) "
    "that captures the essence of this song. No explanation, no punctuation, just the single word."
)

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


def _client():
    key = os.environ.get("HF_TOKEN", "")
    if not key:
        raise RuntimeError("HF_TOKEN is not configured on this Space (needed for the composer LLM).")
    return OpenAI(base_url=ROUTER_URL, api_key=key, max_retries=0)


def _chat(system, user, *, temperature=0.8, max_tokens=6000, models=COMPOSER_MODELS, passes=2):
    client = _client()
    last_error = None
    for attempt in range(passes):
        for model, timeout, extra in models:
            try:
                completion = client.with_options(timeout=timeout).chat.completions.create(
                    model=model,
                    messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
                    temperature=temperature,
                    max_tokens=max_tokens,
                    extra_body=extra,
                )
                text = completion.choices[0].message.content or ""
                text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
                if not text:
                    raise ValueError(f"empty reply (finish_reason={completion.choices[0].finish_reason})")
                return text
            except Exception as e:  # noqa: BLE001
                print(f"[composer] {model} failed (pass {attempt + 1}): {type(e).__name__}: {e}", flush=True)
                last_error = e
        if attempt < passes - 1:
            time.sleep(3)
    raise RuntimeError("The composer model is overloaded right now; please try again in a moment.") from last_error


def _parse_json(text, required):
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end <= start:
        raise ValueError("no JSON object in composer reply")
    data = json.loads(text[start : end + 1], strict=False)  # tolerate raw newlines inside strings
    missing = [k for k in required if not str(data.get(k, "")).strip()]
    if missing:
        raise ValueError(f"composer reply missing keys: {missing}")
    return data


def _section_tag(label: str) -> str:
    """Canonical YuE2 tag: "verse 2" -> [Verse], "pre chorus" -> [Pre-Chorus], "hook" -> [Chorus]."""
    key = re.sub(r"\s+", " ", re.sub(r"[\s_:-]*\d+$", "", label.strip().lower()))
    return f"[{_SECTION_NAMES.get(key, label.strip().title())}]"


def normalize_lyrics(lyrics: str) -> str:
    """Canonical tags, each alone on its own line, with a blank line between sections."""
    out = []
    text = (lyrics or "").replace("\r", "").replace("\\n", "\n")
    for line in text.split("\n"):
        m = _TAG_LINE_RE.match(line)
        if m and m.group(1).strip():
            out += ["", _section_tag(m.group(1))]
            if m.group(2).strip():
                out.append(m.group(2).strip())
        else:
            out.append(line.rstrip())
    text = "\n".join(out)
    return re.sub(r"\n{3,}", "\n\n", text).strip()


def has_words(lyrics: str) -> bool:
    return any(line.strip() and not _TAG_LINE_RE.match(line) for line in (lyrics or "").split("\n"))


def compose(description: str, duration: float, instrumental: bool = False) -> dict:
    """Return {title, tags, style, lyrics}."""
    user = f"Song description: {description.strip()}\nTarget duration: {int(duration)} seconds."
    if instrumental:
        user += "\nThis song is INSTRUMENTAL: no vocals at all. The lyrics are section tags only, with no words."
    else:
        user += "\nThe song has sung vocals (pick a fitting singer if the description does not specify one)."
    required = ("title", "style", "lyrics")
    last_error = None
    for _ in range(2):
        try:
            data = _parse_json(_chat(COMPOSER_SYSTEM, user), required)
            lyrics = normalize_lyrics(str(data["lyrics"]))
            if not instrumental and not has_words(lyrics):
                raise ValueError("no sung words in the lyrics of a vocal song")
            break
        except (ValueError, json.JSONDecodeError) as e:
            print(f"[composer] bad reply: {e}", flush=True)
            last_error = e
    else:
        raise RuntimeError("The composer returned an unusable reply; please try again.") from last_error

    style = " ".join(str(data["style"]).split()).strip().strip('"\'')
    if instrumental:
        # Section tags only: the engine renders tag-only lyrics without a vocal line.
        tags_only = [l for l in lyrics.split("\n") if _TAG_LINE_RE.match(l)]
        lyrics = "\n\n".join(tags_only) or "[Instrumental]"
        if not style.lower().startswith("instrumental"):
            style = f"Instrumental, {style}"
        if "no vocals" not in style.lower():
            style = f"{style.rstrip(' ,.')}, {NO_VOCALS}"
    title = str(data["title"]).strip().strip('"\'')[:80] or "Untitled"
    tags = str(data.get("tags", "")).strip().strip(".")[:120]
    return {"title": title, "tags": tags, "style": style, "lyrics": lyrics}


def visual_word(title: str, tags: str, lyrics: str, description: str) -> str:
    """One evocative noun for the cover-art prompt; falls back to the description's first words."""
    fallback = " ".join((description or title or "music").split()[:2])
    try:
        text = _chat(
            VISUAL_SYSTEM,
            f"Title: {title}\nTags: {tags}\nLyrics: {lyrics[:300]}",
            temperature=0.7,
            max_tokens=2000,
            models=VISUAL_MODELS,
            passes=1,
        )
        word = text.split()[0].strip('."\'!,') if text.split() else ""
        return word or fallback
    except Exception as e:  # noqa: BLE001
        print(f"[composer] visual word failed: {e}", flush=True)
        return fallback