"""Tonight's Tale — interactive illustrated bedtime stories from your child's day.

HF "Build Small" hackathon entry (Backyard AI track).
Small models doing a big job: Qwen3.5-4B (story) + Z-Image-Turbo 6B (pictures)
+ Kokoro-82M (narration) ≈ 10B parameters total, well under the 32B limit.
"""

from __future__ import annotations

import random
import time

import gradio as gr

import imagegen
import story as st
import tts

try:  # `spaces` exists on Hugging Face Spaces; locally we run in mock mode without it
    import spaces

    GPU = spaces.GPU
except ImportError:  # pragma: no cover

    def GPU(*args, **kwargs):  # noqa: N802 - mirrors spaces.GPU
        if args and callable(args[0]):
            return args[0]
        return lambda fn: fn


# ONE GPU function for everything: multiple @spaces.GPU functions mean multiple
# GPU workers each packing all 28.9GB of tensors — the second worker then fails
# to attach a GPU ('No CUDA GPUs are available'). A single entry point keeps one
# reusable worker across all events.
@GPU(duration=90)
def gpu_generate(ctx: st.PromptContext, cast: list, seed: int):
    node = st.generate_story_node(ctx)
    prompt = imagegen.build_illustration_prompt(cast, node.illustration)
    image = imagegen.generate_illustration(prompt, seed)
    return node.model_dump(), image


# Load model weights at startup so downloads never eat into the ZeroGPU budget.
if not st.MOCK:
    st.preload()
    imagegen.preload()


# ---------------------------------------------------------------------------
# Branching state machine (per-session, held in gr.State)
# ---------------------------------------------------------------------------

MAX_DEPTH = st.MAX_DEPTH


def _cast_from_state(state) -> list[st.CastMember]:
    return [st.CastMember(**c) for c in state["cast"]]


def _build_ctx(state, path: str) -> st.PromptContext:
    """Context to generate the node AT `path` ('' = opening, 'A' = after choosing A, ...)."""
    steps: list[st.StoryStep] = []
    for i in range(len(path)):
        prefix = path[:i]
        node = state["nodes"][prefix]
        choice = node["choice"]
        chosen = choice["optionA"]["label"] if path[i] == "A" else choice["optionB"]["label"]
        steps.append(st.StoryStep(paragraphs=node["paragraphs"], choice_prompt=choice["prompt"], chosen_label=chosen))

    if path == "":
        instruction = st.NodeInstruction(kind="opening", remaining_choice_points=MAX_DEPTH - 1)
    else:
        parent = state["nodes"][path[:-1]]
        chosen = parent["choice"]["optionA"]["label"] if path[-1] == "A" else parent["choice"]["optionB"]["label"]
        sibling_key = path[:-1] + ("B" if path[-1] == "A" else "A")
        sibling = state["nodes"].get(sibling_key)
        instruction = st.NodeInstruction(
            kind="continuation",
            chosen_label=chosen,
            must_end=len(path) >= MAX_DEPTH,
            is_replay_after_setback=bool(sibling and sibling["isEnding"] and sibling["endingTone"] == "gentle_setback"),
            remaining_choice_points=max(0, MAX_DEPTH - len(path) - 1),
        )
    return st.PromptContext(
        child_name=state["child_name"],
        age_band=state["age_band"],
        language=state["language"],
        cast=_cast_from_state(state),
        day_event=state["day_event"],
        story_so_far=steps,
        instruction=instruction,
    )


def _is_transient_gpu_error(err: Exception) -> bool:
    """ZeroGPU worker-attach failures surface as gr.Error('RuntimeError') —
    the spaces library swallows the worker traceback and re-raises with just
    the class name. Match broadly on infra-looking errors; our own generation
    failures use dedicated exception classes with descriptive messages."""
    text = f"{type(err).__name__}: {err}"
    return any(token in text for token in ("CUDA", "GPU", "RuntimeError", "ZeroGPU"))


def _gpu_retry(fn, *args, attempts: int = 3, wait: float = 25.0):
    """Worker-attach failures happen before any generation ran, so retrying is safe."""
    for i in range(attempts):
        try:
            return fn(*args)
        except st.StoryGenerationError:
            raise
        except Exception as err:  # noqa: BLE001 - see _is_transient_gpu_error
            if not _is_transient_gpu_error(err) or i == attempts - 1:
                raise
            print(f"[gpu_retry] attempt {i + 1} failed ({type(err).__name__}: {err}); retrying in {wait}s")
            time.sleep(wait)


def _ensure_node(state, path: str):
    """Lazy generation with branch reuse: already-generated branches replay instantly."""
    if path in state["nodes"]:
        return state["nodes"][path]
    node, image = _gpu_retry(gpu_generate, _build_ctx(state, path), _cast_from_state(state), state["seed"])
    state["nodes"][path] = node
    state["images"][path] = image
    return node


def _story_markdown(state) -> str:
    path = state["path"]
    lines = []
    title = state["nodes"][""].get("title")
    if title:
        lines.append(f"# 🌙 {title}\n")
    for i in range(len(path) + 1):
        prefix = path[:i]
        node = state["nodes"].get(prefix)
        if not node:
            continue
        for p in node["paragraphs"]:
            lines.append(p + "\n")
        if i < len(path) and node["choice"]:
            chosen = node["choice"]["optionA"]["label"] if path[i] == "A" else node["choice"]["optionB"]["label"]
            lines.append(f"> 💭 *{node['choice']['prompt']}* → **{chosen}**\n")
    current = state["nodes"][path]
    if not current["isEnding"] and current["choice"]:
        lines.append(f"### 💭 {current['choice']['prompt']}")
    return "\n".join(lines)


def _render(state):
    """Maps state -> UI component updates."""
    path = state["path"]
    node = state["nodes"][path]
    image = state["images"].get(path)
    audio = tts.synthesize(" ".join(node["paragraphs"]), state["language"])
    tts_note = "" if audio or state["language"] == "de" else ""
    if state["language"] == "de":
        tts_note = "*(narration not available for German yet — read it aloud together!)*"

    is_choice = not node["isEnding"]
    is_setback = node["isEnding"] and node["endingTone"] == "gentle_setback"
    is_warm = node["isEnding"] and node["endingTone"] == "warm"
    disabled = state.get("disabled_option")

    btn_a = btn_b = gr.Button(visible=False)
    if is_choice:
        choice = node["choice"]
        btn_a = gr.Button(value=f"🅰 {choice['optionA']['label']}", visible=True, interactive=disabled != "A")
        btn_b = gr.Button(value=f"🅱 {choice['optionB']['label']}", visible=True, interactive=disabled != "B")

    setback_md = (
        "### 🌧️ The story ended with a little disappointment...\n"
        "Ask your child softly: *“What do you think would have happened with the other choice?”*"
        if is_setback
        else ""
    )
    moral = node.get("moral")
    finished_md = (
        f"### 🌙 The end — sweet dreams!\n\n**Note for the parent:** {moral}" if is_warm and moral else ("### 🌙 The end — sweet dreams!" if is_warm else "")
    )

    return (
        state,
        gr.Group(visible=False),  # setup panel
        gr.Group(visible=True),  # story panel
        image,
        _story_markdown(state) + ("\n\n" + tts_note if tts_note else ""),
        audio,
        btn_a,
        btn_b,
        gr.Markdown(value=setback_md, visible=is_setback),
        gr.Button(visible=is_setback),  # replay button
        gr.Markdown(value=finished_md, visible=is_warm),
        gr.Button(visible=node["isEnding"]),  # new-story button
    )


# ---------------------------------------------------------------------------
# Event handlers
# ---------------------------------------------------------------------------


def _gpu_friendly(fn):
    """Translate ZeroGPU scheduling hiccups into a friendly toast instead of a raw traceback."""
    import functools

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except gr.Error:
            raise
        except RuntimeError as err:
            if "CUDA" in str(err) or "GPU" in str(err):
                raise gr.Error("The free GPU is busy right now — please try again in a minute 🌙") from err
            raise

    return wrapper


@_gpu_friendly
def start_story(child_name, age_band, language, c1_name, c1_kind, c1_traits, c1_look, c2_name, c2_kind, c2_traits, c2_look, day_event):
    if not day_event.strip():
        raise gr.Error("Tell us what happened today first 🙂")
    cast = [
        {"name": c1_name.strip(), "kind": c1_kind.strip(), "traits": c1_traits.strip(), "appearance": c1_look.strip()},
        {"name": c2_name.strip(), "kind": c2_kind.strip(), "traits": c2_traits.strip(), "appearance": c2_look.strip()},
    ]
    cast = [c for c in cast if c["name"]]
    state = {
        "child_name": child_name.strip() or "Mia",
        "age_band": age_band,
        "language": language,
        "cast": cast,
        "day_event": day_event.strip(),
        "seed": random.randint(1, 2_147_483_646),
        "nodes": {},
        "images": {},
        "path": "",
        "disabled_option": None,
    }
    _ensure_node(state, "")
    return _render(state)


@_gpu_friendly
def choose(state, option: str):
    node = state["nodes"][state["path"]]
    if node["isEnding"] or not node["choice"]:
        return _render(state)
    state["disabled_option"] = None
    state["path"] = state["path"] + option
    _ensure_node(state, state["path"])
    return _render(state)


def choose_a(state):
    return choose(state, "A")


def choose_b(state):
    return choose(state, "B")


def back_to_choice(state):
    if not state["path"]:
        return _render(state)
    taken = state["path"][-1]
    state["path"] = state["path"][:-1]
    state["disabled_option"] = taken
    return _render(state)


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

INTRO = """# 🌙 Tonight's Tale
**Turn what happened in your child's day into tonight's illustrated bedtime story.**

Tell the app what your kid went through today ("didn't want to share toys at kindergarten...").
A **4B story model** weaves it into a fairy tale starring your child's favorite characters — with
**choice points your child decides**, watercolor pictures from a **6B image model**, and narration
by an **82M voice model**. Wrong turns lead to a *gentle*, recoverable setback — then you rewind
together and try the other path. The lesson hides inside the plot, never preached.

*Small models (≈10B total), big bedtime adventure.*
"""

with gr.Blocks(title="Tonight's Tale", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)
    session = gr.State()

    with gr.Group(visible=True) as setup_panel:
        with gr.Row():
            child_name = gr.Textbox(label="Child's name (nickname is fine)", value="Mia")
            age_band = gr.Radio(["3-6", "6-9"], value="3-6", label="Age")
            language = gr.Dropdown(
                choices=[("中文", "zh"), ("English", "en"), ("Deutsch", "de"), ("Français", "fr")],
                value="en",
                label="Story language",
            )
        gr.Markdown("**The recurring cast** — your child's favorite characters (appearance powers the illustrations):")
        with gr.Row():
            c1_name = gr.Textbox(label="Name", value="Hop", scale=1)
            c1_kind = gr.Textbox(label="What is it?", value="little rabbit", scale=1)
            c1_traits = gr.Textbox(label="Personality", value="curious and energetic, a bit impatient", scale=2)
            c1_look = gr.Textbox(label="Appearance (English)", value="a small white rabbit with long floppy ears and a red scarf", scale=2)
        with gr.Row():
            c2_name = gr.Textbox(label="Name", value="Sage", scale=1)
            c2_kind = gr.Textbox(label="What is it?", value="old tortoise", scale=1)
            c2_traits = gr.Textbox(label="Personality", value="calm and wise, loves riddles", scale=2)
            c2_look = gr.Textbox(label="Appearance (English)", value="an old green tortoise with a mossy shell and tiny round glasses", scale=2)
        day_event = gr.Textbox(
            label="What happened today?",
            placeholder="e.g. Didn't want to share toys at kindergarten and quarreled with a best friend...",
            lines=3,
        )
        start_btn = gr.Button("✨ Weave tonight's story", variant="primary")

    with gr.Group(visible=False) as story_panel:
        illustration = gr.Image(label="Illustration", type="pil", interactive=False)
        story_md = gr.Markdown()
        narration = gr.Audio(label="Narration (Kokoro-82M)", autoplay=True, interactive=False)
        with gr.Row():
            btn_a = gr.Button(visible=False, variant="primary")
            btn_b = gr.Button(visible=False, variant="primary")
        setback_md = gr.Markdown(visible=False)
        replay_btn = gr.Button("⏪ Go back and try the other path", visible=False, variant="secondary")
        finished_md = gr.Markdown(visible=False)
        new_story_btn = gr.Button("🌅 New story", visible=False)

    outputs = [
        session, setup_panel, story_panel, illustration, story_md, narration,
        btn_a, btn_b, setback_md, replay_btn, finished_md, new_story_btn,
    ]

    start_btn.click(
        start_story,
        inputs=[child_name, age_band, language, c1_name, c1_kind, c1_traits, c1_look, c2_name, c2_kind, c2_traits, c2_look, day_event],
        outputs=outputs,
    )
    btn_a.click(choose_a, inputs=[session], outputs=outputs)
    btn_b.click(choose_b, inputs=[session], outputs=outputs)
    replay_btn.click(back_to_choice, inputs=[session], outputs=outputs)

    def reset():
        return (
            None,
            gr.Group(visible=True),
            gr.Group(visible=False),
            None, "", None,
            gr.Button(visible=False), gr.Button(visible=False),
            gr.Markdown(visible=False), gr.Button(visible=False),
            gr.Markdown(visible=False), gr.Button(visible=False),
        )

    new_story_btn.click(reset, outputs=outputs)

if __name__ == "__main__":
    demo.launch(show_error=True)