"""LLM composer: plain-English song description -> MiniMax Music 3 inputs.

Produces a title, short display tags, tagged lyrics and the three-part Structured Caption
(Global Metadata / Vocal Details / Arrangement) the model was trained on, following the
official prompting guide. 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),
)

_CAPTION_CONTRACT = """The three caption fields follow the exact labeled style the model was trained on. Be concrete and musical; describe an energy arc and instrument lifecycles, never a static equipment list or decorative adjectives. Never contradict an explicit user constraint: instrumental stays instrumental; never reverse a required vocal gender, tempo limit, required instrument, or exclusion. Do not quote or paraphrase lyric lines inside the caption. Total caption length roughly 250-400 words.

global_metadata: one paragraph, in order: "Basic Attributes: bpm is <number>. key is <letter>, and scale is <major|minor>. <Genre / Subgenre>." then "Global Emotional Progression: <how the emotion evolves from the opening through the final section>." then "Application Scenarios & Imagery: <two or three vivid listening scenarios>." then "Sonics & Production Profile: <soundstage, frequency balance, dynamics, production character>."

vocal_details: one paragraph: "Vocal Gender & Timbre: Singer A (<Male|Female>), <timbre and register>." then "Vocal Style: <delivery, and how it shifts per section>." then "Harmony/Backing Vocals: <where harmonies or doubles appear and their character>." then "Vocal FX: <restrained treatment: reverb, delay, light compression>." For instrumental pieces write "Instrumental, no vocals." and name the instrument or texture carrying the lead melodic role.

arrangement: one paragraph: "Instrument Lifecycle Description (Primary/Secondary Layering): Primary: <core instruments present start to finish and their role>. Secondary: <instruments that enter, exit or intensify, and in which sections>." then "Groove & Foundation Progression: <how drums, bass and groove develop across sections>." then "Embellishments, Textures & Spatial FX: <fills, textures, transitional gestures, stereo and space treatment where relevant>." State what enters, exits, changes or intensifies for every section of the song, aligned with the lyric section tags."""

_LYRICS_RULES = """lyrics: singable lyrics using ONLY these section tags, each ALWAYS ALONE on its own line: [intro] [verse] [pre-chorus] [chorus] [post-chorus] [bridge] [instrumental] [solo] [outro]. Never put words on the same line as a tag. Size the structure to the duration: <=30s: one verse + one chorus; ~60s: verse/pre-chorus/chorus/verse/chorus; >=120s: full structure with bridge and outro. Roughly 12-16 sung words per 10 seconds. Musical instructions (tempo, instruments, dynamics) never belong in the lyrics. If the song is instrumental, use [instrumental] sections with no words. Write the lyrics in the language the user asks for (default: English)."""

COMPOSER_SYSTEM = f"""You write inputs for MiniMax Music 3, a lyrics+description music generation model.
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. {_LYRICS_RULES}
4-6. global_metadata, vocal_details, arrangement: a structured caption. {_CAPTION_CONTRACT}
Unless the user explicitly asks for an instrumental, the song HAS a singer: vocal_details must describe that singer (gender, timbre, style) and must never say "Instrumental".
Answer with ONLY a JSON object with keys: title, tags, lyrics, global_metadata, vocal_details, arrangement. 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."
)

SECTION_TAG_RE = re.compile(r"^\s*\[(intro|verse|pre-chorus|chorus|post-chorus|bridge|instrumental|solo|outro)\]\s*$", re.I)


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 normalize_lyrics(lyrics: str) -> str:
    """Put every section tag alone on its own line (the model drops words that share a tag's line)."""
    out = []
    text = (lyrics or "").replace("\r", "").replace("\\n", "\n")
    for line in text.split("\n"):
        m = re.match(r"^\s*(\[[^\]]+\])\s*(.*)$", line)
        if m and m.group(2).strip():
            out.append(m.group(1).lower())
            out.append(m.group(2).strip())
        elif m:
            out.append(m.group(1).lower())
        else:
            out.append(line.rstrip())
    text = "\n".join(out)
    return re.sub(r"\n{3,}", "\n\n", text).strip()


def compose(description: str, duration: float, instrumental: bool = False) -> dict:
    """Return {title, tags, lyrics, caption, global_metadata, vocal_details, arrangement}."""
    user = f"Song description: {description.strip()}\nTarget duration: {int(duration)} seconds."
    if instrumental:
        user += "\nThis song is INSTRUMENTAL: no vocals at all. Use [instrumental] sections only, with no words."
    else:
        user += "\nThe song has sung vocals (pick a fitting singer if the description does not specify one)."
    required = ("title", "lyrics", "global_metadata", "vocal_details", "arrangement")
    last_error = None
    for _ in range(2):
        try:
            data = _parse_json(_chat(COMPOSER_SYSTEM, user), required)
            vd = str(data["vocal_details"]).lower()
            if not instrumental and vd.startswith("instrumental"):
                raise ValueError("vocal_details says instrumental for 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

    lyrics = normalize_lyrics(str(data["lyrics"]))
    if instrumental and not any(SECTION_TAG_RE.match(l) for l in lyrics.split("\n")):
        lyrics = "[instrumental]\n\n" + lyrics
    title = str(data["title"]).strip().strip('"\'')[:80] or "Untitled"
    tags = str(data.get("tags", "")).strip().strip(".")[:120]
    global_metadata = str(data["global_metadata"]).strip()
    vocal_details = str(data["vocal_details"]).strip()
    arrangement = str(data["arrangement"]).strip()
    caption = "\n".join(s for s in (global_metadata, vocal_details, arrangement) if s)
    return {
        "title": title,
        "tags": tags,
        "lyrics": lyrics,
        "caption": caption,
        "global_metadata": global_metadata,
        "vocal_details": vocal_details,
        "arrangement": arrangement,
    }


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