app.py

"""
Lullaby Llama — personalized bedtime music, generated on-device.

Three-layer audio:
  1. Rhythm layer  — user-selected instrument plays the chord progression.
  2. Melody layer  — user-selected instrument plays a melodic line over the chords.
  3. Voice layer   — Kokoro reads the lyrics, soft preset, slow pacing.

The user picks both layers with image buttons in the UI.

Aesthetic: children's drawing — crayon textures, wobbly hand-drawn borders.
"""

import base64
import glob
import os
import random
import re
import tempfile
import time
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
from llama_cpp import Llama

from synths.guitar import GuitarSynth, ACOUSTIC_PRESET, SR
from synths.piano import PianoSynth
from synths.harp import HarpSynth
from synths.xylophone import XylophoneSynth
from synths.ocarina import WhistleSynth
from synths.musicbox import MusicBoxSynth
from synths.voice import speak_lyrics, SR_TARGET
from utils import trace as lola_trace
try:
    from utils.safety import screen_inputs
except ImportError:
    # safety.py absent — degrade to a no-op screen rather than crash.
    def screen_inputs(loves, fears):
        return (True, "", [])

# Vision module (drawing → "loves" phrase). Falls back to stroke analysis
# internally if MiniCPM-V can't load, so the drawing input always produces
# SOMETHING for the lullaby to be about.
#
# We catch Exception (not just ImportError) so a heavy transitive failure
# doesn't crash the whole app, but we PRINT the traceback — a silent swallow
# here previously hid the real reason the vision module wasn't loading.
try:
    import draw.vision as vision # vision.describe(image) -> {"loves": str, ...}
except Exception as _vision_import_err:
    import traceback as _tb
    print("WARNING: could not import draw.vision — vision disabled, "
          "falling back to stroke analysis. Real error follows:")
    _tb.print_exc()
    vision = None


# ---------- model loading ----------

MODEL_REPO = os.environ.get(
    "LULLABY_MODEL_REPO",
    "build-small-hackathon/lolaby-llama-3b",
)
SKIP_LLM = os.environ.get("LULLABY_SKIP_LLM", "").lower() in ("1", "true", "yes")

SYSTEM_PROMPT = (
    "You write personalized lullabies for small children, with chord markers "
    "and a tempo/meter header so a guitar accompaniment can be rendered. "
    "Output only the lullaby — no preamble.\n"
    "Weave the child's loves and fears into the imagery NATURALLY and "
    "sensibly: a loved thing should appear doing what it really does (a dog "
    "curls up beside them, keeps them company, wags its tail), never forced "
    "into a metaphor that makes no sense (NOT 'dogs are your blanket'). "
    "Comfort away fears gently. Every line must make literal sense and read "
    "like a real, soothing lullaby."
)

# A canned lullaby used when LULLABY_SKIP_LLM is set — so you can iterate on
# the audio pipeline without waiting for inference each time.
DEFAULT_LULLABY = """\
Tempo: 60bpm, 6/8
Progression: C - Am - F - G

[C] Little one, little [Am] one,
[F] close your eyes for [G] me,
[C] the moon is on the [Am] water,
[F] the world is fast a-[G]sleep...

[C] Stars are softly [Am] shining,
[F] dreams are on their [G] way,
[C] tomorrow is to-[Am]morrow,
[F] but tonight is [G] today..."""

if SKIP_LLM:
    print("LULLABY_SKIP_LLM set — using canned lyrics, skipping model load.")
    llm = None
else:
    print(f"Loading {MODEL_REPO}...")
    try:
        # Thread count is the biggest CPU lever for llama.cpp — token gen
        # scales ~linearly with physical cores up to the memory-bandwidth
        # ceiling. The old n_threads=2 left most of the box idle. Default to
        # all cores; override with LULLABY_N_THREADS (on hyperthreaded CPUs,
        # physical-core count sometimes beats logical-core count).
        _n_threads = int(os.environ.get("LULLABY_N_THREADS", os.cpu_count() or 4))
        _llama_kwargs = dict(
            repo_id=MODEL_REPO,
            filename="*Q4_K_M.gguf",
            n_ctx=1024,
            n_threads=_n_threads,
            n_threads_batch=_n_threads,   # also parallelize prompt ingestion
            n_batch=512,                  # whole prompt in one eval batch
            n_gpu_layers=0,               # CPU only
            chat_format="llama-3",
            verbose=False,
        )
        try:
            # flash_attn trims prompt-eval + decode where the build supports
            # it; older llama-cpp-python rejects the kwarg with TypeError, so
            # fall back to a load without it rather than crashing.
            llm = Llama.from_pretrained(flash_attn=True, **_llama_kwargs)
        except TypeError:
            print("[llm] flash_attn unsupported on this llama-cpp-python; "
                  "loading without it")
            llm = Llama.from_pretrained(**_llama_kwargs)
        print(f"Model loaded (n_threads={_n_threads}).")
    except Exception as e:
        # If the GGUF is missing or won't load, don't take the whole app
        # down with a cryptic import-time crash. Leave llm=None; the UI still
        # loads and make_lullaby surfaces a friendly message.
        print(f"WARNING: could not load model at {MODEL_REPO}: {e}")
        llm = None


# ---------------------------------------------------------------------------
# Pre-warm everything else at app startup, so the first user click is fast.
#
# Honest tradeoff: pre-warming makes COLD STARTUP slower (the app takes
# longer to come online) in exchange for the FIRST USER GENERATION being
# much faster — which is the moment that actually matters for a judge who
# clicks the link and waits to see something happen.
#
# What we pre-warm and why:
#   - Llama lyric model: already eager-loaded above.
#   - Kokoro TTS: ~340 MB voice model + downloads. ~5-10s saved on first
#     generation.
#   - MiniCPM-V vision: ~3 GB download (first run ever) + model load
#     (~10-30s on CPU). Loads weights into RAM at startup so the first user
#     click skips the load. Saves ~20-30s on first user click.
#
# Both pre-warms are wrapped in try/except so a failure during warmup
# doesn't crash the app — the lazy fallback paths inside each module
# still work if warmup fails for any reason.
# ---------------------------------------------------------------------------

if not SKIP_LLM:
    # Warm the lyric model with one tiny throwaway generation. Weights are
    # already resident (loaded at import), but the FIRST create_chat_completion
    # still pays a one-off prompt-eval/graph warmup; doing it here moves that
    # cost off the first real user click.
    if llm is not None:
        try:
            print("Pre-warming lyric model...")
            llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": "Write a lullaby for: test, age 3"},
                ],
                max_tokens=1,
                temperature=0.0,
            )
            print("Lyric model ready.")
        except Exception as e:
            print(f"WARNING: lyric model pre-warm failed (non-fatal): {e}")

    try:
        print("Pre-warming Kokoro TTS voice...")
        # warmup() (not just _load_kokoro()) so the af_nicole.pt voice tensor
        # downloads and the first-inference path compiles at startup — that
        # was happening on the user's first click otherwise.
        from synths.voice import warmup as _warmup_voice
        _warmup_voice()
        print("Kokoro ready.")
    except Exception as e:
        # Don't take the app down — Kokoro will lazy-load on first call.
        print(f"WARNING: Kokoro pre-warm failed (will retry lazily): {e}")

    try:
        print("Pre-warming MiniCPM-V vision (CPU load only)...")
        # Triggers _try_load() inside vision.py, which downloads + loads the
        # model into RAM so the first real describe() call doesn't pay for it.
        import draw.vision as _vision_warmup
        _vision_warmup._try_load()
        print("Vision ready.")
    except Exception as e:
        # Vision is optional — strokes fallback covers the empty-vision case.
        print(f"WARNING: Vision pre-warm failed (will fall back to strokes): {e}")


# ---------- music theory helpers ----------

NOTE_TO_SEMI = {
    "C": 0, "C#": 1, "Db": 1, "D": 2, "D#": 3, "Eb": 3, "E": 4, "F": 5,
    "F#": 6, "Gb": 6, "G": 7, "G#": 8, "Ab": 8, "A": 9, "A#": 10, "Bb": 10, "B": 11,
}


def chord_root_midi(name, octave=5):
    root = name.strip()
    if len(root) >= 2 and root[:2] in NOTE_TO_SEMI:
        s = NOTE_TO_SEMI[root[:2]]
    elif root[:1] in NOTE_TO_SEMI:
        s = NOTE_TO_SEMI[root[:1]]
    else:
        s = 0
    return 12 * (octave + 1) + s


def chord_tones(name, octave=5):
    root = chord_root_midi(name, octave)
    is_minor = "m" in name and "maj" not in name
    third = root + (3 if is_minor else 4)
    fifth = root + 7
    return [root, third, fifth]


def chord_color_tones(name, octave=5):
    """
    Chord tones PLUS nearby scale degrees that work melodically over the
    chord — gives the picker a wider palette than just root/3/5.

    NOTE: This is the legacy non-key-aware version. It stacks intervals from
    each chord's root and can produce notes outside the song's key. Prefer
    scale_color_tones() when the song's key is known — it constrains the
    palette to the diatonic scale and produces musical, not random, choices.

    Returns a list of MIDI numbers, roughly weighted toward stable tones
    first (they appear earlier in the list).
    """
    root = chord_root_midi(name, octave)
    is_minor = "m" in name and "maj" not in name
    # Stable chord tones
    third = root + (3 if is_minor else 4)
    fifth = root + 7
    # Scale-friendly approaches: 2nd, 6th, and a passing tone above the fifth
    second = root + 2
    sixth = root + (8 if is_minor else 9)
    seventh = root + (10 if is_minor else 11)   # b7 in minor, maj7 in major
    octave_up = root + 12
    return [root, third, fifth, second, sixth, octave_up, seventh]


# ---------- diatonic scale helpers ----------
#
# A riff or melody that "sounds wrong" usually has notes that are in some
# scale but not the SONG's scale. Stacking intervals (root + 2, root + 5,
# root + 9) on each chord independently produces notes that vary chord-by-
# chord and frequently leave the key. Walking by raw semitones (±1, ±2)
# similarly ignores key — a "2-semitone step" up from B in C major lands
# on C#, which doesn't belong.
#
# The fix: derive the song's diatonic scale once, then constrain all
# melodic picks (chord-tone selection, color-tone palette, riff runs) to
# that scale. Walking "by scale step" means moving to the adjacent member
# of the scale, not adding a fixed semitone count.


MAJOR_INTERVALS = (0, 2, 4, 5, 7, 9, 11)
NATURAL_MINOR_INTERVALS = (0, 2, 3, 5, 7, 8, 10)


def parse_key(key_string):
    """
    Parse a key string like 'C major', 'A minor', 'F# major'. Returns
    (tonic_pitch_class, mode) where mode is 'major' or 'minor'. Defaults
    to ('C', 'major') if unparseable.
    """
    if not key_string:
        return 0, "major"
    s = key_string.strip()
    parts = s.split()
    if not parts:
        return 0, "major"
    root_name = parts[0]
    mode = "minor" if len(parts) > 1 and parts[1].lower().startswith("min") else "major"
    if len(root_name) >= 2 and root_name[:2] in NOTE_TO_SEMI:
        tonic_pc = NOTE_TO_SEMI[root_name[:2]]
    elif root_name[:1] in NOTE_TO_SEMI:
        tonic_pc = NOTE_TO_SEMI[root_name[:1]]
    else:
        tonic_pc = 0
    return tonic_pc, mode


def scale_pcs(tonic_pc, mode):
    """Set of 7 pitch classes (0-11) belonging to a diatonic scale."""
    intervals = NATURAL_MINOR_INTERVALS if mode == "minor" else MAJOR_INTERVALS
    return {(tonic_pc + step) % 12 for step in intervals}


def scale_pitches_in_range(tonic_pc, mode, low, high):
    """
    Sorted list of MIDI pitches in [low, high] that belong to the scale.
    Used as the index space for "scale-step" motion in riffs.
    """
    pcs = scale_pcs(tonic_pc, mode)
    return [m for m in range(low, high + 1) if (m % 12) in pcs]


def snap_to_scale(midi, scale_pitches, prefer_up=False):
    """
    If `midi` is not in `scale_pitches`, snap to the nearest scale pitch.
    Ties broken by `prefer_up` direction.
    """
    if not scale_pitches:
        return midi
    if midi in scale_pitches:
        return midi
    # Find insertion point
    nearest = min(scale_pitches, key=lambda x: (abs(x - midi),
                                                   -1 if (prefer_up and x > midi) else 0))
    return nearest


def scale_color_tones(chord_name, key_string, octave=5):
    """
    Return a palette of MIDI notes for melodic use over `chord_name` in the
    given `key_string`. Palette = chord tones + scale tones near them,
    constrained to the diatonic scale of the key. Always musical.

    Ordering biases the result toward stable tones (chord tones first).
    """
    tonic_pc, mode = parse_key(key_string)
    pcs = scale_pcs(tonic_pc, mode)

    # Chord tones (root, 3rd, 5th of the chord itself)
    chord_pitches = chord_tones(chord_name, octave=octave)

    # Add the scale-tone neighbours around each chord tone (±1, ±2 semitones)
    # but ONLY if they're in the key's scale.
    palette = list(chord_pitches)
    for cp in chord_pitches:
        for step in (-2, -1, 1, 2, 7, 12):
            n = cp + step
            if (n % 12) in pcs and n not in palette:
                palette.append(n)

    return palette


def midi_to_name(midi):
    names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
    octave = midi // 12 - 1
    return f"{names[midi % 12]}{octave}"


# ---------- LLM ----------

def build_prompt(name, age, loves, fears, mood, key, meter):
    lines = [f"Write a lullaby for: {name}, age {age}"]
    if loves.strip():
        lines.append(f"Loves: {loves} (weave these in naturally — show them "
                     f"doing what they really do, don't force them into odd "
                     f"metaphors)")
    if fears.strip():
        lines.append(f"Fears: {fears} (gently soothe these away)")
    lines.append(f"Mood: {mood}")
    lines.append(f"Key: {key}")
    lines.append(f"Meter: {meter}")
    return "\n".join(lines)


def generate_lullaby(prompt, temperature=0.75, max_tokens=256):
    if SKIP_LLM or llm is None:
        lola_trace.stage("lyric",
            model="(skipped — using canned default lyric)",
            system_prompt=SYSTEM_PROMPT, user_prompt=prompt,
            temperature=temperature, raw_completion=DEFAULT_LULLABY)
        return DEFAULT_LULLABY
    # max_tokens=256 (was 512): a lullaby with its header is ~80-150 tokens,
    # and decode time is linear in tokens PRODUCED — 512 just gave room to
    # ramble. Stop sequences end generation the moment the model starts a
    # second song or adds commentary instead of burning tokens to the cap.
    resp = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=0.9,
        top_k=40,
        repeat_penalty=1.1,
        stop=["<|eot_id|>", "\n\n\n", "Note:", "Here is", "I hope"],
    )
    completion = resp["choices"][0]["message"]["content"].strip()
    lola_trace.stage("lyric",
        model="lolaby-llama-3b (fine-tuned, Q4_K_M, llama.cpp)",
        system_prompt=SYSTEM_PROMPT,
        user_prompt=prompt,
        temperature=temperature,
        top_p=0.9,
        max_tokens=max_tokens,
        raw_completion=completion)
    return completion


# ---------- parser ----------

TEMPO_RE = re.compile(r"Tempo:\s*(\d+)\s*bpm,\s*(\d+)/(\d+)", re.IGNORECASE)
PROG_RE = re.compile(r"Progression:\s*([A-G][^\n]+)", re.IGNORECASE)
CHORD_RE = re.compile(r"\[([A-G][^\]]*)\]\s*([^\[\n]*)")


def _count_lyric_lines(text):
    """Count chord-marked lyric lines in raw generated text (before capping)."""
    return sum(1 for ln in text.splitlines() if ln.strip().startswith("["))


def parse_lullaby(text):
    tempo_match = TEMPO_RE.search(text)
    if not tempo_match:
        raise ValueError("No tempo line found")
    prog_match = PROG_RE.search(text)
    if not prog_match:
        raise ValueError("No progression line found")
    progression = [c.strip() for c in re.split(r"\s*-\s*", prog_match.group(1)) if c.strip()]

    lines = []
    for raw in text.splitlines():
        raw = raw.strip()
        if not raw.startswith("["):
            continue
        fragments = CHORD_RE.findall(raw)
        if fragments:
            lines.append(fragments)
    if not lines:
        raise ValueError("No chord-marked lyric lines")

    # Render the FULL song — the audio scales to however many lines the model
    # writes (body_seconds = n_lines * chord_seconds in render_rhythm), so a
    # longer lullaby just produces a longer, complete song. We keep only a
    # high safety ceiling to guard against a pathological runaway generation
    # (far above any real lullaby); normal long songs pass through untouched.
    SAFETY_MAX_LINES = 24
    if len(lines) > SAFETY_MAX_LINES:
        keep = (SAFETY_MAX_LINES // 4) * 4  # trim to a whole verse
        lines = lines[:keep]

    plain_lyrics = "\n".join(
        " ".join(frag[1].strip() for frag in line if frag[1].strip())
        for line in lines
    )

    return {
        "bpm": int(tempo_match.group(1)),
        "progression": progression,
        "lines": lines,
        "plain_lyrics": plain_lyrics,
    }


# ---------- arrangement variety ----------
#
# Each generation picks a random "arrangement plan": which intro to use, which
# rhythm pattern, chord rate, and outro style. This is what keeps successive
# lullabies from sounding identical even when the LLM writes very different
# lyrics. Three knobs of variety:
#
#   - macro structure (does it have an intro? a held final chord?)
#   - rhythmic feel   (block chords vs strum vs arpeggio)
#   - chord rate      (subtly different pacing)

GUITAR_PATTERNS = ["block", "down_up", "bass_strum", "arpeggio"]
PIANO_PATTERNS  = ["block", "bass_chord", "arpeggio"]
# Music box patterns. Real music boxes do TWO things:
#   - "block": all tines pluck simultaneously (the pins on the cylinder
#     are aligned vertically for the chord)
#   - "roll":  pins are arranged in a diagonal so notes pluck sequentially
#     as the cylinder rotates — this is the iconic music-box arpeggio
MUSICBOX_PATTERNS = ["block", "roll", "alternating", "broken"]
INTRO_STYLES    = ["none", "single_chord", "two_chord_arp", "melody_pickup"]
OUTRO_STYLES    = ["clean", "held_final", "descending_arp"]
CHORD_RATES     = [3.0, 3.6, 4.2]
MELODIC_SHAPES  = ["calm", "walking", "embellished", "riff"]
# Dropped "sparse" — its "every-other-chord" pattern produced melody textures
# that felt too minimal. The remaining four all keep the melody active across
# the song. Riff still has rest-then-burst within itself (that's its identity),
# but never silent for long.


def make_arrangement_plan(rhythm_instrument, melody_instrument, has_melody,
                            key="C major"):
    """
    Random arrangement choices for one generation. Same lyrics + key with
    different plans produce noticeably different songs.

    `key` is stashed in the plan so renderers can constrain melodic picks
    to the diatonic scale of the song.
    """
    plan = {
        "chord_seconds":  random.choice(CHORD_RATES),
        "intro_style":    random.choice(INTRO_STYLES),
        "outro_style":    random.choice(OUTRO_STYLES),
        "melodic_shape":  random.choice(MELODIC_SHAPES),
        "key":            key,
    }
    if rhythm_instrument == "guitar":
        plan["guitar_pattern"] = random.choice(GUITAR_PATTERNS)
    elif rhythm_instrument == "piano":
        plan["piano_pattern"]  = random.choice(PIANO_PATTERNS)
    elif rhythm_instrument == "musicbox":
        plan["musicbox_pattern"] = random.choice(MUSICBOX_PATTERNS)

    # "melody_pickup" intro needs a melody instrument to actually play.
    # If no melody is selected, fall back to something else.
    if plan["intro_style"] == "melody_pickup" and not has_melody:
        plan["intro_style"] = random.choice(["none", "single_chord", "two_chord_arp"])

    return plan


def _intro_seconds(plan):
    """How long the intro section is, in seconds."""
    style = plan["intro_style"]
    cs = plan["chord_seconds"]
    return {
        "none":           0.0,
        "single_chord":   cs * 0.7,
        "two_chord_arp":  cs * 1.4,
        "melody_pickup":  cs * 0.9,
    }[style]


def _outro_seconds(plan):
    """How long the outro section is, after the last vocal line."""
    style = plan["outro_style"]
    cs = plan["chord_seconds"]
    return {
        "clean":           0.0,
        "held_final":      cs * 0.8,
        "descending_arp":  cs * 1.0,
    }[style]


# ---------- rhythm layer ----------

# Chord-quality intervals, matching synths.musicbox.CHORD_INTERVALS so the
# music box's individual-note patterns (alternating, broken) use the same
# voicing the synth's chord() would.
_MB_NOTE = {
    "C": 0, "C#": 1, "Db": 1, "D": 2, "D#": 3, "Eb": 3, "E": 4, "F": 5,
    "F#": 6, "Gb": 6, "G": 7, "G#": 8, "Ab": 8, "A": 9, "A#": 10, "Bb": 10, "B": 11,
}
_MB_INTERVALS = {
    "": [0, 4, 7, 12], "m": [0, 3, 7, 12], "7": [0, 4, 7, 10],
    "m7": [0, 3, 7, 10], "maj7": [0, 4, 7, 11],
    "sus2": [0, 2, 7, 12], "sus4": [0, 5, 7, 12],
}


def _musicbox_chord_midis(chord_name, octave=5):
    """Return the four chord-tone MIDI numbers (root, third, fifth, octave)
    for a chord name, matching the music box synth's voicing."""
    name = chord_name.strip()
    if len(name) >= 2 and name[1] in ("#", "b"):
        root, quality = name[:2], name[2:]
    else:
        root, quality = name[:1], name[1:]
    intervals = _MB_INTERVALS.get(quality, _MB_INTERVALS[""])
    root_midi = 12 * (octave + 1) + _MB_NOTE.get(root, 0)
    return [root_midi + i for i in intervals]


def render_rhythm(parsed, instrument, plan, min_body_seconds=None):
    """
    Render the rhythm/accompaniment track for the given instrument according
    to the arrangement plan.

    Returns (audio, body_seconds, intro_seconds) where:
        body_seconds:  duration of the vocal section
        intro_seconds: where the vocal section starts within the audio
    The total audio length is intro + body + outro.

    min_body_seconds: if given, the body is extended (extra chord cycles
        appended, looping the progression) so the music covers AT LEAST this
        long. Used to guarantee instruments play under a voice that ran longer
        than the n_lines estimate — so long lyrics are never left over silence.
    """
    progression = parsed["progression"]
    n_lines = len(parsed["lines"])
    chord_seconds = plan["chord_seconds"]
    # Extend the number of chord cycles if the voice needs more time than the
    # lyric-line estimate. Each cycle is one chord_seconds slot.
    n_cycles = n_lines
    if min_body_seconds is not None:
        needed = int(np.ceil(min_body_seconds / chord_seconds))
        n_cycles = max(n_cycles, needed)
    n_lines = n_cycles
    body_seconds = n_lines * chord_seconds
    intro_seconds = _intro_seconds(plan)

    if instrument == "guitar":
        synth = GuitarSynth()
        events = _guitar_events(progression, n_lines, plan, intro_seconds)
        return synth.sequence(events, effects=ACOUSTIC_PRESET), body_seconds, intro_seconds

    if instrument == "piano":
        synth = PianoSynth()
        events = _piano_events(progression, n_lines, plan, intro_seconds)
        return synth.sequence(events), body_seconds, intro_seconds

    if instrument == "xylophone":
        synth = XylophoneSynth()
        events = []
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = intro_seconds + i * chord_seconds
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": chord_seconds * 0.5,
                "octave": 4, "volume": 0.7,
                "brightness": 0.5, "arpeggio_ms": 55, "direction": "up",
            })
            events.append({
                "type": "chord", "name": chord, "time": t + chord_seconds * 0.5,
                "duration": chord_seconds * 0.5,
                "octave": 4, "volume": 0.5,
                "brightness": 0.45, "arpeggio_ms": 70, "direction": "down",
            })
        events.extend(_intro_events(progression, plan, instrument))
        events.extend(_outro_events(progression, n_lines, plan, intro_seconds, instrument))
        return synth.sequence(events), body_seconds, intro_seconds

    if instrument == "ocarina":
        synth = WhistleSynth()
        events = []
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = intro_seconds + i * chord_seconds
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": chord_seconds,
                "octave": 4, "volume": 0.65,
            })
        events.extend(_intro_events(progression, plan, instrument))
        events.extend(_outro_events(progression, n_lines, plan, intro_seconds, instrument))
        return synth.sequence(events), body_seconds, intro_seconds

    if instrument == "musicbox":
        synth = MusicBoxSynth()
        pattern = plan.get("musicbox_pattern", "roll")
        events = []
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = intro_seconds + i * chord_seconds
            if pattern == "block":
                # All tines pluck together — chord-aligned cylinder pins.
                events.append({
                    "type": "chord", "name": chord, "time": t,
                    "duration": chord_seconds * 1.3,
                    "octave": 5, "volume": 0.65,
                    "spread_ms": 18, "direction": "up",
                })
            elif pattern == "roll":
                # Sequential cylinder rotation through the chord — the
                # iconic music-box arpeggio sweep.
                events.append({
                    "type": "chord", "name": chord, "time": t,
                    "duration": chord_seconds * 1.4,
                    "octave": 5, "volume": 0.60,
                    "spread_ms": 220, "direction": "up",
                })
            elif pattern == "alternating":
                # Waltz-like "oom-pah-pah": the bass tine alone on the beat,
                # then the upper chord tones rolled in the second half of the
                # bar. Gentle and rocking — very music-box / lullaby.
                midis = _musicbox_chord_midis(chord, octave=5)
                # Bass tine on the downbeat (one octave below the chord root)
                events.append({
                    "type": "note", "name": midi_to_name(midis[0] - 12),
                    "time": t, "duration": chord_seconds * 0.5,
                    "volume": 0.55,
                })
                # Upper tones rolled in, slightly later, softer
                events.append({
                    "type": "chord", "name": chord,
                    "time": t + chord_seconds * 0.5,
                    "duration": chord_seconds * 0.9,
                    "octave": 5, "volume": 0.52,
                    "spread_ms": 90, "direction": "up",
                })
            else:  # "broken"
                # The chord split into a gentle rocking back-and-forth figure
                # across the bar (root, fifth, third, top, third, fifth) —
                # like cylinder pins arranged for a lilting pattern.
                midis = _musicbox_chord_midis(chord, octave=5)
                root, third, fifth, top = midis
                figure = [root, fifth, third, top, third, fifth]
                step = (chord_seconds * 1.1) / len(figure)
                for k, m in enumerate(figure):
                    events.append({
                        "type": "note", "name": midi_to_name(m),
                        "time": t + k * step,
                        "duration": step * 2.2,
                        "volume": 0.50 + (0.08 if k == 0 else 0.0),
                    })
        events.extend(_intro_events(progression, plan, instrument))
        events.extend(_outro_events(progression, n_lines, plan, intro_seconds, instrument))
        return synth.sequence(events), body_seconds, intro_seconds

    if instrument == "harp":
        # Harp rhythm: slow rolled chords. The harp synth handles the
        # arpeggio internally via arpeggio_ms.
        synth = HarpSynth()
        events = []
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = intro_seconds + i * chord_seconds
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": chord_seconds * 1.3,
                "octave": 3, "volume": 0.65, "brightness": 0.4,
                "arpeggio_ms": 130, "direction": "up",
            })
        events.extend(_intro_events(progression, plan, instrument))
        events.extend(_outro_events(progression, n_lines, plan, intro_seconds, instrument))
        return synth.sequence(events), body_seconds, intro_seconds

    raise ValueError(f"Unknown rhythm instrument: {instrument}")


# ---------- guitar rhythm patterns ----------

def _guitar_events(progression, n_lines, plan, t_offset):
    """
    Build the guitar event list for the chosen pattern. t_offset shifts
    everything by the intro length.
    """
    pattern = plan["guitar_pattern"]
    cs = plan["chord_seconds"]
    events = []

    if pattern == "block":
        # Single mellow down-strum per chord — ballad feel
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 1.4,
                "direction": "down", "spread_ms": 28,
                "volume": 0.62, "brightness": 0.34,
            })

    elif pattern == "down_up":
        # Down-up-down-up pattern, 4 strokes per chord
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 1.4,
                "direction": "down", "spread_ms": 22,
                "volume": 0.55, "brightness": 0.34,
            })
            events.append({
                "type": "chord", "name": chord, "time": t + cs * 0.5,
                "duration": cs * 0.9,
                "direction": "up", "spread_ms": 16,
                "volume": 0.35, "brightness": 0.36,
            })

    elif pattern == "bass_strum":
        # "Boom-chuck" folk pattern: bass note, then strum, repeated
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            # Bass note on beat 1 — render as single-note chord
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 0.4,
                "direction": "down", "spread_ms": 8,
                "volume": 0.55, "brightness": 0.32, "octave": 2,
            })
            # Light strum on beat 2
            events.append({
                "type": "chord", "name": chord, "time": t + cs * 0.4,
                "duration": cs * 0.6,
                "direction": "down", "spread_ms": 26,
                "volume": 0.42, "brightness": 0.36,
            })
            # Bass on beat 3
            events.append({
                "type": "chord", "name": chord, "time": t + cs * 0.55,
                "duration": cs * 0.35,
                "direction": "down", "spread_ms": 8,
                "volume": 0.40, "brightness": 0.32, "octave": 2,
            })

    elif pattern == "arpeggio":
        # Slow rolling arpeggio across each chord — fingerpicked feel
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 1.5,
                "direction": "down", "spread_ms": 220,   # long roll = arpeggio
                "volume": 0.55, "brightness": 0.36,
            })

    # Append intro and outro
    events.extend(_intro_events(progression, plan, "guitar"))
    events.extend(_outro_events(progression, n_lines, plan, t_offset, "guitar"))
    return events


# ---------- piano rhythm patterns ----------

def _piano_events(progression, n_lines, plan, t_offset):
    """Build the piano event list for the chosen pattern."""
    pattern = plan["piano_pattern"]
    cs = plan["chord_seconds"]
    events = []

    if pattern == "block":
        # Single chord strike per chord — current behavior
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 1.3,
                "octave": 3, "direction": "down",
                "spread_ms": 35, "volume": 0.65, "brightness": 0.4,
            })

    elif pattern == "bass_chord":
        # Bass note on beat 1, chord on beat 3 — gentle half-time feel
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 0.5,
                "octave": 2, "direction": "down",
                "spread_ms": 10, "volume": 0.60, "brightness": 0.38,
            })
            events.append({
                "type": "chord", "name": chord, "time": t + cs * 0.5,
                "duration": cs * 0.9,
                "octave": 3, "direction": "down",
                "spread_ms": 30, "volume": 0.55, "brightness": 0.42,
            })

    elif pattern == "arpeggio":
        # Slow rolling arpeggio — broken-chord ballad
        for i in range(n_lines):
            chord = progression[i % len(progression)]
            t = t_offset + i * cs
            events.append({
                "type": "chord", "name": chord, "time": t,
                "duration": cs * 1.4,
                "octave": 3, "direction": "up",
                "spread_ms": 280,   # long roll
                "volume": 0.60, "brightness": 0.42,
            })

    events.extend(_intro_events(progression, plan, "piano"))
    events.extend(_outro_events(progression, n_lines, plan, t_offset, "piano"))
    return events


# ---------- intro / outro event builders ----------

def _intro_events(progression, plan, instrument):
    """
    Events that go BEFORE the vocal section (t < intro_seconds).
    Builds the intro figure for the rhythm instrument.
    """
    style = plan["intro_style"]
    if style == "none":
        return []
    cs = plan["chord_seconds"]
    first_chord = progression[0]

    if style == "single_chord":
        # One soft chord rings out before voice enters
        if instrument == "guitar":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.7,
                "direction": "down", "spread_ms": 30,
                "volume": 0.45, "brightness": 0.32,
            }]
        if instrument == "piano":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.7,
                "octave": 3, "direction": "down",
                "spread_ms": 30, "volume": 0.55, "brightness": 0.38,
            }]
        if instrument == "xylophone":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.6, "octave": 4,
                "volume": 0.55, "brightness": 0.5,
                "arpeggio_ms": 70, "direction": "up",
            }]
        if instrument == "ocarina":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.6, "octave": 4, "volume": 0.55,
            }]
        if instrument == "musicbox":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.7, "octave": 5,
                "volume": 0.55, "spread_ms": 18, "direction": "up",
            }]
        if instrument == "harp":
            return [{
                "type": "chord", "name": first_chord, "time": 0.0,
                "duration": cs * 0.8, "octave": 3, "volume": 0.55,
                "brightness": 0.4, "arpeggio_ms": 130, "direction": "up",
            }]

    if style == "two_chord_arp":
        # Two chords of arpeggiation — a pickup figure
        second_chord = progression[1 % len(progression)]
        if instrument == "guitar":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.7, "direction": "down", "spread_ms": 180,
                 "volume": 0.45, "brightness": 0.34},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.7, "direction": "down", "spread_ms": 200,
                 "volume": 0.45, "brightness": 0.34},
            ]
        if instrument == "piano":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.7, "octave": 3, "direction": "up",
                 "spread_ms": 220, "volume": 0.55, "brightness": 0.4},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.7, "octave": 3, "direction": "up",
                 "spread_ms": 220, "volume": 0.55, "brightness": 0.4},
            ]
        if instrument == "xylophone":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.5, "octave": 4, "volume": 0.55,
                 "brightness": 0.5, "arpeggio_ms": 80, "direction": "up"},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.5, "octave": 4, "volume": 0.55,
                 "brightness": 0.5, "arpeggio_ms": 80, "direction": "up"},
            ]
        if instrument == "ocarina":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.7, "octave": 4, "volume": 0.55},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.7, "octave": 4, "volume": 0.55},
            ]
        if instrument == "musicbox":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.6, "octave": 5, "volume": 0.55,
                 "spread_ms": 220, "direction": "up"},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.6, "octave": 5, "volume": 0.55,
                 "spread_ms": 220, "direction": "up"},
            ]
        if instrument == "harp":
            return [
                {"type": "chord", "name": first_chord, "time": 0.0,
                 "duration": cs * 0.7, "octave": 3, "volume": 0.55,
                 "brightness": 0.4, "arpeggio_ms": 160, "direction": "up"},
                {"type": "chord", "name": second_chord, "time": cs * 0.7,
                 "duration": cs * 0.7, "octave": 3, "volume": 0.55,
                 "brightness": 0.4, "arpeggio_ms": 160, "direction": "up"},
            ]

    if style == "melody_pickup":
        # The rhythm instrument plays a single soft chord. The actual melody
        # pickup riff is added separately by render_melody when this style
        # is active. So the rhythm intro is just a single soft chord.
        if instrument == "guitar":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "direction": "down", "spread_ms": 30,
                "volume": 0.40, "brightness": 0.32,
            }]
        if instrument == "piano":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "octave": 3, "direction": "down",
                "spread_ms": 30, "volume": 0.50, "brightness": 0.38,
            }]
        if instrument == "xylophone":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "octave": 4,
                "volume": 0.50, "brightness": 0.5,
                "arpeggio_ms": 70, "direction": "up",
            }]
        if instrument == "ocarina":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "octave": 4, "volume": 0.55,
            }]
        if instrument == "musicbox":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "octave": 5, "volume": 0.50,
                "spread_ms": 18, "direction": "up",
            }]
        if instrument == "harp":
            return [{
                "type": "chord", "name": first_chord, "time": cs * 0.5,
                "duration": cs * 0.5, "octave": 3, "volume": 0.50,
                "brightness": 0.4, "arpeggio_ms": 130, "direction": "up",
            }]

    return []


def _outro_events(progression, n_lines, plan, t_offset, instrument):
    """
    Events that go AFTER the vocal section. They ring out past the body.
    """
    style = plan["outro_style"]
    if style == "clean":
        return []

    cs = plan["chord_seconds"]
    last_chord = progression[(n_lines - 1) % len(progression)]
    end_t = t_offset + n_lines * cs

    if style == "held_final":
        # One held strike of the final chord, lets it ring
        if instrument == "guitar":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5,
                "direction": "down", "spread_ms": 32,
                "volume": 0.50, "brightness": 0.32,
            }]
        if instrument == "piano":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5,
                "octave": 3, "direction": "down",
                "spread_ms": 40, "volume": 0.55, "brightness": 0.38,
            }]
        if instrument == "xylophone":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.0, "octave": 4,
                "volume": 0.55, "brightness": 0.5,
                "arpeggio_ms": 80, "direction": "up",
            }]
        if instrument == "ocarina":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.0, "octave": 4, "volume": 0.55,
            }]
        if instrument == "musicbox":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.4, "octave": 5,
                "volume": 0.55, "spread_ms": 20, "direction": "up",
            }]
        if instrument == "harp":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5, "octave": 3, "volume": 0.55,
                "brightness": 0.4, "arpeggio_ms": 130, "direction": "up",
            }]

    if style == "descending_arp":
        # Slow downward arpeggio on the last chord — resolution figure
        if instrument == "guitar":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5,
                "direction": "up", "spread_ms": 280,
                "volume": 0.50, "brightness": 0.34,
            }]
        if instrument == "piano":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5,
                "octave": 3, "direction": "down",
                "spread_ms": 280, "volume": 0.55, "brightness": 0.40,
            }]
        if instrument == "xylophone":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.0, "octave": 4,
                "volume": 0.50, "brightness": 0.45,
                "arpeggio_ms": 110, "direction": "down",
            }]
        if instrument == "ocarina":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.0, "octave": 4, "volume": 0.50,
            }]
        if instrument == "musicbox":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.4, "octave": 5,
                "volume": 0.55, "spread_ms": 260, "direction": "down",
            }]
        if instrument == "harp":
            return [{
                "type": "chord", "name": last_chord, "time": end_t,
                "duration": cs * 1.5, "octave": 3, "volume": 0.55,
                "brightness": 0.4, "arpeggio_ms": 220, "direction": "down",
            }]

    return []



# ---------- melody layer ----------

def design_melody_pitches(progression, n_lines, low=67, high=79,
                            key="C major"):
    """
    Pick MIDI pitches per line — one pitch per chord. Stochastic, but
    constrained to the diatonic scale of the song's key so the melody
    never lands on out-of-key notes.

    Each chord gets a randomly-chosen pitch from a scale-aware palette
    (chord tones plus nearby scale tones), weighted toward stable tones.
    The previous pitch is used as the target for nearest-neighbour
    selection, so consecutive picks tend to be stepwise within the scale.

    Resolves to the tonic of the last chord at the end.
    """
    tonic_pc, mode = parse_key(key)
    scale_pitches = scale_pitches_in_range(tonic_pc, mode, low, high)

    pitches = []
    prev = None
    for i in range(n_lines):
        chord = progression[i % len(progression)]
        palette = scale_color_tones(chord, key, octave=5)

        # Weight: chord tones (first 3) heavy, color tones moderate
        weights = []
        for k, _ in enumerate(palette):
            if k < 3:
                weights.append(3)   # chord tones
            else:
                weights.append(1)   # scale color tones
        chosen_tone = random.choices(palette, weights=weights, k=1)[0]

        # Bring into [low, high]; pick the octave closest to prev (or middle)
        target = prev if prev is not None else (low + high) // 2
        candidates = []
        for shift in (-24, -12, 0, 12, 24):
            n = chosen_tone + shift
            if low <= n <= high:
                candidates.append(n)
        if not candidates:
            n = chosen_tone
            while n < low: n += 12
            while n > high: n -= 12
            candidates = [n]

        # Small chance of a deliberate leap for melodic interest, but the
        # leap target must still be in the scale.
        if prev is not None and random.random() < 0.18:
            pick = random.choice(candidates)
        else:
            pick = min(candidates, key=lambda x: abs(x - target))

        # Final safety: snap to the diatonic scale of the key
        if scale_pitches and pick not in scale_pitches:
            pick = snap_to_scale(pick, scale_pitches)
        pitches.append(pick)
        prev = pick

    # Final note: tonic of the song's key (not just the last chord), so the
    # phrase resolves harmonically. Bring into range.
    if pitches:
        final = 12 * 6 + tonic_pc   # C6-ish
        while final < low:
            final += 12
        while final > high:
            final -= 12
        pitches[-1] = final
    return pitches


def _scale_step(midi, step, scale_pitches):
    """
    Move `step` scale degrees from `midi` along `scale_pitches`. If `midi`
    isn't in the scale, snap to the nearest scale pitch first.

    Positive step = up, negative = down. Returns a pitch in scale_pitches,
    or `midi` unchanged if the move would leave the range.
    """
    if not scale_pitches:
        return midi
    if midi not in scale_pitches:
        midi = snap_to_scale(midi, scale_pitches)
    try:
        idx = scale_pitches.index(midi)
    except ValueError:
        return midi
    new_idx = idx + step
    if new_idx < 0 or new_idx >= len(scale_pitches):
        # Reflect — go the other way instead of escaping range
        new_idx = idx - step
    new_idx = max(0, min(len(scale_pitches) - 1, new_idx))
    return scale_pitches[new_idx]


def _make_scale_run(start_midi, n_notes, total_seconds, scale_pitches,
                     direction="random"):
    """
    Build a riff run that walks BY SCALE STEPS, not raw semitones — so every
    note belongs to the song's key. Returns list of (midi, time_in_window,
    duration) tuples.

    Pattern is biased toward stepwise motion (1 scale step at a time) with
    occasional thirds (2 scale steps) for melodic interest.
    """
    if direction == "random":
        direction = random.choice(["up", "down", "arc_up", "arc_down", "wave"])

    # Snap start to scale
    if scale_pitches:
        if start_midi not in scale_pitches:
            start_midi = snap_to_scale(start_midi, scale_pitches)

    notes = [start_midi]
    for k in range(1, n_notes):
        prev = notes[-1]
        # 1 scale step is the most common; 2 (a third in the scale) for color
        magnitude = random.choices([1, 2], weights=[4, 1], k=1)[0]
        if direction == "up":
            step = +magnitude
        elif direction == "down":
            step = -magnitude
        elif direction == "arc_up":
            step = +magnitude if k <= n_notes // 2 else -magnitude
        elif direction == "arc_down":
            step = -magnitude if k <= n_notes // 2 else +magnitude
        else:  # wave
            step = random.choice([-1, +1, -2, +2]) * 1
            step = step if abs(step) == magnitude else (step // abs(step)) * magnitude

        nxt = _scale_step(prev, step, scale_pitches)
        if nxt == prev:
            # Would have left range — flip direction
            nxt = _scale_step(prev, -step, scale_pitches)
        notes.append(nxt)

    step_t = total_seconds / n_notes
    out = []
    for k, m in enumerate(notes):
        out.append((m, k * step_t, step_t * 1.15))
    return out


def design_melody_events(pitches, chord_seconds, shape, low, high,
                          key="C major"):
    """
    Convert a list of one-pitch-per-chord into a list of timed melody events
    according to the chosen rhythmic shape.

    Returns list of dicts: {time_offset, midi, duration_s, volume}
    where time_offset is in seconds relative to the start of the body
    (caller adds start_delay).

    All melodic choices are constrained to the diatonic scale of `key` —
    riffs walk by scale steps, neighbour tones snap to the scale,
    ornaments use scale steps not raw semitones.

    Shapes:
      calm:        one held note per chord, slow
      walking:     two notes per chord, stepwise (in-scale)
      sparse:      one note every other chord, longer holds
      embellished: main note + 1-2 ornament notes per chord (occasional run)
      riff:        bursts of fast scale-walking notes on selected chords +
                   rest or held tone on others. The riff goes IN-CHORD —
                   notes pack into the chord's window, not at chord boundaries.
    """
    events = []
    n = len(pitches)
    if n == 0:
        return events

    tonic_pc, mode = parse_key(key)
    scale_pitches = scale_pitches_in_range(tonic_pc, mode, low, high)

    # ---- RIFF shape ----
    # Bursts of fast scale-walking notes on selected chords. Other chords
    # get a held tone or rest. This breaks the chord-grid feel because
    # the burst sits INSIDE a single chord, several notes-per-second.
    if shape == "riff":
        n_runs = 3 if n >= 6 else 2
        eligible = list(range(1, n - 1))   # avoid first chord (let it breathe)
        random.shuffle(eligible)
        run_indices = set(eligible[:n_runs])
        # Non-run chords get a held tone with HIGH probability — used to be
        # 55% which produced 8-11 second silent gaps on bad rolls. At 85%,
        # almost every non-run chord gets at least one held note, so the
        # melody never disappears for too long. Riff bursts still feel
        # distinctive because they have ~6 notes vs. the held tone's 1.
        held_indices = set()
        for k in range(n - 1):
            if k not in run_indices and random.random() < 0.85:
                held_indices.add(k)

        for i, midi in enumerate(pitches):
            chord_t = i * chord_seconds
            is_last = (i == n - 1)

            if is_last:
                # Final chord: resolve with held tonic
                events.append({
                    "time": chord_t + chord_seconds * 0.10,
                    "midi": midi,
                    "duration": chord_seconds * 1.6,
                    "volume": 0.85,
                })
                continue

            if i in run_indices:
                n_notes = random.choice([4, 5, 6, 7])
                run_duration = chord_seconds * random.uniform(0.55, 0.75)
                run_offset = chord_t + chord_seconds * random.uniform(0.15, 0.30)
                run_notes = _make_scale_run(midi, n_notes, run_duration,
                                              scale_pitches)
                for k, (m, dt, dur) in enumerate(run_notes):
                    events.append({
                        "time": run_offset + dt,
                        "midi": m,
                        "duration": dur,
                        "volume": 0.80 if k == 0 else 0.65,
                    })
            elif i in held_indices:
                events.append({
                    "time": chord_t + chord_seconds * 0.10,
                    "midi": midi,
                    "duration": chord_seconds * 0.7,
                    "volume": 0.75,
                })
            # else: rest (silence on this chord)
        return events

    # ---- other shapes ----
    for i, midi in enumerate(pitches):
        is_last = (i == n - 1)
        chord_t = i * chord_seconds

        if shape == "calm":
            offset = chord_t + chord_seconds * 0.05
            dur = chord_seconds * (1.6 if is_last else 0.8)
            events.append({"time": offset, "midi": midi, "duration": dur, "volume": 0.85})

        elif shape == "walking":
            offset1 = chord_t + chord_seconds * 0.05
            dur1 = chord_seconds * (1.6 if is_last else 0.40)
            events.append({"time": offset1, "midi": midi, "duration": dur1, "volume": 0.85})
            if not is_last:
                next_midi = pitches[i + 1] if i + 1 < n else midi
                step = 1 if next_midi > midi else -1
                neighbour = _scale_step(midi, step, scale_pitches)
                if neighbour == midi:
                    neighbour = _scale_step(midi, -step, scale_pitches)
                if low <= neighbour <= high and neighbour != midi:
                    events.append({
                        "time": chord_t + chord_seconds * 0.55,
                        "midi": neighbour,
                        "duration": chord_seconds * 0.40,
                        "volume": 0.65,
                    })

        elif shape == "sparse":
            if i % 2 == 0 or is_last:
                offset = chord_t + chord_seconds * 0.10
                dur = chord_seconds * (1.8 if is_last else 1.4)
                events.append({"time": offset, "midi": midi, "duration": dur,
                                "volume": 0.85})

        elif shape == "embellished":
            offset_main = chord_t + chord_seconds * 0.05
            dur_main = chord_seconds * (1.6 if is_last else 0.55)
            events.append({"time": offset_main, "midi": midi,
                            "duration": dur_main, "volume": 0.85})
            if not is_last:
                pattern = random.choice(["grace", "turn", "run", "run", "skip"])
                if pattern == "skip":
                    continue
                if pattern == "grace":
                    step = random.choice([1, -1, 2, -2])
                    ornament = _scale_step(midi, step, scale_pitches)
                    if low <= ornament <= high and ornament != midi:
                        events.append({
                            "time": chord_t + chord_seconds * 0.60,
                            "midi": ornament,
                            "duration": chord_seconds * 0.30,
                            "volume": 0.60,
                        })
                elif pattern == "turn":
                    up = _scale_step(midi, +1, scale_pitches)
                    down = _scale_step(midi, -1, scale_pitches)
                    if low <= up <= high and up != midi:
                        events.append({
                            "time": chord_t + chord_seconds * 0.62,
                            "midi": up,
                            "duration": chord_seconds * 0.15,
                            "volume": 0.55,
                        })
                    if low <= down <= high and down != midi:
                        events.append({
                            "time": chord_t + chord_seconds * 0.78,
                            "midi": down,
                            "duration": chord_seconds * 0.18,
                            "volume": 0.55,
                        })
                elif pattern == "run":
                    # 3-4 quick scale-walking notes in the back half of the chord
                    n_notes = random.choice([3, 4])
                    run_dur = chord_seconds * 0.45
                    run_offset = chord_t + chord_seconds * 0.50
                    run = _make_scale_run(midi, n_notes, run_dur, scale_pitches)
                    for k, (m, dt, dur) in enumerate(run):
                        events.append({
                            "time": run_offset + dt,
                            "midi": m,
                            "duration": dur,
                            "volume": 0.60,
                        })
    return events


def render_melody(parsed, instrument, plan, intro_seconds, natural_offset=1.5):
    """
    Render the melody layer with the chosen instrument.

    intro_seconds: time the intro takes — melody is shifted by this much
    natural_offset: short delay inside the vocal section before the melody
                    enters. 1.5s lets the rhythm + voice register first
                    without leaving a noticeable hole at the start. Used
                    to be 4.0s which produced ~1 chord of dead air at
                    slow tempos.

    If plan["intro_style"] == "melody_pickup", a short pickup riff is added
    BEFORE the vocal section.
    """
    progression = parsed["progression"]
    n_lines = len(parsed["lines"])
    chord_seconds = plan["chord_seconds"]
    start_delay = intro_seconds + natural_offset

    if instrument == "ocarina":
        low, high = 67, 79
    elif instrument == "xylophone":
        low, high = 72, 84   # bright high register
    elif instrument == "harp":
        low, high = 67, 81   # mid-high; harp sustains, so room above ocarina range
    elif instrument == "piano":
        low, high = 67, 79
    elif instrument == "musicbox":
        low, high = 72, 88   # C5-E6 — the natural music-box register
    else:  # guitar melody — E4-E5, sits above the keyboard chord rhythm
        low, high = 64, 76

    # When the same instrument plays both layers (musicbox solo "both" case),
    # the melody needs to sit above the chord voicing or it gets masked by
    # the same-timbre rhythm. Bump the range up by an octave (clamped).
    octave_shift = plan.get("melody_octave_shift", 0)
    if octave_shift:
        low  = min(low  + octave_shift, 100)
        high = min(high + octave_shift, 100)

    key = plan.get("key", "C major")
    pitches = design_melody_pitches(progression, n_lines, low=low, high=high,
                                      key=key)
    shape = plan.get("melodic_shape", "calm")
    melody_events_spec = design_melody_events(pitches, chord_seconds,
                                                shape, low, high, key=key)

    # Optional melody pickup: a 4-note phrase before the vocal section.
    # Vary the pickup contour too — three different shapes.
    pickup_events = []
    if plan["intro_style"] == "melody_pickup" and pitches:
        first_pitch = pitches[0]
        pickup_shape = random.choice(["ascending", "turn", "leap_back"])
        if pickup_shape == "ascending":
            pickup_notes = [first_pitch - 4, first_pitch - 2, first_pitch - 1, first_pitch]
        elif pickup_shape == "turn":
            pickup_notes = [first_pitch - 2, first_pitch + 1, first_pitch - 1, first_pitch]
        else:  # leap_back
            pickup_notes = [first_pitch + 5, first_pitch + 2, first_pitch - 1, first_pitch]
        pickup_notes = [max(low, min(high, p)) for p in pickup_notes]
        pickup_duration = chord_seconds * 0.9 / len(pickup_notes)
        pickup_start = chord_seconds * 0.1
        for k, midi in enumerate(pickup_notes):
            pickup_events.append({
                "type": "note", "name": midi_to_name(midi),
                "time": pickup_start + k * pickup_duration,
                "duration": pickup_duration * 1.3,
                "volume": 0.65,
            })

    if instrument == "ocarina":
        synth = WhistleSynth()
        events = list(pickup_events)
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"],
                "volume": ev["volume"],
            })
        return synth.sequence(events)

    if instrument == "xylophone":
        synth = XylophoneSynth()
        events = list(pickup_events)
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"],
                "volume": ev["volume"] * 1.3,    # xylo plays a bit hotter
                "brightness": 0.55,
            })
        return synth.sequence(events)

    if instrument == "piano":
        synth = PianoSynth()
        events = list(pickup_events)
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"],
                "volume": ev["volume"] * 0.82,
                "brightness": 0.45,
            })
        return synth.sequence(events)

    if instrument == "guitar":
        # Guitar melody: emit each pitch as a SINGLE fingerpicked string,
        # NOT a chord voicing. This used to render full 6-string chord
        # shapes for every melody note, which:
        #   (a) stacked chord voicings on top of the keyboard's chords
        #   (b) added a small 10ms strum sweep per note, smearing timing
        # Now each melody note is a single-string note event — clean,
        # in time, and reads as a melody line over the keyboard.
        synth = GuitarSynth()
        events = list(pickup_events)
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"] * 1.2,   # ring slightly past
                "volume": ev["volume"] * 0.75,
                "brightness": 0.5,
            })
        return synth.sequence(events, effects=ACOUSTIC_PRESET)

    if instrument == "harp":
        synth = HarpSynth()
        events = list(pickup_events)
        # Harp KS already rings ~1.5s — over-extending duration for fast riff
        # notes muddies the run. Use shorter multiplier for riff shape.
        dur_mult = 1.05 if plan.get("melodic_shape") == "riff" else 1.6
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"] * dur_mult,
                "volume": ev["volume"],
                "brightness": 0.4,
            })
        return synth.sequence(events)

    if instrument == "musicbox":
        synth = MusicBoxSynth()
        events = list(pickup_events)
        # Music box tines have their own slow decay (~1-2s); we don't need
        # to extend duration much. For riff shape the tines need to release
        # quickly enough that fast runs don't muddy together.
        dur_mult = 1.0 if plan.get("melodic_shape") == "riff" else 1.3
        for ev in melody_events_spec:
            events.append({
                "type": "note",
                "name": midi_to_name(ev["midi"]),
                "time": start_delay + ev["time"],
                "duration": ev["duration"] * dur_mult,
                "volume": ev["volume"],
            })
        return synth.sequence(events)

    raise ValueError(f"Unknown melody instrument: {instrument}")


# ---------- mixing ----------

def _active_rms(x, thresh_ratio=0.1):
    """RMS over only the 'active' (above-threshold) samples. This tracks
    PERCEIVED loudness far better than peak or full-signal RMS: speech has
    gaps between words, so its full-signal RMS is misleadingly low, while a
    continuous instrument bed has high full-signal RMS. Measuring only the
    active regions compares like-for-like (how loud each is *when sounding*)."""
    a = np.abs(x)
    peak = a.max()
    if peak < 1e-9:
        return 0.0
    active = x[a > peak * thresh_ratio]
    if len(active) == 0:
        return 0.0
    return float(np.sqrt(np.mean(active ** 2) + 1e-12))


def mix_tracks(rhythm, melody, voice, melody_instrument="ocarina",
               rhythm_instrument=None,
               rhythm_gain=0.7, voice_gain=1.4, same_instrument=False):
    """
    Mix the three layers so the VOICE clearly leads and the instruments sit
    underneath as a soft cushion — the lullaby norm.

    The old approach used fixed gain multipliers + a final peak-normalize.
    That made instruments feel too loud because loudness is about average
    energy (RMS), not peak: the voice is intermittent (gaps between words)
    while the instruments play continuously, so even at a lower peak the
    instruments dominated the *perceived* mix, especially in the gaps.

    New approach: measure each layer's ACTIVE RMS (loudness when actually
    sounding) and scale the instrument bed so it sits at a fixed, gentle
    fraction of the voice's loudness. The voice is the reference; everything
    else is balanced relative to it.
    """
    n = max(len(rhythm), len(melody), len(voice))
    def pad(x):
        out = np.zeros(n, dtype=np.float32)
        out[:len(x)] = x
        return out
    rhythm, melody, voice = pad(rhythm), pad(melody), pad(voice)

    # Relative timbre weights between rhythm and melody (some instruments are
    # naturally softer/louder per-voice). These set the BALANCE between the
    # two instrument layers; the overall instrument level is set afterward by
    # RMS targeting against the voice.
    melody_rel = {
        "ocarina":   0.45,
        "xylophone": 1.00,
        "harp":      0.85,
        "musicbox":  0.80,
        "guitar":    0.75,
    }.get(melody_instrument, 0.75)
    if same_instrument:
        melody_rel *= 0.7
    has_melody = melody_instrument is not None and np.any(melody)

    # Build the raw instrument bed (rhythm + weighted melody). The music box
    # reads a touch hot per-voice, so trim it slightly when it's the rhythm
    # layer too (the melody case is handled by melody_rel above).
    # Per-instrument rhythm weight (the continuous accompaniment layer). The
    # baseline is 0.7; music box reads hot so it's trimmed, piano sits a touch
    # quiet as accompaniment so it's nudged up a tad.
    rhythm_rel = 0.7 * {
        "musicbox": 0.85,
        "piano":    1.15,
    }.get(rhythm_instrument, 1.0)
    bed = rhythm * rhythm_rel + (melody * melody_rel if has_melody else 0.0)

    # Measure loudness of voice and bed on their active regions.
    v_rms = _active_rms(voice)
    bed_rms = _active_rms(bed)

    # TARGET: instruments sit at this fraction of the voice's active
    # loudness — present as a cushion but below the voice. Slightly higher
    # when there's no melody (so a lone rhythm doesn't feel thin).
    # Voice sits a touch louder in the mix: the instrument bed is balanced to
    # this fraction of the voice's active-RMS. Lowering it nudges the voice up
    # relative to the bed (0.47/0.55 → 0.42/0.50).
    target_ratio = 0.50 if not has_melody else 0.42

    if v_rms > 1e-6 and bed_rms > 1e-6:
        bed_scale = (v_rms * target_ratio) / bed_rms
    else:
        bed_scale = 0.5  # fallback if we can't measure (e.g. silent voice)

    # Clamp the scale to sane bounds so a quirk in measurement can't blow the
    # bed up or crush it to nothing.
    bed_scale = float(np.clip(bed_scale, 0.05, 1.2))

    # Voice is the reference layer at a steady gain.
    v_gain = 1.0

    mix = bed * bed_scale + voice * v_gain

    # Final peak safety (keep headroom; don't normalize UP, only down).
    peak = float(np.max(np.abs(mix)))
    if peak > 0.95:
        mix *= 0.95 / peak
    return mix


# ---------- end-to-end ----------

# Single unified instrument grid. Order matters — this is the display order
# in the picker (user-specified).
ALL_INSTRUMENTS = [
    ("musicbox",  "Music Box"),
    ("guitar",    "Guitar"),
    ("piano",     "Keyboard"),
    ("ocarina",   "Ocarina"),
    ("harp",      "Harp"),
    ("xylophone", "Xylophone"),
]

# Role classification — drives the picker rules and role resolution.
#   - "rhythm_first": when paired with something else, this plays rhythm.
#     Guitar+keyboard is the exception (keyboard plays rhythm, guitar plays
#     melody) — handled explicitly in resolve_roles.
#   - "melody_first": when paired with anything else, this plays melody.
#     When solo, falls back to rhythm (the song needs SOMETHING).
#   - "dual": role decided by pairing partner / random when solo.
RHYTHM_FIRST = {"guitar", "piano"}
MELODY_FIRST = {"ocarina", "harp", "xylophone"}
DUAL_ROLE    = {"musicbox"}

INSTRUMENT_META  = ALL_INSTRUMENTS
VALID_INSTRUMENTS = {k for k, _ in ALL_INSTRUMENTS}
# Backwards-compat aliases so any older code in the file doesn't break.
RHYTHM_INSTRUMENTS = ALL_INSTRUMENTS
MELODY_INSTRUMENTS = ALL_INSTRUMENTS
VALID_RHYTHM       = VALID_INSTRUMENTS
VALID_MELODY       = VALID_INSTRUMENTS


def _validation_error(message):
    """Return outputs that show an error banner instead of generating audio."""
    banner_html = f'<div class="error-banner-inner">⚠️ {message}</div>'
    return (
        gr.update(value=None, playback_position=0),  # audio_out
        gr.update(value=""), # lyrics_out
        banner_html,         # error_banner HTML content
        "",                  # raw_state (cleared on error)
    )


def resolve_roles(instrument_1, instrument_2):
    """
    Convert the user's ordered selection into (rhythm_instrument,
    melody_instrument). Returns (None, None, error_message) on invalid combo.

    Rules:
      Solo (just instrument_1):
        - musicbox → 50/50: (musicbox, None) or (musicbox, musicbox)
        - guitar/piano → (that, None) — no melody
        - ocarina/harp/xylophone → (that, None) — they cover both via the
          rhythm renderer (held chord tones), since these instruments make
          for thin solos
      Pair (both set):
        - guitar + piano → keyboard rhythm, guitar melody (the one specific
          exception where guitar plays melody)
        - rhythm_first + melody_first → rhythm_first as rhythm, the other as melody
        - rhythm_first + musicbox → rhythm_first as rhythm, musicbox as melody
        - musicbox + melody_first → musicbox as rhythm, melody_first as melody
        - melody_first + melody_first → INVALID (both want melody)
        - same instrument twice → INVALID (the picker should already prevent this)
    """
    a = instrument_1
    b = instrument_2 if instrument_2 in VALID_INSTRUMENTS else None

    if a not in VALID_INSTRUMENTS:
        return None, None, "Please pick at least one instrument."

    # Solo
    if b is None:
        if a == "musicbox":
            # Coin flip: rhythm-only OR rhythm + melody (both musicbox)
            if random.random() < 0.5:
                return "musicbox", None, None
            return "musicbox", "musicbox", None
        # All other solos: rhythm only.
        return a, None, None

    # Same instrument twice — shouldn't happen via the picker but be safe
    if a == b:
        return None, None, "Pick two different instruments."

    # Special: guitar + keyboard → keyboard rhythm, guitar melody
    if {a, b} == {"guitar", "piano"}:
        return "piano", "guitar", None

    # Both melody-first instruments — invalid combo (two-melody isn't a thing)
    if a in MELODY_FIRST and b in MELODY_FIRST:
        return None, None, "Pick at most one of harp, ocarina, xylophone."

    # General pairing logic. Priority order for rhythm: piano > guitar >
    # musicbox > (else the rhythm_first one).
    rhythm_priority = ["piano", "guitar", "musicbox"]
    rhythm = None
    melody = None
    for cand in rhythm_priority:
        if cand in (a, b):
            rhythm = cand
            break
    if rhythm is None:
        # Neither is in rhythm_priority — only happens if both are melody_first
        # which we already caught above. Defensive fallback.
        rhythm = a
        melody = b
    else:
        melody = b if rhythm == a else a

    return rhythm, melody, None


def _render_lyrics_view(raw, just_the_words):
    """Render the lyrics box. just_the_words=True → clean lyrics only (chords
    and tempo/progression headers stripped). False → the full technical view
    (tempo, progression, [chord] markers) as generated."""
    if not raw:
        return ""
    if not just_the_words:
        return raw  # full technical view as generated
    try:
        return parse_lullaby(raw)["plain_lyrics"]
    except ValueError:
        return raw  # unparseable — show raw rather than an empty box


def _friendly_instrument(key):
    """Map internal instrument keys to friendly names for progress messages."""
    return {
        "musicbox": "music box",
        "guitar": "guitar",
        "piano": "piano",
        "harp": "harp",
        "xylophone": "xylophone",
        "ocarina": "ocarina",
    }.get(key, key or "instrument")


def _validate_basics(name, loves, fears, instrument_1, instrument_2):
    """
    Run the cheap, deterministic validation checks that should fail FAST,
    before any expensive work (vision, lyric model, audio synth).

    Returns a `_validation_error(...)` tuple if any check fails, or
    `None` if everything checks out and the caller should proceed.

    Mirrors the same checks `make_lullaby` does at its top, so if
    those pass here, they'll pass there too. We don't safety-screen
    vision output here (that still happens in make_lullaby) — only the
    typed inputs get pre-screened.
    """
    if not name or not name.strip():
        return _validation_error("Please enter a name.")

    # Safety on TYPED inputs only — vision output is screened later in
    # make_lullaby. We still want to catch a user typing something
    # unsafe before paying for vision.
    safe, safe_msg, _bad = screen_inputs(loves, fears)
    if not safe:
        return _validation_error(safe_msg)

    _r, _m, role_err = resolve_roles(instrument_1, instrument_2)
    if role_err:
        return _validation_error(role_err)

    return None


def _status_html(pct, label):
    """Fake progress bar — amber Citrus theme colours, blinking label."""
    if pct is None:
        return ""
    return (
        f"<div class='lola-progress-wrap'>"
        f"<div class='lola-progress-bar' style='width:{pct}%'></div>"
        f"<span class='lola-progress-label'>{label} — {pct}%</span>"
        f"</div>"
    )


def _clear_outputs():
    """Clear all outputs instantly on click before the generator fires."""
    IDLE_BANNER = ("<div class='success-banner-idle'>"
                   "Fill in the details and press the button, "
                   "then Lola will sing here.</div>")
    return (
        None,
        gr.update(value=""),
        "",
        "",
        "",
        IDLE_BANNER,
        _status_html(5, "Looking at your drawing"),
    )


def make_lullaby_from_drawing(drawing, uploaded, name, age, loves, fears, mood, key,
                                meter, instrument_1, instrument_2,
                                just_the_words=True):
    """
    Drawing-aware front door for the same generation pipeline.

    Accepts TWO image sources because a 4-year-old can't draw on a laptop
    trackpad — but their parent or teacher can snap a photo of a paper
    drawing on their phone and upload it. The two inputs are alternatives,
    not additions, to the same "what did the child draw" signal:

        - Upload provided (photo of paper drawing) → use the upload.
          This is the higher-fidelity path and wins when both are present.
        - Canvas only → use the canvas drawing.
        - Neither → fall back to typed `loves` field below.

    Behaviour rules for the `loves` field, decided by the user:
        - Drawing only        → drawing's interpretation IS the loves.
        - Typed only          → typed value IS the loves.
        - BOTH drawing+typed  → COMBINE: drawing first, then typed
                                ("a cosy fire, trains") so both flavour
                                the lyrics.
        - Neither             → empty loves (`make_lullaby` handles it).

    Empty-canvas detection: `vision.describe()` returns None for a truly
    blank canvas (`_to_pil` checks all-white), and the stroke fallback also
    returns None on near-empty input. So "drawing_loves is None" means
    "user didn't actually draw anything" — and we just use the typed value.

    Fears and mood always come from the form — a drawing can't really
    convey "scared of thunder".

    Returns 5 values: (audio_path, lyrics_update, error_html, raw_text,
    saw_hint_html). The 5th is the human-readable "what Lola saw" line
    shown above the audio output so the user understands HOW the drawing
    became part of the lyrics.
    """
    drawing_loves = None
    drawing_source = None
    typed_loves = (loves or "").strip()

    # ── trace: begin a new run and capture the user-facing inputs ──
    lola_trace.begin()
    lola_trace.set_inputs(
        name=name, age=age, typed_loves=typed_loves, fears=fears,
        mood=mood, key=key, meter=meter,
        instrument_1=instrument_1, instrument_2=instrument_2,
        canvas_used=(drawing is not None),
        upload_used=(uploaded is not None),
    )

    # ── Early validation: catch cheap failures BEFORE running vision.
    # The vision call is the most expensive stage (10-60s); making a
    # user wait through it only to be told "please enter a name" is bad
    # UX. Same validators that make_lullaby uses internally — we just
    # run them up front when the inputs are obviously incomplete.
    early_err = _validate_basics(name, typed_loves, fears,
                                  instrument_1, instrument_2)
    if early_err is not None:
        audio, lyrics_upd, err_html, raw = early_err
        idle_msg = ("<div class='success-banner-idle'>"
                    "Fill in the details and press the button, "
                    "then Lola will sing here.</div>")
        # Trace finalize for the error path so the run still gets recorded.
        try:
            lola_trace.finalize(error=err_html)
        except Exception:
            pass
        yield audio, lyrics_upd, err_html, raw, "", idle_msg, ""
        return

    # Pick the image source: upload wins over canvas (higher fidelity).
    image_to_describe = None
    if uploaded is not None:
        image_to_describe = uploaded
        print("[drawing→loves] using UPLOADED image")
    elif drawing is not None:
        image_to_describe = drawing
        print("[drawing→loves] using CANVAS drawing")

    # ── trace: stash the input drawing for the trace folder ──
    if image_to_describe is not None:
        lola_trace.save_input_drawing(image_to_describe)

    if image_to_describe is not None and vision is not None:
        try:
            interp = vision.describe(image_to_describe, prefer="auto")
            if interp and interp.get("loves"):
                drawing_loves = interp["loves"].strip() or None
                drawing_source = interp.get("source", "?")
                print(f"[drawing→loves] ({drawing_source}): {drawing_loves!r}")
            else:
                print("[drawing→loves] canvas empty — using typed loves only")
            # ── trace: vision stage (success or empty) ──
            lola_trace.stage("vision",
                model="openbmb/MiniCPM-V-4.6 (with strokes fallback)",
                input_source=("uploaded" if uploaded is not None else "canvas"),
                interpreted_loves=drawing_loves,
                source=drawing_source,
                raw_output=(interp.get("raw") if interp else None))
        except Exception as e:
            print(f"[drawing→loves] failed: {type(e).__name__}: {e}; "
                  f"using typed loves only")
            lola_trace.stage("vision",
                model="openbmb/MiniCPM-V-4.6",
                error=f"{type(e).__name__}: {e}")

    # Compose the effective loves per the rules above.
    if drawing_loves and typed_loves:
        effective_loves = f"{drawing_loves}, {typed_loves}"
        print(f"[drawing→loves] combining drawing + typed: {effective_loves!r}")
    elif drawing_loves:
        effective_loves = drawing_loves
    else:
        effective_loves = typed_loves           # may be empty string — OK

    # Build the "what Lola saw" hint — shown above the audio so the user can
    # see why the lyrics came out the way they did. Plain-language, honest:
    # we don't claim the model saw a specific object unless it actually did.
    saw_hint = _build_saw_hint(drawing_loves, drawing_source, typed_loves,
                                drawing_present=(image_to_describe is not None))

    IDLE_BANNER = ("<div class='success-banner-idle'>"
                   "Fill in the details and press the button, "
                   "then Lola will sing here.</div>")

    # ── YIELD 0: vision done — 30%
    yield (None, "", "", "", saw_hint, IDLE_BANNER,
           _status_html(30, "Writing the lullaby"))

    STEPS = [
        (60, "Teaching Lola to sing the words"),
        (85, "Playing the instruments"),
    ]

    audio = None
    err_html = ""
    raw = ""
    lyrics_upd = ""
    gen = make_lullaby(
        name, age, effective_loves, fears, mood, key, meter,
        instrument_1, instrument_2,
        just_the_words=just_the_words,
    )
    for step_idx, partial in enumerate(gen):
        audio, lyrics_upd, err_html, raw = partial
        hint_to_show = "" if err_html else saw_hint
        if err_html:
            yield (audio, lyrics_upd, err_html, raw, hint_to_show, IDLE_BANNER, "")
        else:
            pct, label = STEPS[step_idx] if step_idx < len(STEPS) else (85, "Almost there")
            yield (audio, lyrics_upd, err_html, raw, hint_to_show, IDLE_BANNER,
                   _status_html(pct, label))

    if audio and not err_html:
        success_html = (
            "<div class='success-banner-inner'>"
            "Your lullaby is ready — press play to hear it."
            "</div>"
        )
    else:
        success_html = IDLE_BANNER

    try:
        if audio and not err_html:
            import soundfile as sf
            wav, sr = sf.read(audio)
            lola_trace.finalize(output_audio=wav, sample_rate=sr, error=None)
        else:
            lola_trace.finalize(error=(err_html or "no audio produced"))
    except Exception as _trace_err:
        print(f"[trace] finalize skipped: {_trace_err}")

    yield (audio, lyrics_upd, err_html, raw,
           "" if err_html else saw_hint, success_html, "")


def _build_saw_hint(drawing_loves, drawing_source, typed_loves, drawing_present):
    """Compose the user-facing 'what informed this lullaby' line, written as
    LOLA speaking to the user in first person. Consistent character voice
    across all five cases (vision sees, strokes feels, typed only, empty
    canvas, nothing at all) — so the hint reads like a small voice talking
    to you, not a system label.

    The vision/strokes distinction is preserved in the verb: vision really
    saw objects ('I'm seeing...'), strokes only read mood ('I can feel...').
    Don't claim vision when it was really strokes."""
    has_drawing = bool(drawing_loves)
    has_typed = bool(typed_loves)

    # Verb depends honestly on which engine read the drawing.
    if drawing_source == "vision":
        drawing_verb = "I'm seeing"
    else:                                          # "strokes"
        drawing_verb = "I can feel"

    if has_drawing and has_typed:
        # Lyrics use BOTH (combine happens in the wrapper), but the hint
        # focuses on the drawing interpretation — the magic moment.
        return (f"<div class='saw-hint'>"
                f"<span class='saw-label'>Lola:</span> "
                f"<em>{drawing_verb} {drawing_loves}.</em></div>")
    if has_drawing:
        return (f"<div class='saw-hint'>"
                f"<span class='saw-label'>Lola:</span> "
                f"<em>{drawing_verb} {drawing_loves}.</em></div>")
    if has_typed:
        return (f"<div class='saw-hint'>"
                f"<span class='saw-label'>Lola:</span> "
                f"<em>I'll sing about {typed_loves}.</em></div>")
    # Neither — user clicked Sing with nothing
    if drawing_present:
        return ("<div class='saw-hint saw-hint-empty'>"
                "<span class='saw-label'>Lola:</span> "
                "<em>I couldn't see anything on the canvas, so I sang a "
                "gentle one. Try drawing something or typing in "
                "<u>What do they love?</u></em></div>")
    return ("<div class='saw-hint saw-hint-empty'>"
            "<span class='saw-label'>Lola:</span> "
            "<em>Nothing to go on this time, so I sang a gentle one.</em>"
            "</div>")


def make_lullaby(name, age, loves, fears, mood, key, meter,
                 instrument_1, instrument_2, just_the_words=True):
    """
    Generate a lullaby from the unified picker.

    Yields partial results so the UI can update incrementally:
      1. After lyrics are written  → yields lyrics (audio still loading)
      2. After audio is ready      → yields audio + success banner

    instrument_1: first instrument clicked
    instrument_2: second instrument clicked (or empty)

    Role resolution is delegated to `resolve_roles` and follows the rules
    set by the user — see that function's docstring.
    """
    # Sentinel tuple shape: (audio, lyrics_upd, err_html, raw)
    # We yield gr.update() for outputs that aren't ready yet so Gradio
    # leaves them in their current state (spinner running) until we fill them.
    if not name or not name.strip():
        yield _validation_error("Please enter a name.")
        return

    # Content safety (Fix A): screen the user's free-text loves/fears so a
    # lullaby can't be built around inappropriate themes (death, weapons,
    # violence, etc.). Reject with a gentle, kid-app-appropriate message.
    safe, safe_msg, _bad = screen_inputs(loves, fears)
    lola_trace.stage("safety_input",
        module="safety.py", screened_loves=loves, screened_fears=fears,
        result=("passed" if safe else "blocked"),
        bad_terms=_bad if not safe else [])
    if not safe:
        yield _validation_error(safe_msg)
        return

    rhythm_instrument, melody_instrument, role_err = resolve_roles(
        instrument_1, instrument_2)
    if role_err:
        yield _validation_error(role_err)
        return

    has_melody = melody_instrument is not None
    # Same-instrument-both-layers happens specifically when musicbox solo
    # rolls the "both" outcome.
    same_instrument = (rhythm_instrument == melody_instrument
                        and melody_instrument is not None)

    # Resolve "Random" choices to actual values so every generation differs.
    if key == "Random" or not key:
        key = random.choice([
            "C major", "G major", "D major", "F major", "A major",
            "A minor", "E minor", "D minor",
        ])
    if meter == "Random" or not meter:
        meter = random.choice(["6/8", "3/4", "4/4"])

    try:
        prompt = build_prompt(name, age, loves, fears, mood, key, meter)
        raw = generate_lullaby(prompt)
        try:
            parsed = parse_lullaby(raw)
        except ValueError:
            raw = generate_lullaby(prompt, temperature=0.4)
            parsed = parse_lullaby(raw)

        # ── YIELD 1: lyrics ready
        yield (None,
               gr.update(value=_render_lyrics_view(raw, just_the_words)),
               "",
               raw)

        plan = make_arrangement_plan(rhythm_instrument, melody_instrument,
                                       has_melody, key=key)
        # When the same instrument plays both layers (musicbox solo "both"
        # case), the melody needs to sit a bit above the rhythm voicing.
        plan["melody_octave_shift"] = 12 if same_instrument else 0

        # Render the VOICE first so we know its ACTUAL length. The n_lines
        # estimate (body_seconds = n_lines * chord_seconds) can underestimate
        # how long Kokoro takes (more words per line, slower pacing), which
        # would leave a long lyric playing over silence. We measure the real
        # voice length and extend the instrument bed to cover it.
        intro_seconds_est = _intro_seconds(plan)
        n_lines_est = len(parsed["lines"])
        body_seconds_est = n_lines_est * plan["chord_seconds"]
        voice_body = speak_lyrics(parsed["plain_lyrics"],
                                   target_seconds=body_seconds_est)
        # Actual sung duration (without the intro padding that speak_lyrics may
        # have added — we add our own intro below).
        voice_seconds = len(voice_body) / SR_TARGET

        # Extend the instrument body to cover the longer of (estimate, actual
        # voice) so there is ALWAYS music under every word.
        min_body = max(body_seconds_est, voice_seconds)
        rhythm_audio, body_seconds, intro_seconds = render_rhythm(
            parsed, rhythm_instrument, plan, min_body_seconds=min_body)
        if has_melody:
            melody_audio = render_melody(parsed, melody_instrument,
                                          plan, intro_seconds)
        else:
            melody_audio = np.zeros(len(rhythm_audio), dtype=np.float32)

        if intro_seconds > 0:
            pad_samples = int(intro_seconds * SR)
            voice_audio = np.concatenate([
                np.zeros(pad_samples, dtype=np.float32),
                voice_body.astype(np.float32),
            ])
        else:
            voice_audio = voice_body.astype(np.float32)

        mix = mix_tracks(rhythm_audio, melody_audio, voice_audio,
                         melody_instrument=melody_instrument,
                         rhythm_instrument=rhythm_instrument,
                         same_instrument=same_instrument)
        lola_trace.stage("audio",
            voice="hexgrad/Kokoro-82M (af_nicole)",
            rhythm_instrument=rhythm_instrument,
            melody_instrument=melody_instrument,
            tempo_bpm=plan.get("bpm"),
            key=key, meter=meter,
            voice_seconds=round(voice_seconds, 2),
            total_seconds=round(len(mix) / SR, 2),
            sample_rate=SR)

        # Clean up older generated files so they don't pile up in the Space's
        # temp dir over a long session.
        try:
            now = time.time()
            for old in glob.glob(str(Path(tempfile.gettempdir()) / "lullaby_*.wav")):
                if now - os.path.getmtime(old) > 600:  # older than 10 min
                    os.remove(old)
        except OSError:
            pass

        # Unique filename per generation. Reusing one fixed path
        # ("lullaby_output.wav") makes the browser audio player keep the OLD
        # file cached — so a newly generated song doesn't reload and the
        # playhead doesn't reset to the start. A fresh name forces a reload.
        out_path = (Path(tempfile.gettempdir())
                    / f"lullaby_{uuid.uuid4().hex}.wav")
        sf.write(str(out_path), mix, SR, subtype="PCM_16")
    except Exception as e:
        yield _validation_error(f"Something broke while making the lullaby: {e}")
        return

    # ── YIELD 2 (final): audio is ready.
    # playback_position=0 explicitly resets the Gradio audio component's
    # tracked position — without this, Gradio 6 restores whatever position
    # the previous song was paused at, even with a fresh filename.
    yield (gr.update(value=str(out_path), playback_position=0),
           gr.update(value=_render_lyrics_view(raw, just_the_words)),
           "",
           raw)


# ---------- UI ----------

ASSETS_DIR = Path(__file__).parent / "assets"


def _img_data_uri(name):
    """Load assets/<name>.png as a data URI so it inlines into HTML.
    Falls back to a small SVG with the first letter of the name if the
    PNG isn't present yet — keeps the cell from rendering broken when
    a new instrument is added without an asset."""
    # Filename aliases: internal key → preferred PNG filename.
    # (Some assets are named with hyphens for readability.)
    filename_aliases = {
        "musicbox": ["music-box", "musicbox"],
    }
    candidates = filename_aliases.get(name, [name])
    for stem in candidates:
        path = ASSETS_DIR / f"{stem}.png"
        if path.exists():
            data = path.read_bytes()
            b64 = base64.b64encode(data).decode("ascii")
            return f"data:image/png;base64,{b64}"

    # SVG fallback: a small circle with the instrument's first letter
    initial = name[:1].upper() if name else "?"
    palette = {
        "guitar":    "#e15a4e",
        "piano":     "#2d3142",
        "harp":      "#c98c52",
        "xylophone": "#e85a8c",
        "ocarina":   "#6fb05e",
        "musicbox":  "#a86fc1",
    }
    color = palette.get(name, "#4a90d9")
    svg = (
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 56 56">'
        f'<circle cx="28" cy="28" r="24" fill="{color}"/>'
        f'<text x="28" y="36" text-anchor="middle" font-family="Fredoka One, Arial Black, sans-serif" '
        f'font-size="24" fill="white" font-weight="700">{initial}</text>'
        f'</svg>'
    )
    b64 = base64.b64encode(svg.encode("utf-8")).decode("ascii")
    return f"data:image/svg+xml;base64,{b64}"



def instrument_grid_html(group_id, instruments, selected=None):
    """Build the HTML for one instrument-selection grid."""
    cells = []
    for key, label in instruments:
        uri = _img_data_uri(key)
        sel_class = "selected" if key == selected else ""
        cells.append(f"""
        <button type="button"
                class="instr-cell {sel_class}"
                data-instr="{key}"
                data-group="{group_id}"
                onclick="pickInstrument(this)">
          <img src="{uri}" alt="{label}" />
          <span>{label}</span>
        </button>
        """)
    return f'<div class="instr-grid">{"".join(cells)}</div>'


ATTRIBUTION_HTML = """
<div class="footer">
  <!--<div class="attribution">
    Icons:
    <a href="https://www.flaticon.com/free-icons/music-and-multimedia" target="_blank">Guitar</a>,
    <a href="https://www.flaticon.com/free-icons/piano" target="_blank">Keyboard</a>,
    <a href="https://www.flaticon.com/free-icons/harp" target="_blank">Harp</a>,
    <a href="https://www.flaticon.com/free-icons/xylophone" target="_blank">Xylophone</a>,
    <a href="https://www.flaticon.com/free-icons/ocarina" target="_blank">Ocarina</a>,
    <a href="https://www.flaticon.com/free-icons/music-box" title="music box icons" target="_blank">Music Box</a>
    · created by Freepik, Smashicons &amp; iconixar on Flaticon
  </div>-->
  <div class="powered-by">
    Powered by
    <a href="https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct" target="_blank" rel="noopener">Llama 3.2</a>
    ·
    <a href="https://huggingface.co/openbmb/MiniCPM-V-4_6" target="_blank" rel="noopener">MiniCPM-V 4.6</a>
    ·
    <a href="https://huggingface.co/hexgrad/Kokoro-82M" target="_blank" rel="noopener">Kokoro 82M</a>
    ·
    <a href="https://github.com/ggml-org/llama.cpp" target="_blank" rel="noopener">llama.cpp</a>
  </div>
</div>
"""


# ---------------------------------------------------------------------------
# EARLY_THEME_JS — must run as the FIRST thing in <head>, synchronously,
# before any CSS/markup is painted.
# ---------------------------------------------------------------------------
EARLY_THEME_JS = """
(function() {
    try {
        var url = new URL(window.location.href);
        var isDark;
        if (url.searchParams.get('__theme') === 'dark') {
            isDark = url.searchParams.get('__theme') === 'dark';
        } else {
            isDark = !!(window.matchMedia &&
                         window.matchMedia('(prefers-color-scheme: dark)').matches);
        }
        if (isDark) {
            document.documentElement.classList.add('dark');
        }
    } catch (e) {
        // If anything above throws (e.g. matchMedia unavailable), just fall
        // back to the light-mode default rather than breaking the page.
    }
})();
"""

# JS injected once; the buttons call pickInstrument() which writes the chosen
# value into the hidden Gradio textbox for that group and triggers its change.
SELECTION_JS = """
// ── Dark mode toggle ──────────────────────────────────────────────────────
(function() {
    function applyTheme(dark) {
        document.documentElement.classList.toggle('dark', dark);
        document.body.classList.toggle('dark', dark);
        const container = document.querySelector('.gradio-container');
        if (container) container.classList.toggle('dark', dark);
    }

    function toggleDark() {
        const isDark = document.documentElement.classList.contains('dark');
        const newTheme = isDark ? 'light' : 'dark';

        applyTheme(!isDark);

        const url = new URL(window.location);
        url.searchParams.set('__theme', newTheme);
        history.replaceState(null, '', url);
    }

    window._lolaby_applyTheme = applyTheme;
    window._lolaby_toggleDark = toggleDark;

    function syncThemeClasses() {
        const isDark = document.documentElement.classList.contains('dark');
        applyTheme(isDark);
    }

    if (document.readyState === 'loading') {
        document.addEventListener('DOMContentLoaded', syncThemeClasses);
    } else {
        syncThemeClasses();
    }
})();

// Ordered multi-select picker. Up to 2 instruments.
// Slot 1 = first click, slot 2 = second click.
// Click an already-selected one to deselect; if you remove slot 1 with
// slot 2 still selected, the remaining selection slides into slot 1.
//
// Combo validation:
//   - Cannot pick two of {harp, ocarina, xylophone} (both want to be melody)
//   - Cannot select the same instrument twice (it's the same button anyway)
//   - At most 2 instruments
//
// A blocked click flashes the picker briefly so the user notices.

const MELODY_FIRST = new Set(["ocarina", "harp", "xylophone"]);

function pickInstrument(btn) {
    const instr = btn.dataset.instr;
    const cells = Array.from(document.querySelectorAll('.instr-cell[data-group="instr"]'));

    // Current ordered selection
    let selected = [];
    cells.forEach(el => {
        if (el.classList.contains('selected')) {
            const slot = parseInt(el.dataset.slot || '0', 10);
            if (slot > 0) selected[slot - 1] = el.dataset.instr;
        }
    });
    selected = selected.filter(x => x);

    const idx = selected.indexOf(instr);
    if (idx >= 0) {
        // Deselect — always allowed. Also clears any active error.
        selected.splice(idx, 1);
        clearPickerError();
    } else if (selected.length >= 2) {
        // Already at max — reject with a brief shake
        flashLimitHint();
        return;
    } else {
        // Validate the proposed combo
        const proposed = [...selected, instr];
        const melodyFirstCount = proposed.filter(i => MELODY_FIRST.has(i)).length;
        if (melodyFirstCount > 1) {
            // Two of harp/ocarina/xylophone — not allowed
            flashComboHint();
            return;
        }
        selected = proposed;
        // A successful selection also clears any lingering error
        clearPickerError();
    }

    // Re-apply visual state across all cells
    cells.forEach(el => {
        el.classList.remove('selected');
        el.dataset.slot = '0';
        const badge = el.querySelector('.slot-badge');
        if (badge) badge.remove();
    });
    selected.forEach((instrKey, i) => {
        const cell = cells.find(el => el.dataset.instr === instrKey);
        if (!cell) return;
        cell.classList.add('selected');
        cell.dataset.slot = String(i + 1);
        const badge = document.createElement('span');
        badge.className = 'slot-badge';
        badge.textContent = String(i + 1);
        cell.appendChild(badge);
    });

    // Push to hidden Gradio textboxes
    setHidden('pick-1', selected[0] || '');
    setHidden('pick-2', selected[1] || '');
}

function setHidden(eid, val) {
    const holder = document.getElementById(eid);
    if (!holder) return;
    const input = holder.querySelector('input, textarea');
    if (!input) return;
    input.value = val;
    input.dispatchEvent(new Event('input', { bubbles: true }));
    input.dispatchEvent(new Event('change', { bubbles: true }));
}

// Shared timeout id — so a new error cancels any pending revert and the
// previous error's text doesn't get captured as the new "original" text.
let pickerErrorTimeout = null;

function clearPickerError() {
    const hint = document.getElementById('picker-hint');
    if (!hint) return;
    if (pickerErrorTimeout !== null) {
        clearTimeout(pickerErrorTimeout);
        pickerErrorTimeout = null;
    }
    hint.textContent = '';
    hint.classList.remove('hint-flash');
    hint.classList.remove('hint-error');
}

function showPickerError(message, useErrorWeight) {
    const hint = document.getElementById('picker-hint');
    if (!hint) return;
    // Cancel any in-flight revert so the message stays visible for its
    // full window and doesn't inherit text from a previous flash.
    if (pickerErrorTimeout !== null) {
        clearTimeout(pickerErrorTimeout);
        pickerErrorTimeout = null;
    }
    hint.textContent = message;
    hint.classList.add('hint-flash');
    if (useErrorWeight) hint.classList.add('hint-error');
    pickerErrorTimeout = setTimeout(() => {
        hint.textContent = '';
        hint.classList.remove('hint-flash');
        hint.classList.remove('hint-error');
        pickerErrorTimeout = null;
    }, 1600);
}

function flashLimitHint() {
    showPickerError("You can only select up to two instruments.", true);
}

function flashComboHint() {
    showPickerError("Only one of harp, ocarina, xylophone at a time", false);
}

// ── Blink the audio/lyrics widgets while Lola is generating ────────────────
// We watch #lola-status: when it gets content the generation is running;
// when it empties (cleared by _clear_outputs or the final yield) it stops.
// The class .lola-generating is placed on the right-card column so the CSS
// rules inside it target the audio player and lyrics textarea precisely.
(function() {
    function getRightCard() {
        return document.querySelector('.form-card-right');
    }

    function syncGenerating() {
        const status = document.getElementById('lola-status');
        const card = getRightCard();
        if (!card) return;
        if (status && status.textContent.trim() !== '') {
            card.classList.add('lola-generating');
        } else {
            card.classList.remove('lola-generating');
        }
    }

    function attachObserver() {
        const status = document.getElementById('lola-status');
        if (!status) {
            // DOM not ready yet — retry
            setTimeout(attachObserver, 300);
            return;
        }
        syncGenerating();
        new MutationObserver(syncGenerating).observe(status, {
            childList: true, characterData: true, subtree: true, attributes: true
        });
    }

    if (document.readyState === 'loading') {
        document.addEventListener('DOMContentLoaded', attachObserver);
    } else {
        attachObserver();
    }
})();
"""

HEAD_HTML = f"<script>{EARLY_THEME_JS}</script><script>{SELECTION_JS}</script>"


with gr.Blocks(css_paths="style.css", title="Lolaby", theme=gr.themes.Citrus(),
                head=HEAD_HTML) as demo:
    gr.HTML("""
    <div style="text-align:right; padding: 12px 18px 0;">
    <label class="dark-toggle-label">
        <span id="theme-emoji">🌗</span>
        <input type="checkbox" id="dark-toggle-btn" role="switch"
            onchange="window._lolaby_toggleDark && window._lolaby_toggleDark()">
    </label>
    </div>
    """)
    gr.HTML("""
    <div id="title-block">
      <h1>Lolaby</h1>
      <div class="subtitle">~ AI-powered lullabies ~</div>
    </div>
    <div id="title-explainer">
      Meet Lola, your personal bedtime singer
    </div>
    <span class="explainer-roles"></span>
    <!-- Three-step diagram. Acts as wayfinding for non-power users who
         need to see the sequence at a glance before scrolling. Step 1
         (Draw) is dashed/lighter to signal "optional" — reinforces the
         copy elsewhere that says drawing isn't required. -->
    <div id="step-diagram" aria-label="How to use Lolaby in three steps">
      <div class="step step-1 step-optional">
        <div class="step-num">1</div>
        <div class="step-icon">🎨</div>
        <div class="step-label">Draw or upload a doodle<span class="step-sub">Or skip it and go to the next step</span></div>
      </div>
      <div class="step-connector"></div>
      <div class="step step-2">
        <div class="step-num">2</div>
        <div class="step-icon">🖍️</div>
        <div class="step-label">Fill in the details<span class="step-sub">Tell us about the little one</span></div>
      </div>
      <div class="step-connector"></div>
      <div class="step step-3">
        <div class="step-num">3</div>
        <div class="step-icon">🎶</div>
        <div class="step-label">Sing with Lola<span class="step-sub">Hit the button to play your lullaby</span></div>
      </div>
    </div>
    <span class="explainer-roles"></span>
    """)

    # ---- DRAWING CANVAS (the hero input) ---------------------------------
    # Above the form: kid draws what they love, Lola sings about it.
    # Drawing is OPTIONAL — if blank, the typed "loves" field is used.
    #
    # The card has TWO image inputs side-by-side because a 4-year-old can't
    # draw on a laptop trackpad. The upload path (right) is how a teacher or
    # parent gets a photo of a paper drawing into the app. The canvas (left)
    # is there for whoever IS comfortable drawing on screen.
    with gr.Row():
        with gr.Column(scale=1, elem_classes=["draw-card"]):
            gr.HTML(
                '<div class="card-tab card-tab-kid"><span class="step-badge">1</span> For the little ones</div>'
                '<div class="card-heading">🎨 Draw or upload a doodle here</div>'
                '<div class="draw-hint">This step is optional, '
                'but a drawing helps Lola make the song more personal. '
                'Snap a phone photo of their paper drawing on the right, or '
                'doodle on the left if they\'re comfortable with a trackpad. '
                'Don\'t worry about the details.</div>'
            )

            with gr.Row(equal_height=True, elem_id="draw-row"):
                # ── LEFT HALF: live canvas (trackpad/mouse drawing) ──────
                with gr.Column(scale=1, elem_classes=["draw-half", "draw-half-canvas"]):
                    gr.HTML('<div class="draw-half-label">Doodle here</div>')
                    canvas = gr.Sketchpad(
                        label="",
                        show_label=False,
                        type="numpy",
                        height=320,
                        brush=gr.Brush(
                            default_size=14,
                            colors=[
                                "#2d3142",   # ink (default — first in list)
                                "#e15a4e",   # crayon red
                                "#f0934a",   # crayon orange
                                "#f5c842",   # crayon yellow
                                "#6fb05e",   # crayon green
                                "#4a90d9",   # crayon blue
                                "#9b6db5",   # crayon purple
                                "#f08ab0",   # crayon pink
                                "#8a6a4a",   # crayon brown
                            ],
                        ),
                        elem_id="draw-canvas",
                    )
                    # Clearing a Sketchpad reliably is harder than it should be.
                    # Five attempts via the Python API all hit known Gradio bugs
                    # (#9978, #888, #501, #7816): None-returns fail on repeat,
                    # gr.update(value=None) is swallowed, fixed-size arrays change
                    # the zoom mode, and even gr.ClearButton stops working after the
                    # canvas participates in a function call.
                    #
                    # The path that actually works: trigger Sketchpad's OWN built-in
                    # clear button (which the framework hides by default but still
                    # renders in the DOM) via a JS click. The framework's internal
                    # clear operates directly on the canvas state and doesn't suffer
                    # the "value didn't change" or "initial value got mutated"
                    # problems the public API has.
                    reset_drawing_btn = gr.Button(
                        "Clear",
                        elem_id="reset-drawing-btn",
                        elem_classes=["secondary-btn"],
                    )

                # ── RIGHT HALF: photo upload (paper drawing → phone photo) ──
                with gr.Column(scale=1, elem_classes=["draw-half", "draw-half-upload"]):
                    gr.HTML('<div class="draw-half-label">Or upload a photo</div>')
                    uploaded = gr.Image(
                        label="",
                        show_label=False,
                        type="numpy",
                        sources=["upload"],
                        height=320,
                        elem_id="draw-upload",
                    )
                    gr.HTML(
                        '<div class="draw-half-hint">PNG or JPG. A phone '
                        'photo of their paper drawing works great.</div>'
                    )

            # Gentle notice that appears ONLY when both inputs are
            # populated. Tells the user upload will win (because it's
            # higher fidelity) before they're surprised by the result.
            # Hidden via inline display:none and toggled by JS-side
            # change handlers below.
            both_inputs_notice = gr.HTML(
                value="",
                elem_id="both-inputs-notice",
            )

    with gr.Row():
        with gr.Column(scale=1, elem_classes=["form-card"]):
            gr.HTML(
                '<div class="card-tab card-tab-grownup"><span class="step-badge">2</span> For grown-ups</div>'
                '<div class="card-heading">🖍️ Tell us about the little one</div>'
                '<div class="form-role-hint">Only the name and at least one '
                'instrument are required. Everything '
                'else is optional, but it helps Lola make the song fit them.</div>'
            )
            name = gr.Textbox(label="Name *", placeholder="Lucy")
            age = gr.Slider(label="How old?", minimum=0, maximum=8, step=1, value=3)
            loves = gr.Textbox(
                label="What do they love?",
                placeholder="her stuffed elephant Pip, the moon",
                lines=2,
            )
            fears = gr.Textbox(
                label="Anything scary?",
                placeholder="the dark",
            )
            mood = gr.Dropdown(
                label="How are they feeling tonight?",
                choices=[
                    "sleepy and comforted",
                    "restless but settling",
                    "tearful, needs soothing",
                    "wide awake, needs winding down",
                    "cosy and content",
                ],
                value="sleepy and comforted",
            )

            gr.HTML(
                '<div class="section-title">Pick your instruments *</div>'
                '<div id="picker-hint"></div>'
            )
            gr.HTML(value=instrument_grid_html("instr", ALL_INSTRUMENTS))
            # Hidden state holders — the JS sets these.
            # pick-1 = first instrument clicked, pick-2 = second.
            # Role resolution happens server-side in resolve_roles().
            pick_1 = gr.Textbox(value="", visible=True, elem_id="pick-1",
                                  interactive=True, label="instrument 1",
                                  elem_classes=["pick-state"])
            pick_2 = gr.Textbox(value="", visible=True, elem_id="pick-2",
                                  interactive=True, label="instrument 2",
                                  elem_classes=["pick-state"])

            with gr.Accordion("⚙️ Advanced settings", open=False,
                              elem_id="advanced-accordion"):
                with gr.Row():
                    key = gr.Dropdown(
                        label="Key",
                        choices=["Random", "C major", "G major", "D major", "F major",
                                 "A major", "A minor", "E minor", "D minor"],
                        value="Random",
                    )
                    meter = gr.Dropdown(
                        label="Beat",
                        choices=["Random", "6/8", "3/4", "4/4"],
                        value="Random",
                    )

            # Footnote moved to below the Sing button (next block).

            error_banner = gr.HTML(value="", elem_id="error-banner")

            gr.HTML("""<style>
/* ── Progress bar ── */
.lola-progress-wrap {
  position: relative;
  background: #fef3c7;
  border-radius: 999px;
  height: 26px;
  overflow: hidden;
  margin: 4px 0 6px;
}
.dark .lola-progress-wrap {
  background: #2a2d4a;
}
.lola-progress-bar {
  position: absolute;
  left: 0; top: 0; bottom: 0;
  background: linear-gradient(90deg, #fcd34d, #f59e0b);
  border-radius: 999px;
  transition: width 0.7s cubic-bezier(0.4, 0, 0.2, 1);
}
.dark .lola-progress-bar {
  background: linear-gradient(90deg, #c9a020, #d97706);
}
@keyframes lola-blink {
  0%, 100% { opacity: 1; }
  50%       { opacity: 0.45; }
}
.lola-progress-label {
  position: absolute;
  inset: 0;
  display: flex;
  align-items: center;
  justify-content: center;
  font-size: 0.76rem;
  font-weight: 600;
  color: #78350f;
  animation: lola-blink 1.6s ease-in-out infinite;
  white-space: nowrap;
}
.dark .lola-progress-label {
  color: #f5d080;
}

/* ── Blink audio player and lyrics while generating ── */
@keyframes lola-widget-blink {
  0%, 100% { opacity: 1; }
  50%       { opacity: 0.38; }
}
.lola-generating #output-lyrics textarea,
.lola-generating audio,
.lola-generating .waveform-container,
.lola-generating .gr-audio,
.lola-generating #audio-out,
.lola-generating #audio-out > div,
.lola-generating #audio-out .wrap,
.lola-generating #audio-out .block {
  animation: lola-widget-blink 1.4s ease-in-out infinite;
  pointer-events: none;
}
</style>""")
            status_out = gr.HTML(value="", elem_id="lola-status")
            gr.HTML('<div class="step-3-hint"><span class="step-badge step-badge-inline">3</span> Ready? Press the button.</div>')
            btn = gr.Button("Sing with Lola", elem_id="generate-btn")

        with gr.Column(scale=1, elem_classes=["form-card-right"]):
            gr.HTML('<div class="card-heading">🎶 Your lullaby</div>')
            # Success banner — appears below the section heading, above the
            # audio player. Has an IDLE message ("waiting for you...")
            # shown before the first generation, then gets replaced with
            # the green success state once a song is ready.
            success_banner = gr.HTML(
                value=("<div class='success-banner-idle'>"
                       "Fill in the details and press the button, "
                       "then Lola will sing here."
                       "</div>"),
                elem_id="success-banner",
            )
            audio_out = gr.Audio(label="", type="filepath", buttons=['download'], elem_id="audio-out")
            # "What Lola saw" hint — sits between the song output and the
            # lyrics, explaining how the drawing/typed inputs informed
            # what's playing. Stays blank until the first generation; hides
            # itself when empty.
            saw_hint_out = gr.HTML(value="", elem_id="saw-hint")
            #gr.HTML('<div class="card-heading">💬 Your words</div>')
            lyrics_out = gr.Textbox(
                label="",
                lines=14,
                interactive=False,
                elem_id="output-lyrics",
                show_label=False,
            )
            just_words = gr.Checkbox(
                value=True,
                label="Just the words",
                info="Show only the lyrics. Turn off to see tempo, key, "
                     "and chords for musicians.",
                elem_id="just-words-toggle",
            )
            # Holds the full generated text so toggling the view doesn't
            # require regenerating the song.
            raw_state = gr.State("")

    gr.HTML(ATTRIBUTION_HTML)

    ALL_OUTPUTS = [audio_out, lyrics_out, error_banner, raw_state,
                   saw_hint_out, success_banner, status_out]

    btn.click(
        _clear_outputs,
        inputs=None,
        outputs=ALL_OUTPUTS,
        queue=False,
        show_progress="hidden",
    ).then(
        make_lullaby_from_drawing,
        inputs=[canvas, uploaded, name, age, loves, fears, mood, key, meter,
                pick_1, pick_2, just_words],
        outputs=ALL_OUTPUTS,
    ).then(
        # After generation completes, scroll the lullaby section into view
        # IF it's not already visible. We use scrollIntoView with
        # block:'nearest' so a user already looking at the result isn't
        # yanked anywhere; only off-screen users get auto-scrolled.
        fn=None,
        inputs=None,
        outputs=None,
        js="""
        () => {
            // Wait one tick so the audio element has had time to mount.
            setTimeout(() => {
                const card = document.querySelector('.form-card-right');
                if (!card) return;
                const rect = card.getBoundingClientRect();
                const fullyVisible = (
                    rect.top >= 0 &&
                    rect.bottom <= (window.innerHeight ||
                                    document.documentElement.clientHeight)
                );
                if (!fullyVisible) {
                    card.scrollIntoView({
                        behavior: 'smooth',
                        block: 'start',
                    });
                }
            }, 200);
        }
        """,
    )

    # ── Both-inputs notice: reactively show/hide based on whether BOTH
    # canvas and upload have content. Gentle informational message, not
    # an error — upload wins silently per the wrapper's logic; we just
    # surface that fact so the user isn't surprised.
    def _both_inputs_notice(canv, up):
        if canv is not None and up is not None:
            return ("<div class='both-inputs-inner'>"
                    "<span class='both-inputs-icon'>💡</span> "
                    "If you upload a drawing and a photo, Lola will use the "
                    "photo to get more detail."
                    "</div>")
        return ""

    canvas.change(_both_inputs_notice, inputs=[canvas, uploaded],
                  outputs=[both_inputs_notice])
    uploaded.change(_both_inputs_notice, inputs=[canvas, uploaded],
                    outputs=[both_inputs_notice])

    # Trigger Sketchpad's OWN built-in clear by JS-clicking its hidden
    # internal button. This path works reliably even after the canvas has
    # been read by a function call, where the public ClearButton/value-reset
    # APIs all fail. No Python handler needed — pure JS.
    reset_drawing_btn.click(
        fn=None,
        inputs=None,
        outputs=None,
        js="""
        () => {
            const card = document.querySelector('#draw-canvas');
            if (!card) return;
            // Sketchpad renders its own toolbar buttons; the clear/erase
            // button is identifiable by aria-label. Try the most common
            // labels Gradio has used across versions.
            const labels = ['Clear', 'Clear canvas', 'Erase', 'Remove image'];
            let btn = null;
            for (const lbl of labels) {
                btn = card.querySelector(`button[aria-label="${lbl}"]`);
                if (btn) break;
            }
            // Fallback: any button whose title attribute or text mentions clear
            if (!btn) {
                const allBtns = card.querySelectorAll('button');
                for (const b of allBtns) {
                    const txt = (b.getAttribute('title') || b.textContent || '').toLowerCase();
                    if (txt.includes('clear') || txt.includes('erase') || txt.includes('remove')) {
                        btn = b;
                        break;
                    }
                }
            }
            if (btn) {
                btn.click();
            } else {
                // Last-resort: wipe the visible canvas pixels. Server state
                // may stay stale, but the next "Sing" will see a near-white
                // canvas and our _to_pil emptiness check (all-white) catches it.
                card.querySelectorAll('canvas').forEach(c => {
                    const ctx = c.getContext('2d');
                    if (ctx) {
                        ctx.fillStyle = '#ffffff';
                        ctx.fillRect(0, 0, c.width, c.height);
                    }
                });
            }
        }
        """,
    )

    # Flipping the toggle re-renders the stored text instantly (no model call).
    just_words.change(
        _render_lyrics_view,
        inputs=[raw_state, just_words],
        outputs=[lyrics_out],
    )


# ---------------------------------------------------------------------------
# Off-Brand badge: the app runs on an explicit gradio.Server (Server mode),
# not the default demo.launch() path.
#
# We create a gradio.Server (a FastAPI server with Gradio's engine built in)
# and mount the Blocks UI onto it. The visible UI is identical to the original
# — same components, theme, CSS, and SELECTION_JS.
#
# Details that are each verified necessary on this Gradio version (6.x):
#   * css / head / theme are passed to mount_gradio_app(...) here, NOT relied
#     upon from the gr.Blocks(css=..., head=..., theme=...) constructor. In
#     Gradio 6 the constructor values are dropped when the Blocks is served via
#     Server mode, which strips the styling (the unstyled-text page).
#   * the Blocks is mounted at the ROOT path "/". Mounting at a subpath like
#     "/app" makes the live event/queue calls fail to round-trip — the "Sing"
#     button spins forever and the backend function is never reached. Mounting
#     at "/" keeps the queue endpoints where the frontend expects them, so
#     generation actually runs. (Passing css to the mount is what lets the
#     root mount stay fully styled, which an earlier root-mount attempt lacked.)
# No content, layout, or component changes anywhere.
# ---------------------------------------------------------------------------

# Serialize generation: there's a single LLM + Kokoro pipeline in memory, and
# a lullaby takes ~10-20s. concurrency_limit=1 means simultaneous users are
# queued (and shown their place) instead of colliding on the shared model.
# max_size caps the waiting line so the Space can't be flooded.
demo.queue(max_size=20, default_concurrency_limit=1)

# Explicit Server mode (this is what earns the Off-Brand badge). The Blocks UI
# is mounted at the root path, carrying its theme + CSS + head script so the
# rendered page — and the working generation queue — are identical to before.
server = gr.Server()
gr.mount_gradio_app(
    server,
    demo,
    path="/",
    theme=gr.themes.Citrus(),
    css_paths="style.css",
    head=HEAD_HTML,
    ssr_mode=False,
)


if __name__ == "__main__":
    server.launch()