"""
Lulluna — AI Bedtime Story Weaver
===================================
Generates personalized bedtime stories for children using:
  • MiniCPM5-1B via llama.cpp  (story generation, 1B params)
  • Kokoro-82M TTS             (narration)
  • SD Turbo ~860M             (illustration)

All models run locally — no cloud APIs.
Built for the HuggingFace Build Small Hackathon 2026.
Track: Backyard AI  |  Badges: 🦙 Llama Champion · 🔌 Off the Grid
"""

import os
import json
import re
import subprocess
import time
import logging
from pathlib import Path
from typing import Optional

import gradio as gr
from dotenv import load_dotenv

from prompts import build_system_prompt, build_user_prompt
from engine import StoryEngine

load_dotenv()
logging.basicConfig(level=logging.INFO, format="%(levelname)s  %(message)s")
log = logging.getLogger("lulluna")

# ── Boot the engine once ──────────────────────────────────────
engine = StoryEngine()
engine.load()

# ── Narration (Kokoro-82M) ────────────────────────────────────
# HF Spaces (unified env):  imports narrate.py directly — no subprocess needed.
# Local dev (separate .venv_tts): falls back to subprocess automatically.

_TTS_PYTHON        = Path(__file__).parent / ".venv_tts/bin/python3.11"
_NARRATE_SCRIPT    = Path(__file__).parent / "narrate.py"
_ILLUSTRATE_SCRIPT = Path(__file__).parent / "illustrate.py"

# Lazy-import sentinels: None = not tried yet, False = not available, module = loaded.
_narrate_mod  = None
_illus_mod    = None


def _run_narration(text: str) -> Optional[dict]:
    """
    Narrate *text*. Tries direct import first (HF Spaces unified env);
    falls back to subprocess (local dev with separate .venv_tts).
    """
    global _narrate_mod
    if _narrate_mod is None:
        try:
            import narrate as _m
            _narrate_mod = _m
            log.info("Narration mode: direct import (unified env)")
        except (ImportError, ModuleNotFoundError):
            _narrate_mod = False
            log.info("Narration mode: subprocess (.venv_tts)")

    if _narrate_mod:
        try:
            return _narrate_mod.narrate(text)
        except Exception as e:
            log.error(f"Narration (inline) failed: {e}")
            return None

    # Subprocess fallback
    if not _TTS_PYTHON.exists():
        log.error("TTS venv not found — run setup_m5.sh or install kokoro in the main env")
        return None
    try:
        result = subprocess.run(
            [str(_TTS_PYTHON), str(_NARRATE_SCRIPT)],
            input=text,
            capture_output=True,
            text=True,
            timeout=120,
        )
        if result.returncode != 0:
            log.error(f"Narration subprocess failed: {result.stderr.strip()}")
            return None
        return json.loads(result.stdout)
    except subprocess.TimeoutExpired:
        log.error("Narration timed out (120 s)")
        return None
    except json.JSONDecodeError as e:
        log.error(f"Narration subprocess returned invalid JSON: {e}")
        return None


def narrate_story(text: str) -> str:
    """API endpoint handler — returns JSON string {audio: base64, mime: audio/wav}."""
    result = _run_narration(text)
    if result is None:
        raise gr.Error("Narration failed — check the backend logs.")
    return json.dumps(result)


# ── Illustration (SD Turbo ~860 M) ───────────────────────────

_ILLUS_PROMPT_SYSTEM = """You extract a short visual scene description from a children's bedtime story.
Output ONLY a single short sentence (max 12 words) describing the most visually interesting moment or character in the story.
Focus on animals, magical objects, nature, or the main character. No dialogue, no abstract concepts."""

def _extract_image_prompt(title: str, body: str) -> str:
    """Use MiniCPM to extract a concise visual scene for image generation."""
    if not engine.ready:
        return title
    try:
        snippet = body[:600]
        raw = engine.generate(
            _ILLUS_PROMPT_SYSTEM,
            f"Story title: {title}\n\nOpening:\n{snippet}\n\nVisual scene:",
        )
        prompt = raw.strip().split("\n")[0].split(".")[0].strip()
        return prompt if len(prompt) > 5 else title
    except Exception as e:
        log.warning(f"Prompt extraction failed: {e}")
        return title


def _run_illustration(prompt: str) -> Optional[dict]:
    """
    Generate an illustration for *prompt*. Tries direct import first (HF Spaces);
    falls back to subprocess (local dev with separate .venv_tts).
    """
    global _illus_mod
    if _illus_mod is None:
        try:
            import illustrate as _m
            _illus_mod = _m
            log.info("Illustration mode: direct import (unified env)")
        except (ImportError, ModuleNotFoundError):
            _illus_mod = False
            log.info("Illustration mode: subprocess (.venv_tts)")

    if _illus_mod:
        try:
            return _illus_mod.generate(prompt)
        except Exception as e:
            log.error(f"Illustration (inline) failed: {e}")
            return None

    # Subprocess fallback
    if not _TTS_PYTHON.exists():
        log.error("TTS venv not found")
        return None
    try:
        result = subprocess.run(
            [str(_TTS_PYTHON), str(_ILLUSTRATE_SCRIPT)],
            input=json.dumps({"prompt": prompt}),
            capture_output=True,
            text=True,
            timeout=180,
        )
        if result.returncode != 0:
            log.error(f"Illustration subprocess failed: {result.stderr.strip()[:300]}")
            return None
        return json.loads(result.stdout)
    except subprocess.TimeoutExpired:
        log.error("Illustration timed out (180 s)")
        return None
    except json.JSONDecodeError as e:
        log.error(f"Illustration subprocess returned invalid JSON: {e}")
        return None


def illustrate_story(title: str, body: str) -> str:
    """API endpoint — returns JSON string {image: base64, mime: image/png}."""
    prompt = _extract_image_prompt(title, body)
    log.info(f"Illustration prompt: {prompt!r}")
    result = _run_illustration(prompt)
    if result is None:
        raise gr.Error("Illustration failed — check the backend logs.")
    return json.dumps(result)


# ── Core generation function ──────────────────────────────────

def generate_story(
    age: str,
    value: str,
    tradition: str,
    length: str,
    child_name: str,
    interests: str,
) -> dict:
    """
    Called by the Gradio API.
    Returns a dict {title, emoji, body} — same shape as Lovable AI Gateway.
    """
    t0 = time.time()
    log.info(f"generate | age={age} value={value} tradition={tradition} length={length}")

    system = build_system_prompt(age, value, tradition, length, child_name, interests)
    user   = "Write tonight's bedtime story."

    raw = engine.generate(system, user)
    log.info(f"Raw output ({len(raw)} chars) in {time.time()-t0:.1f}s")

    result = parse_story_output(raw)
    log.info(f"Parsed: title='{result['title'][:40]}' emoji={result['emoji']} body={len(result['body'])} chars")
    return result


def parse_story_output(raw: str) -> dict:
    """
    Parse model output into {title, emoji, body}.

    Tries 3 strategies in order:
    1. JSON block  ```json { ... } ```
    2. Labelled sections  TITLE: / EMOJI: / STORY:
    3. Heuristic fallback — first line = title, rest = body
    """

    # Strategy 1 — JSON block
    json_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw, re.DOTALL)
    if not json_match:
        json_match = re.search(r"\{[^{}]*\"title\"[^{}]*\"body\"[^{}]*\}", raw, re.DOTALL)
    if json_match:
        try:
            data = json.loads(json_match.group(1) if "```" in raw else json_match.group(0))
            if data.get("title") and data.get("body"):
                return {
                    "title": str(data["title"])[:120],
                    "emoji": str(data.get("emoji", "🌙"))[:4],
                    "body":  str(data["body"]),
                }
        except json.JSONDecodeError:
            pass

    # Strategy 2 — Labelled sections
    title_m = re.search(r"(?:TITLE|Title):\s*(.+)", raw)
    emoji_m = re.search(r"(?:EMOJI|Emoji):\s*(\S+)", raw)
    body_m  = re.search(r"(?:STORY|Story|BODY|Body):\s*([\s\S]+)", raw)
    if title_m and body_m:
        return {
            "title": title_m.group(1).strip()[:120],
            "emoji": emoji_m.group(1).strip()[:4] if emoji_m else "🌙",
            "body":  body_m.group(1).strip(),
        }

    # Strategy 3 — Heuristic fallback
    lines = [l for l in raw.strip().split("\n") if l.strip()]
    if not lines:
        return {"title": "A Bedtime Story", "emoji": "🌙", "body": raw.strip()}

    title = lines[0].strip().lstrip("#").strip()[:120]
    body_lines = lines[1:]
    body = "\n\n".join(
        " ".join(para).strip()
        for para in _group_paragraphs(body_lines)
        if " ".join(para).strip()
    )
    return {
        "title": title,
        "emoji": _pick_emoji(title),
        "body":  body or raw.strip(),
    }


def _group_paragraphs(lines: list[str]) -> list[list[str]]:
    """Group lines into paragraphs split by blank lines."""
    paragraphs, current = [], []
    for line in lines:
        if line.strip():
            current.append(line.strip())
        else:
            if current:
                paragraphs.append(current)
                current = []
    if current:
        paragraphs.append(current)
    return paragraphs


def _pick_emoji(title: str) -> str:
    """Very rough emoji picker based on title keywords."""
    t = title.lower()
    mapping = {
        "lion": "🦁", "mouse": "🐭", "elephant": "🐘", "rabbit": "🐰",
        "fox": "🦊", "wolf": "🐺", "bear": "🐻", "bird": "🐦",
        "crane": "🕊️", "monkey": "🐒", "turtle": "🐢", "snake": "🐍",
        "moon": "🌙", "sun": "☀️", "star": "⭐", "sky": "🌌",
        "river": "🌊", "forest": "🌿", "flower": "🌸", "tree": "🌳",
        "spider": "🕷️", "ant": "🐜", "bee": "🐝",
        "king": "👑", "queen": "👑", "prince": "🤴", "princess": "👸",
        "magic": "✨", "wish": "🌠", "dream": "💫",
    }
    for word, emoji in mapping.items():
        if word in t:
            return emoji
    return "🌙"


# ── UI helper functions ───────────────────────────────────────

def ui_generate(age, value, tradition, length, child_name, interests):
    """
    Gradio UI handler. Returns (story_markdown, story_state, clear_audio, clear_image).
    Clears audio and image on each new generation so stale results don't persist.
    """
    try:
        result = generate_story(age, value, tradition, length, child_name, interests)
        story_md = f"# {result['emoji']} {result['title']}\n\n{result['body']}"
        return story_md, result, None, None
    except Exception as e:
        log.error(f"UI generate error: {e}")
        return f"*Error generating story: {e}*", {}, None, None


def ui_narrate_from_state(story_state: dict) -> Optional[str]:
    """Narrate the story from state dict. Returns a temp WAV file path."""
    import base64
    import tempfile
    if not story_state or not story_state.get("body"):
        gr.Warning("Generate a story first, then click Narrate.")
        return None
    text = f"{story_state.get('title', '')}\n\n{story_state['body']}"
    result = _run_narration_subprocess(text)
    if not result:
        gr.Warning("Narration failed — check the backend logs.")
        return None
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False, prefix="lulluna_ui_", dir="/tmp")
    tmp.write(base64.b64decode(result["audio"]))
    tmp.close()
    return tmp.name


def ui_illustrate_from_state(story_state: dict):
    """Generate illustration from state dict. Returns a PIL Image."""
    import base64
    import io
    from PIL import Image as PILImage
    if not story_state or not story_state.get("title"):
        gr.Warning("Generate a story first, then click Illustrate.")
        return None
    try:
        json_str = illustrate_story(story_state["title"], story_state.get("body", ""))
        data = json.loads(json_str)
        img_bytes = base64.b64decode(data["image"])
        return PILImage.open(io.BytesIO(img_bytes))
    except gr.Error:
        raise
    except Exception as e:
        log.error(f"UI illustration failed: {e}")
        gr.Warning("Illustration failed — check the backend logs.")
        return None


def api_endpoint(age, value, tradition, length, child_name, interests):
    """Named API endpoint — called by the React frontend."""
    return generate_story(age, value, tradition, length, child_name, interests)


# ── Constants ─────────────────────────────────────────────────

TRADITIONS = [
    "Any", "Aesop", "Panchatantra", "Japanese", "African",
    "Norse", "Native American", "Sufi", "Celtic", "Arabian",
    "Chinese", "Grimm", "Jataka",
]
VALUES = [
    "Kindness", "Courage", "Love", "Honesty", "Patience",
    "Generosity", "Friendship", "Perseverance", "Wisdom", "Humility",
]
AGES = ["3-4", "5-6", "7-8", "9-10"]
LENGTHS = [
    ("Short  (~2 min)", "short"),
    ("Medium (~5 min)", "medium"),
    ("Long   (~8 min)", "long"),
]


# ── Theme & CSS ───────────────────────────────────────────────

LULLUNA_THEME = gr.themes.Soft(
    primary_hue="violet",
    secondary_hue="blue",
    neutral_hue="slate",
)

NIGHT_CSS = """
/* ── Night sky background ─────────────────────────────────── */
gradio-app,
.gradio-container {
    background: linear-gradient(170deg, #060614 0%, #0c1230 55%, #160b2e 100%) !important;
    min-height: 100vh;
}

/* Blocks / panels */
.block, .gr-block, .svelte-vt3r6s {
    background: rgba(13, 18, 52, 0.88) !important;
    border-color: rgba(80, 60, 150, 0.3) !important;
}

/* Input fields */
input[type="text"],
textarea,
.svelte-1gfkn6j {
    background: #0c1030 !important;
    border-color: #2a3560 !important;
    color: #e0d0f0 !important;
}
input[type="text"]::placeholder,
textarea::placeholder {
    color: #6858a0 !important;
}

/* Labels */
label > span,
.label-wrap > span {
    color: #9888c0 !important;
}

/* Tab nav */
.tab-nav > button {
    color: #9080b8 !important;
    font-size: 0.97rem !important;
}
.tab-nav > button.selected {
    color: #d8c8f8 !important;
    border-bottom-color: #7c3aed !important;
}

/* ── Story display ────────────────────────────────────────── */
#story-display {
    min-height: 290px;
    max-height: 480px;
    overflow-y: auto;
    padding: 0.25rem 0.5rem;
}
#story-display::-webkit-scrollbar { width: 4px; }
#story-display::-webkit-scrollbar-track { background: rgba(13,18,52,0.5); }
#story-display::-webkit-scrollbar-thumb {
    background: rgba(109, 40, 217, 0.5);
    border-radius: 2px;
}

#story-display .prose,
#story-display [class*="markdown"] {
    font-family: Georgia, 'Palatino Linotype', 'Book Antiqua', serif !important;
    font-size: 1.07rem !important;
    line-height: 1.88 !important;
    color: #f0e8d8 !important;
}
#story-display .prose h1,
#story-display [class*="markdown"] h1 {
    font-size: 1.45rem !important;
    font-weight: 700 !important;
    text-align: center !important;
    color: #d8c8f8 !important;
    margin-bottom: 1rem !important;
    padding-bottom: 0.5rem !important;
    border-bottom: 1px solid rgba(109, 40, 217, 0.3) !important;
}
#story-display .prose p,
#story-display [class*="markdown"] p {
    margin-bottom: 0.8em !important;
    text-indent: 1.4em !important;
    color: #ede5d5 !important;
}
#story-display .prose p:first-of-type,
#story-display [class*="markdown"] p:first-of-type {
    text-indent: 0 !important;
}

/* ── Generate button — violet glow ───────────────────────── */
#generate-btn > button {
    background: linear-gradient(135deg, #6d28d9 0%, #4338ca 100%) !important;
    box-shadow: 0 4px 20px rgba(109, 40, 217, 0.4) !important;
    border: 0 !important;
    font-size: 1.02rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.01em !important;
    border-radius: 12px !important;
    width: 100% !important;
    padding: 0.7rem 1rem !important;
    color: #ffffff !important;
    transition: all 0.2s ease !important;
}
#generate-btn > button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 32px rgba(109, 40, 217, 0.65) !important;
}
#generate-btn > button:active {
    transform: translateY(0) !important;
}

/* ── Secondary action buttons ────────────────────────────── */
#narrate-btn > button,
#illustrate-btn > button {
    background: rgba(20, 15, 60, 0.8) !important;
    border: 1px solid rgba(80, 60, 160, 0.45) !important;
    color: #c8b8e8 !important;
    border-radius: 10px !important;
    font-size: 0.95rem !important;
    transition: all 0.15s ease !important;
}
#narrate-btn > button:hover,
#illustrate-btn > button:hover {
    background: rgba(40, 30, 90, 0.9) !important;
    border-color: rgba(109, 40, 217, 0.65) !important;
    color: #e8d8ff !important;
}

/* ── About tab typography ────────────────────────────────── */
.about-content .prose h2,
.about-content [class*="markdown"] h2 {
    color: #d0c0f0 !important;
    border-bottom: 1px solid rgba(80,60,150,0.4) !important;
    padding-bottom: 0.3rem !important;
    margin-top: 1.5rem !important;
}
.about-content .prose table,
.about-content [class*="markdown"] table {
    border-color: rgba(80,60,150,0.35) !important;
}
.about-content .prose th,
.about-content [class*="markdown"] th {
    background: rgba(40,20,80,0.6) !important;
    color: #c8b8e8 !important;
}
.about-content .prose td,
.about-content [class*="markdown"] td {
    color: #ddd0f0 !important;
    border-color: rgba(80,60,150,0.25) !important;
}

/* ── Hide Gradio footer ──────────────────────────────────── */
footer { display: none !important; }
"""

# ── Hero HTML header ──────────────────────────────────────────

HERO_HTML = """
<div style="
    text-align: center;
    padding: 2rem 1rem 1.5rem;
    background: linear-gradient(180deg, rgba(30,12,60,0.85) 0%, rgba(13,18,48,0.5) 100%);
    border-radius: 18px;
    margin-bottom: 0.25rem;
    border: 1px solid rgba(109, 40, 217, 0.22);
">
    <div style="font-size:3rem; line-height:1; margin-bottom:0.55rem;">🌙</div>
    <h1 style="
        font-size: 2.5rem;
        font-weight: 800;
        margin: 0;
        letter-spacing: -0.025em;
        background: linear-gradient(135deg, #d8c8f8 0%, #a8b8ff 50%, #f8d8c0 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    ">Lulluna</h1>
    <p style="color:#9880c8; font-size:1.05rem; margin:0.4rem 0 1.1rem; font-style:italic; font-family:Georgia,serif;">
        Bedtime stories, woven just for your child
    </p>
    <div style="display:flex; justify-content:center; gap:8px; flex-wrap:wrap;">
        <span style="
            display:inline-flex; align-items:center; gap:5px;
            padding:4px 13px; border-radius:20px; font-size:0.8rem;
            background:rgba(109,40,217,0.22); border:1px solid rgba(109,40,217,0.5);
            color:#c8b4f8;
        ">🦙 llama.cpp · MiniCPM5-1B</span>
        <span style="
            display:inline-flex; align-items:center; gap:5px;
            padding:4px 13px; border-radius:20px; font-size:0.8rem;
            background:rgba(16,110,55,0.22); border:1px solid rgba(30,170,85,0.4);
            color:#88f0a8;
        ">🔌 100% local · zero cloud APIs</span>
        <span style="
            display:inline-flex; align-items:center; gap:5px;
            padding:4px 13px; border-radius:20px; font-size:0.8rem;
            background:rgba(150,85,15,0.2); border:1px solid rgba(195,145,25,0.45);
            color:#f0d08a;
        ">🏆 HF Build Small Hackathon 2026</span>
    </div>
</div>
"""

# ── About tab content ─────────────────────────────────────────

ABOUT_CONTENT = """
## 🌙 What is Lulluna?

Lulluna solves a nightly challenge for busy parents: **"What story do I read tonight?"**

In about 10 seconds, it weaves a personalized bedtime story tailored to your child's **age**, **name**, and **interests** — drawing from a rich tradition (Aesop, Panchatantra, Norse, African, Japanese...) and quietly threading in a value like kindness or courage. Every story ends with a soft, tucking-in line that helps little ones drift off to sleep.

---

## 🤖 AI Stack — 100% Local, Zero Cloud APIs

| Component | Model | Params | Runtime |
|---|---|---|---|
| Story generation | MiniCPM5-1B (OpenBMB) | 1.0 B | llama.cpp + Metal (Apple Silicon) |
| Voice narration | Kokoro-82M (hexgrad) | 82 M | PyTorch |
| Illustration | SD Turbo (Stability AI) | ~860 M | Diffusers (MPS) |
| **Total** | | **~1.96 B** | All models ≤ 32 B ✓ |

Every story is generated, narrated, and illustrated on-device. No API keys. No data leaves your machine.

---

## 🏅 Hackathon Badge Claims

| Badge | How we earn it |
|---|---|
| 🦙 **Llama Champion** | Story engine uses llama.cpp with GGUF Q4_K_M quantization |
| 🔌 **Off the Grid** | Zero cloud API calls — all three models run entirely locally |
| 📓 **Field Notes** | Blog post documenting the build ([read it →](https://huggingface.co/blog/lulluna-build-small)) |

**Track:** Backyard AI — solving a real daily pain point for parents everywhere.

---

## 🔌 REST API (companion React frontend)

```
POST /api/generate
  {"data": [age, value, tradition, length, name, interests]}
  → {"data": [{"title": "...", "emoji": "...", "body": "..."}]}

POST /api/narrate
  {"data": ["story text here"]}
  → {"data": ["{\"audio\":\"<base64 wav>\",\"mime\":\"audio/wav\"}"]}

POST /api/illustrate
  {"data": ["title", "body"]}
  → {"data": ["{\"image\":\"<base64 png>\",\"mime\":\"image/png\"}"]}
```
"""

# ── Gradio interface ──────────────────────────────────────────

with gr.Blocks(
    title="Lulluna — Bedtime Story Weaver",
    theme=LULLUNA_THEME,
    css=NIGHT_CSS,
) as demo:

    story_state = gr.State({})

    gr.HTML(HERO_HTML)

    with gr.Tabs():

        # ── Tab 1: Story creator ──────────────────────────────
        with gr.Tab("✨ Create a Story"):
            with gr.Row(equal_height=False):

                # Left — settings form
                with gr.Column(scale=1, min_width=260):
                    name_in = gr.Textbox(
                        label="Child's name",
                        placeholder="e.g. Aria, Leo, Maya",
                    )
                    with gr.Row():
                        age_in = gr.Dropdown(
                            choices=AGES,
                            value="5-6",
                            label="Age range",
                            scale=1,
                        )
                        length_in = gr.Dropdown(
                            choices=LENGTHS,
                            value="short",
                            label="Length",
                            scale=1,
                        )
                    value_in = gr.Dropdown(
                        choices=VALUES,
                        value="Kindness",
                        label="Value / theme",
                    )
                    tradition_in = gr.Dropdown(
                        choices=TRADITIONS,
                        value="Any",
                        label="Tradition / setting",
                    )
                    interests_in = gr.Textbox(
                        label="Interests (optional)",
                        placeholder="e.g. baby elephants, the ocean, dinosaurs",
                        lines=2,
                    )
                    generate_btn = gr.Button(
                        "✨ Weave tonight's story",
                        variant="primary",
                        elem_id="generate-btn",
                    )
                    gr.Markdown(
                        "<small>*First generation ~10 s on M-series Mac.*</small>"
                    )

                # Right — story output
                with gr.Column(scale=2, min_width=380):
                    story_out = gr.Markdown(
                        value=(
                            "> *Fill in the settings on the left and click "
                            "**✨ Weave tonight's story** to begin.*"
                        ),
                        elem_id="story-display",
                    )
                    with gr.Row():
                        narrate_btn = gr.Button(
                            "🎙️ Narrate",
                            variant="secondary",
                            scale=1,
                            elem_id="narrate-btn",
                        )
                        illustrate_btn = gr.Button(
                            "🎨 Illustrate",
                            variant="secondary",
                            scale=1,
                            elem_id="illustrate-btn",
                        )
                    audio_out = gr.Audio(
                        label="🎙️ Narration — Kokoro-82M",
                        type="filepath",
                        autoplay=True,
                    )
                    image_out = gr.Image(
                        label="🎨 Illustration — SD Turbo",
                        type="pil",
                        show_download_button=True,
                    )

        # ── Tab 2: About ──────────────────────────────────────
        with gr.Tab("📖 About Lulluna"):
            with gr.Column(elem_classes=["about-content"]):
                gr.Markdown(ABOUT_CONTENT)

    # ── Event wiring ─────────────────────────────────────────
    generate_btn.click(
        fn=ui_generate,
        inputs=[age_in, value_in, tradition_in, length_in, name_in, interests_in],
        outputs=[story_out, story_state, audio_out, image_out],
    )
    narrate_btn.click(
        fn=ui_narrate_from_state,
        inputs=[story_state],
        outputs=[audio_out],
    )
    illustrate_btn.click(
        fn=ui_illustrate_from_state,
        inputs=[story_state],
        outputs=[image_out],
    )

    # ── Hidden API endpoints (called by the React frontend) ──
    # POST /api/generate
    gr.Interface(
        fn=api_endpoint,
        inputs=[
            gr.Textbox(visible=False),
            gr.Textbox(visible=False),
            gr.Textbox(visible=False),
            gr.Textbox(visible=False),
            gr.Textbox(visible=False),
            gr.Textbox(visible=False),
        ],
        outputs=gr.JSON(visible=False),
        api_name="generate",
    )
    # POST /api/narrate
    gr.Interface(
        fn=narrate_story,
        inputs=gr.Textbox(visible=False),
        outputs=gr.Textbox(visible=False),
        api_name="narrate",
    )
    # POST /api/illustrate
    gr.Interface(
        fn=illustrate_story,
        inputs=[gr.Textbox(visible=False), gr.Textbox(visible=False)],
        outputs=gr.Textbox(visible=False),
        api_name="illustrate",
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        share=False,
    )