app.py

"""FrogQuest — plain Gradio (gr.Blocks) app, single-page master-detail layout.

Layout (one page, no onboarding gate):
  LEFT   — hero photo (click to upload/change), hero stats, world picker + reset
  CENTER — the selected quest's scene image, its description, and Done / Couldn't actions
  RIGHT  — the quest log: a selectable list of quests (selected one highlighted orange)
  BOTTOM — the "Frog Master" chat: one box that forges quests, adds tasks, and marks quests
           done / couldn't (the LLM classifies each message into an intent — see llm.route_intent)

State lives in gr.BrowserState (per-browser localStorage): the resized photo (base64), the
chosen world/theme, the validated adventure/quests JSON, the selected quest id, and a cache of
generated scene images (JPEG data-urls). The photo goes to the GPU only transiently during
generation and is never persisted server-side.

Images generate LAZILY: selecting a quest with no cached scene generates it on the spot; Done /
Couldn't EDIT the cached initial scene into a success / failure state (never regenerate).
"""
from __future__ import annotations

import base64
import html
import io
import os
import traceback

# Speed up all HF downloads (Nemotron GGUF + FLUX weights). Must precede any huggingface_hub
# import (diffusers/llama pull it in). hf-transfer is in requirements.txt.
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

import gradio as gr
import spaces
from PIL import Image

from schema import THEMES, merge_quests, validate_and_clamp


# ZeroGPU scans for a @spaces.GPU function at startup; this guarantees detection regardless
# of whether the heavy model modules import cleanly.
@spaces.GPU(duration=15)
def _warmup():
    return "ok"


# Defensive import: never let a model-import error hang the whole app.
MODEL_IMPORT_ERROR = None
try:
    from llm import generate_quests_raw, route_intent
    from images import edit_image, initial_image, pil_to_data_url
except Exception:
    MODEL_IMPORT_ERROR = traceback.format_exc()


# ----------------------------- photo / image helpers (PIL only) -----------------------------

PHOTO_MAX_SIDE = 512
IMAGES_CAP_BYTES = 3_500_000  # soft cap on the cached-scene blob to stay under localStorage ~5MB


def _resize_dataurl(pil: Image.Image, max_side: int = PHOTO_MAX_SIDE) -> str:
    img = pil.convert("RGB")
    w, h = img.size
    scale = min(1.0, max_side / max(w, h))
    if scale < 1.0:
        img = img.resize((round(w * scale), round(h * scale)))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=80)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def _dataurl_to_pil(data_url: str) -> Image.Image:
    s = data_url.split(",", 1)[1] if "," in data_url else data_url
    return Image.open(io.BytesIO(base64.b64decode(s))).convert("RGB")


def _import_error_message() -> str:
    last = (MODEL_IMPORT_ERROR or "").strip().splitlines()
    return "Model failed to load: " + (last[-1] if last else "unknown error")


def _trim_images(images: dict, keep_id) -> dict:
    """Bound the cached-scene blob so a BrowserState write stays under localStorage's ~5MB cap.
    Drops oldest non-selected `initial` variants first (keeping success/failure = completed
    progress), then whole entries if still over. Mutates and returns `images`."""
    def total() -> int:
        return sum(len(v) for c in images.values() for v in c.values())

    if total() <= IMAGES_CAP_BYTES:
        return images
    for qid in list(images.keys()):
        if total() <= IMAGES_CAP_BYTES:
            break
        if qid == keep_id:
            continue
        c = images[qid]
        if "initial" in c and ("success" in c or "failure" in c):
            del c["initial"]
    for qid in list(images.keys()):
        if total() <= IMAGES_CAP_BYTES:
            break
        if qid != keep_id:
            del images[qid]
    return images


# ----------------------------- state plumbing -----------------------------

def _default_state() -> dict:
    return {"photo": None, "theme": "fantasy", "adventure": None,
            "quests": [], "selected_id": None, "images": {}}


def _merge(browser: dict | None, **changes) -> dict:
    """Return the browser dict with `changes` applied (other keys preserved for persistence)."""
    b = dict(browser or _default_state())
    b.update(changes)
    return b


def _find(quests: list, qid) -> dict | None:
    return next((q for q in (quests or []) if q.get("id") == qid), None)


def _art_style(theme: str) -> str:
    return f"8-bit / 16-bit retro pixel art, {theme} palette, NES RPG style"


# ----------------------------- HTML builders -----------------------------

def _esc(s) -> str:
    return html.escape(str(s or ""))


def _badges_html(q: dict) -> str:
    badges = []
    if q.get("is_frog"):
        badges.append('<span class="badge frog-badge">🐸 THE FROG</span>')
    if q.get("type") == "bonus":
        badges.append('<span class="badge bonus-badge">✦ BONUS · OPTIONAL</span>')
    if q.get("goal_group"):
        badges.append(f'<span class="badge group-badge">⛓ {_esc(q["goal_group"])}</span>')
    return f'<div class="quest-badges">{"".join(badges)}</div>' if badges else ""


def _desc_html(quest: dict | None, adventure: dict | None, hint: str | None = None) -> str:
    adv = ""
    if adventure and adventure.get("title"):
        adv = f'<div class="adventure-header"><h2>{_esc(adventure["title"])}</h2></div>'
    if not quest:
        body = ('<div class="fq-empty">No quest selected yet.<br>'
                'Tell the Frog Master your plans below to forge your quest log.</div>')
        return f'<div class="fq-desc">{adv}{body}</div>'

    status = quest.get("status", "active")
    state_cls = {"success": " state-success", "failure": " state-failure"}.get(status, "")
    parts = [
        adv,
        _badges_html(quest),
        f'<h3 class="quest-title">{_esc(quest.get("quest_title"))}</h3>',
        f'<p class="quest-narrative">{_esc(quest.get("narrative"))}</p>',
        f'<p class="quest-task">{_esc(quest.get("task"))}</p>',
        f'<div class="quest-foot"><span class="xp">{int(quest.get("xp", 0))} XP</span></div>',
    ]
    if status in ("success", "failure") and quest.get("result_msg"):
        parts.append(f'<p class="result-msg">{_esc(quest["result_msg"])}</p>')
    elif hint:
        parts.append(f'<p class="result-msg">{_esc(hint)}</p>')
    return f'<div class="fq-desc{state_cls}">{"".join(parts)}</div>'


def _stats_html(quests: list) -> str:
    quests = quests or []
    total = len(quests)
    done = sum(1 for q in quests if q.get("status") == "success")
    failed = sum(1 for q in quests if q.get("status") == "failure")
    xp = sum(int(q.get("xp", 0)) for q in quests if q.get("status") == "success")
    return (
        '<div class="fq-stats"><h4>HERO STATS</h4>'
        f'<div class="stat-row"><span>QUESTS</span><b>{done}/{total}</b></div>'
        f'<div class="stat-row"><span>XP EARNED</span><b>{xp}</b></div>'
        f'<div class="stat-row"><span>RETREATS</span><b>{failed}</b></div>'
        '</div>'
    )


# ----------------------------- core actions (reused by clicks + chat) -----------------------------

def select_quest(qid, photo, adventure, images, quests, browser):
    """Select a quest: show its cached scene, or lazily generate the initial scene on first view.
    Returns (selected_id, scene_pil_or_None, desc_html, images, browser)."""
    images = dict(images or {})
    quest = _find(quests, qid)
    if quest is None:
        return qid, None, _desc_html(None, adventure), images, _merge(browser, selected_id=qid)

    cache = dict(images.get(qid) or {})
    want = quest.get("image_state", "initial")
    data = cache.get(want) or cache.get("initial")
    scene = None
    hint = None
    if data is not None:
        scene = _dataurl_to_pil(data)
    else:
        if MODEL_IMPORT_ERROR:
            raise gr.Error(_import_error_message())
        if not adventure:
            raise gr.Error("Forge a quest log first — tell the Frog Master your plans.")
        if not photo:
            hint = "Upload your photo (left) to draw this scene."
        else:
            pil = initial_image(_dataurl_to_pil(photo), adventure["art_style"],
                                quest.get("initial_image_prompt", ""), int(adventure["seed"]))
            cache["initial"] = pil_to_data_url(pil)
            images[qid] = cache
            images = _trim_images(images, qid)
            scene = pil
    return qid, scene, _desc_html(quest, adventure, hint), images, _merge(browser, selected_id=qid, images=images)


def _apply_result(qid, quests, adventure, photo, images, browser, kind, reason):
    """Mark a quest success/failure: EDIT its initial scene into the after-state and update text.
    Returns (quests, images, scene_pil, desc_html, stats_html, browser)."""
    if MODEL_IMPORT_ERROR:
        raise gr.Error(_import_error_message())
    quests = [dict(q) for q in (quests or [])]
    quest = _find(quests, qid)
    if quest is None:
        raise gr.Error("Select a quest first.")
    if not adventure:
        raise gr.Error("Forge a quest log first.")

    images = dict(images or {})
    cache = dict(images.get(qid) or {})
    if "initial" in cache:
        init_pil = _dataurl_to_pil(cache["initial"])
    else:
        if not photo:
            raise gr.Error("Upload your photo (left) and view the scene first.")
        init_pil = initial_image(_dataurl_to_pil(photo), adventure["art_style"],
                                 quest.get("initial_image_prompt", ""), int(adventure["seed"]))
        cache["initial"] = pil_to_data_url(init_pil)

    instruction = reason or (quest["success_edit"] if kind == "success" else quest["failure_edit"])
    edited = edit_image(init_pil, instruction, adventure["art_style"], int(adventure["seed"]))
    cache[kind] = pil_to_data_url(edited)
    images[qid] = cache
    images = _trim_images(images, qid)

    quest["status"] = kind
    quest["image_state"] = kind
    if kind == "success":
        quest["result_msg"] = f"⚔ VICTORY! +{quest['xp']} XP — {quest['quest_title']} conquered."
    else:
        msg = "🌙 Retreat for now — you'll face it another day. No shame in resting."
        if reason:
            msg = f"{reason.strip()}  ·  {msg}"
        quest["result_msg"] = msg

    desc = _desc_html(quest, adventure)
    return quests, images, edited, desc, _stats_html(quests), _merge(browser, quests=quests, images=images)


def apply_done(qid, quests, adventure, photo, images, browser):
    if qid is None:
        raise gr.Error("Select a quest first.")
    return _apply_result(qid, quests, adventure, photo, images, browser, "success", None)


def apply_couldnt(qid, quests, adventure, photo, images, browser):
    if qid is None:
        raise gr.Error("Select a quest first.")
    return _apply_result(qid, quests, adventure, photo, images, browser, "failure", None)


def remove_quest(qid, quests, adventure, images, selected_id, browser):
    """Drop a finished (done/failed) quest from the log and its cached scenes. If it was the
    selected quest, repoint selection to the first remaining quest (no regeneration).
    Returns (quests, images, selected_id, scene_pil, desc_html, stats_html, browser)."""
    quests = [q for q in (quests or []) if q.get("id") != qid]
    images = dict(images or {})
    images.pop(qid, None)
    if selected_id == qid:
        selected_id = quests[0]["id"] if quests else None

    quest = _find(quests, selected_id)
    scene = None
    if quest:
        cache = images.get(quest["id"]) or {}
        data = cache.get(quest.get("image_state", "initial")) or cache.get("initial")
        if data:
            scene = _dataurl_to_pil(data)
    return (quests, images, selected_id, scene, _desc_html(quest, adventure), _stats_html(quests),
            _merge(browser, quests=quests, images=images, selected_id=selected_id))


# ----------------------------- chat (Frog Master, full intent routing) -----------------------------

def _chat_context(quests: list, selected_id) -> str:
    if not quests:
        return "No quest log exists yet."
    lines = []
    for i, q in enumerate(quests):
        sel = " | SELECTED" if q.get("id") == selected_id else ""
        lines.append(f"{i + 1}. id={q.get('id')} | title={q.get('quest_title')} | "
                     f"task={q.get('task')} | status={q.get('status')}{sel}")
    return "A quest log already exists:\n" + "\n".join(lines)


def _resolve_target(quests: list, target: str, selected_id):
    t = (target or "").strip().lower()
    if t:
        for q in quests:
            if t in (str(q.get("id", "")).lower(), q.get("quest_title", "").lower(),
                     q.get("task", "").lower()):
                return q["id"]
        for q in quests:  # substring fallback
            if t in q.get("quest_title", "").lower() or t in q.get("task", "").lower():
                return q["id"]
    return selected_id


def _do_forge(message, theme, photo, browser):
    raw = generate_quests_raw(message, theme)
    adv = validate_and_clamp(raw, theme)
    quests, adventure = adv["quests"], adv["adventure"]
    frog_id = quests[0]["id"]
    _sel, scene, desc, images, _b = select_quest(frog_id, photo, adventure, {}, quests, browser)
    browser2 = _merge(browser, theme=theme, adventure=adventure, quests=quests,
                      selected_id=frog_id, images=images)
    return quests, adventure, frog_id, images, scene, desc, _stats_html(quests), browser2


# chat_send outputs (order): chat_input, quests_state, adventure_state, selected_id_state,
#                            images_state, scene_image, desc_html, stats_html, browser
def chat_send(message, quests, adventure, theme, photo, selected_id, images, browser):
    message = (message or "").strip()
    nochange = (gr.update(),) * 8  # everything except the cleared input
    if not message:
        return ("",) + nochange
    if MODEL_IMPORT_ERROR:
        raise gr.Error(_import_error_message())
    theme = theme if theme in THEMES else "fantasy"

    # Skip the (GPU) intent classifier when there's no log yet — the only sensible action is to
    # forge one. Saves a whole GPU reservation on the most common first message.
    if not quests:
        intent, kind = {}, "forge"
    else:
        intent = route_intent(message, _chat_context(quests, selected_id))
        kind = intent.get("intent", "unknown")

    if kind == "forge":
        q, a, fid, im, sc, de, st, b2 = _do_forge(message, theme, photo, browser)
        return ("", q, a, fid, im, sc, de, st, b2)

    if kind == "add_tasks":
        raw = generate_quests_raw(message, theme)
        quests2 = merge_quests(quests, raw, theme)
        gr.Info(f"Added {len(quests2) - len(quests)} quest(s) to your log.")
        return ("", quests2, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(),
                _stats_html(quests2), _merge(browser, quests=quests2))

    if kind in ("mark_done", "mark_couldnt"):
        qid = _resolve_target(quests, intent.get("target_task", ""), selected_id)
        if qid is None:
            gr.Info("Which quest? Select it on the right, or name it.")
            return ("",) + nochange
        result_kind = "success" if kind == "mark_done" else "failure"
        reason = intent.get("reason") or (message if kind == "mark_couldnt" else None)
        q2, im2, sc, de, st, b2 = _apply_result(qid, quests, adventure, photo, images, browser,
                                                 result_kind, reason)
        return ("", q2, gr.update(), qid, im2, sc, de, st, b2)

    gr.Info("Frog Master: tell me your plans to forge quests, what you finished, or what you couldn't do.")
    return ("",) + nochange


# ----------------------------- other handlers -----------------------------

def upload_photo(pil, browser):
    if pil is None:
        return gr.update(), gr.update()
    b64 = _resize_dataurl(pil)
    return b64, _merge(browser, photo=b64)


def change_theme(theme, adventure, images, browser):
    """The world picker drives image generation. Changing it with an adventure present re-themes
    art_style and clears the cached scenes so they regenerate in the new world on next select."""
    theme = theme if theme in THEMES else "fantasy"
    if adventure and adventure.get("theme") != theme:
        adventure = dict(adventure)
        adventure["theme"] = theme
        adventure["art_style"] = _art_style(theme)
        desc = (f'<div class="fq-desc"><p class="fq-empty">World changed to {theme.upper()}. '
                'Reselect a quest to redraw its scene in the new style.</p></div>')
        return adventure, {}, None, desc, _merge(browser, theme=theme, adventure=adventure, images={})
    return gr.update(), gr.update(), gr.update(), gr.update(), _merge(browser, theme=theme)


def reset_all():
    d = _default_state()
    return (d, None, None, [], None, None, {}, None,
            _desc_html(None, None), _stats_html([]), "fantasy")


def boot(browser):
    """Restore everything from BrowserState on page load. Shows cached scenes only (never
    generates) so reloads are instant; uncached scenes regenerate when the quest is clicked."""
    b = browser or _default_state()
    photo = b.get("photo")
    quests = b.get("quests") or []
    adventure = b.get("adventure")
    images = b.get("images") or {}
    theme = b.get("theme") or "fantasy"

    quest = _find(quests, b.get("selected_id")) or (quests[0] if quests else None)
    selected = quest["id"] if quest else None
    scene = None
    if quest:
        cache = images.get(quest["id"]) or {}
        data = cache.get(quest.get("image_state", "initial")) or cache.get("initial")
        if data:
            scene = _dataurl_to_pil(data)
    photo_pil = _dataurl_to_pil(photo) if photo else None
    return (photo_pil, photo, quests, adventure, selected, images,
            scene, _desc_html(quest, adventure), _stats_html(quests), theme)


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

_DIR = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(_DIR, "theme.css"), "r", encoding="utf-8") as _f:
    THEME_CSS = _f.read()

with gr.Blocks(title="FrogQuest") as demo:
    browser = gr.BrowserState(_default_state(), storage_key="frogquest")
    photo_state = gr.State(None)          # resized photo as base64 data URL
    quests_state = gr.State([])           # drives the quest-log render
    adventure_state = gr.State(None)      # title / art_style / seed / theme
    selected_id_state = gr.State(None)    # currently selected quest id
    images_state = gr.State({})           # {quest_id: {initial/success/failure: jpeg data-url}}

    gr.HTML('<div class="fq-topbar"><h1 class="fq-logo">FROG<b>QUEST</b></h1></div>')

    with gr.Row(elem_classes=["fq-app-grid"]):
        # ---------- LEFT: hero ----------
        with gr.Column(elem_classes=["fq-left"]):
            gr.HTML('<h2 class="panel-title">HERO</h2>')
            photo_image = gr.Image(
                type="pil", sources=["upload"], show_label=False, height=190,
                elem_classes=["fq-photo"],
            )
            gr.HTML('<p class="hint">Click to upload / change your photo. It stays in your '
                    'browser — sent to the GPU only to draw your scenes, never stored.</p>')
            stats_html = gr.HTML(_stats_html([]))
            gr.HTML('<h3 class="panel-subtitle">WORLD</h3>')
            theme_radio = gr.Radio(
                choices=list(THEMES), value="fantasy", show_label=False,
                elem_classes=["fq-theme"],
            )
            reset_btn = gr.Button("⟲ RESET", elem_classes=["pix-btn", "small"])

        # ---------- CENTER: selected scene + actions ----------
        with gr.Column(elem_classes=["fq-center"]):
            scene_image = gr.Image(
                show_label=False, interactive=False, container=False, height=380,
                elem_classes=["fq-scene"],
            )
            with gr.Row(elem_classes=["fq-detail"]):
                with gr.Column(scale=3):
                    desc_html = gr.HTML(_desc_html(None, None))
                with gr.Column(scale=1, min_width=136, elem_classes=["fq-actions"]):
                    done_btn = gr.Button("✓ DONE", elem_classes=["pix-btn", "btn-done"])
                    couldnt_btn = gr.Button("✗ COULDN'T", elem_classes=["pix-btn", "btn-couldnt"])

        # ---------- RIGHT: quest log ----------
        with gr.Column(elem_classes=["fq-right"]):
            gr.HTML('<h2 class="panel-title">QUEST LOG</h2>')

            @gr.render(inputs=[quests_state, selected_id_state])
            def render_tasklist(quests, selected):
                if not quests:
                    gr.HTML('<p class="hint">No quests yet.<br>Forge them with the Frog '
                            'Master chat below.</p>')
                    return
                with gr.Column(elem_classes=["fq-tasklist"]):
                    for q in quests:
                        classes = ["fq-task-item"]
                        if q.get("id") == selected:
                            classes.append("selected")
                        if q.get("is_frog"):
                            classes.append("is-frog")
                        if q.get("status") == "success":
                            classes.append("done")
                        elif q.get("status") == "failure":
                            classes.append("failed")
                        if q.get("status") == "success":
                            mark = "✓"
                        elif q.get("status") == "failure":
                            mark = "✗"
                        elif q.get("is_frog"):
                            mark = "🐸"
                        elif q.get("type") == "bonus":
                            mark = "✦"
                        else:
                            mark = "•"
                        finished = q.get("status") in ("success", "failure")
                        with gr.Row(elem_classes=["fq-task-row"]):
                            btn = gr.Button(f"{mark} {q.get('quest_title', '')}", elem_classes=classes)
                            btn.click(
                                (lambda qid: (lambda photo, adv, imgs, qs, br:
                                              select_quest(qid, photo, adv, imgs, qs, br)))(q["id"]),
                                inputs=[photo_state, adventure_state, images_state, quests_state, browser],
                                outputs=[selected_id_state, scene_image, desc_html, images_state, browser],
                            )
                            # Clear button appears only once a quest is done/failed, to remove it.
                            if finished:
                                clr = gr.Button("✕", elem_classes=["pix-btn", "small", "fq-task-clear"])
                                clr.click(
                                    (lambda qid: (lambda qs, adv, imgs, sel, br:
                                                  remove_quest(qid, qs, adv, imgs, sel, br)))(q["id"]),
                                    inputs=[quests_state, adventure_state, images_state, selected_id_state, browser],
                                    outputs=[quests_state, images_state, selected_id_state, scene_image,
                                             desc_html, stats_html, browser],
                                )

    # ---------- BOTTOM: Frog Master chat ----------
    with gr.Row(elem_classes=["fq-chatbar"]):
        chat_input = gr.Textbox(
            show_label=False, lines=2, elem_classes=["compose", "fq-chat-input"],
            placeholder="Tell the Frog Master your plans, what you finished, or what you couldn't do...",
        )
        send_btn = gr.Button("SEND ▶", elem_classes=["pix-btn", "primary", "fq-chat-send"])
    gr.HTML('<p class="fq-chat-hint">e.g. "Finish the report, reply to emails, book dentist" · '
            '"I finished the report" · "Couldn\'t do the gym — too tired"</p><p>Where to Start, Where to Break and let Go,<br>Where to change <br>Where to look and to grow </p>')

    # ---------- wiring ----------
    photo_image.upload(upload_photo, [photo_image, browser], [photo_state, browser])
    theme_radio.change(
        change_theme, [theme_radio, adventure_state, images_state, browser],
        [adventure_state, images_state, scene_image, desc_html, browser],
    )
    reset_btn.click(
        reset_all, None,
        [browser, photo_image, photo_state, quests_state, adventure_state, selected_id_state,
         images_state, scene_image, desc_html, stats_html, theme_radio],
    )
    done_btn.click(
        apply_done, [selected_id_state, quests_state, adventure_state, photo_state, images_state, browser],
        [quests_state, images_state, scene_image, desc_html, stats_html, browser],
    )
    couldnt_btn.click(
        apply_couldnt, [selected_id_state, quests_state, adventure_state, photo_state, images_state, browser],
        [quests_state, images_state, scene_image, desc_html, stats_html, browser],
    )

    _chat_inputs = [chat_input, quests_state, adventure_state, theme_radio, photo_state,
                    selected_id_state, images_state, browser]
    _chat_outputs = [chat_input, quests_state, adventure_state, selected_id_state, images_state,
                     scene_image, desc_html, stats_html, browser]
    send_btn.click(chat_send, _chat_inputs, _chat_outputs)
    chat_input.submit(chat_send, _chat_inputs, _chat_outputs)

    demo.load(
        boot, [browser],
        [photo_image, photo_state, quests_state, adventure_state, selected_id_state, images_state,
         scene_image, desc_html, stats_html, theme_radio],
    )


if __name__ == "__main__":
    demo.launch(
        css=THEME_CSS,
        theme=gr.themes.Base(font=[gr.themes.GoogleFont("Press Start 2P"), "monospace"]),
    )