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 SAME full state is also written through to a
per-user JSON file server-side (store.py), keyed by a random "Hero Code" uid that lives in
BrowserState — pasting the code on a new/wiped browser restores everything. No login.

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
import uuid

# 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 MAX_CAMPAIGNS, THEMES, merge_quests, validate_and_clamp, validate_campaign
import store


# 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_campaign_raw, generate_quests_raw, route_intent
    from images import edit_image, initial_image, initial_images, 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 {"uid": None, "photo": None, "theme": "fantasy", "adventure": None,
            "quests": [], "campaigns": [], "selected_id": None, "images": {},
            "stats": {"xp": 0, "done": 0, "retreats": 0}}


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"


def _world_for(quest: dict | None, adventure: dict | None, campaigns: list | None) -> dict | None:
    """The art world (title/art_style/seed) a quest renders in: its campaign's own cohesive world
    if it belongs to one, else the day adventure's. None when neither exists yet."""
    cid = (quest or {}).get("campaign_id")
    if cid:
        c = next((c for c in (campaigns or []) if c.get("id") == cid), None)
        if c:
            return {"title": c.get("title", ""), "art_style": c.get("art_style", _art_style("fantasy")),
                    "seed": c.get("seed", 0)}
    return adventure


# ----------------------------- 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("campaign_id"):
        badges.append('<span class="badge group-badge">⚑ CAMPAIGN</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 _seed_stats(quests: list | None) -> dict:
    """One-time migration: derive lifetime counters from an old record's quest list."""
    quests = quests or []
    return {
        "done": sum(1 for q in quests if q.get("status") == "success"),
        "retreats": 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"),
    }


def _stats_html(quests: list, stats: dict | None = None) -> str:
    """Hero stats = LIFETIME counters (stored in state; survive quest removal and day re-forges)
    plus a live count of currently active quests."""
    s = stats or {"xp": 0, "done": 0, "retreats": 0}
    active = sum(1 for q in (quests or []) if q.get("status") == "active")
    return (
        '<div class="fq-stats"><h4>HERO STATS</h4>'
        f'<div class="stat-row"><span>QUESTS DONE</span><b>{int(s.get("done", 0))}</b></div>'
        f'<div class="stat-row"><span>XP EARNED</span><b>{int(s.get("xp", 0))}</b></div>'
        f'<div class="stat-row"><span>RETREATS</span><b>{int(s.get("retreats", 0))}</b></div>'
        f'<div class="stat-row"><span>ACTIVE</span><b>{active}</b></div>'
        '</div>'
    )


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

def select_quest(qid, photo, adventure, images, quests, campaigns, browser):
    """Select a quest: show its cached scene, or lazily generate the initial scene on first view.
    Campaign quests render in their campaign's world. 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,
                _persist_browser(_merge(browser, selected_id=qid)))

    world = _world_for(quest, adventure, campaigns)
    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 world:
            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), world["art_style"],
                                quest.get("initial_image_prompt", ""), int(world["seed"]))
            cache["initial"] = pil_to_data_url(pil)
            images[qid] = cache
            images = _trim_images(images, qid)
            scene = pil
    return (qid, scene, _desc_html(quest, world or adventure, hint), images,
            _persist_browser(_merge(browser, selected_id=qid, images=images)))


def _apply_result(qid, quests, adventure, photo, images, campaigns, 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.")
    world = _world_for(quest, adventure, campaigns)
    if not world:
        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), world["art_style"],
                                 quest.get("initial_image_prompt", ""), int(world["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, world["art_style"], int(world["seed"]))
    cache[kind] = pil_to_data_url(edited)
    images[qid] = cache
    images = _trim_images(images, qid)

    # Bank LIFETIME stats — but only on the FIRST resolution of this quest (a re-mark of an
    # already-finished quest must not double-count XP or retreats).
    stats = dict((browser or {}).get("stats") or {"xp": 0, "done": 0, "retreats": 0})
    if quest.get("status", "active") == "active":
        if kind == "success":
            stats["done"] = int(stats.get("done", 0)) + 1
            stats["xp"] = int(stats.get("xp", 0)) + int(quest.get("xp", 0))
        else:
            stats["retreats"] = int(stats.get("retreats", 0)) + 1

    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, world)
    return (quests, images, edited, desc, _stats_html(quests, stats),
            _merge(browser, quests=quests, images=images, stats=stats))


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


def apply_couldnt(qid, quests, adventure, photo, images, campaigns, browser):
    if qid is None:
        raise gr.Error("Select a quest first.")
    return _apply_result(qid, quests, adventure, photo, images, campaigns, 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, (browser or {}).get("stats")),
            _persist_browser(_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 _batch_initials(photo, art_style, seed, quests, images, keep_id):
    """ONE batched GPU call generating every quest's initial scene into the images cache.
    Mutates and returns `images`. No-op without a photo (lazy/hint path takes over)."""
    if not photo or not quests:
        return images
    pils = initial_images(_dataurl_to_pil(photo), art_style,
                          [q.get("initial_image_prompt", "") for q in quests], int(seed))
    for q, pil in zip(quests, pils):
        images.setdefault(q["id"], {})["initial"] = pil_to_data_url(pil)
    return _trim_images(images, keep_id)


def _do_forge(message, theme, photo, quests_old, images_old, campaigns, browser):
    """Forge a fresh DAY log. Campaign quests (and their cached scenes) survive the re-forge —
    only standalone day quests are replaced. With a photo present, ALL day scenes are generated
    eagerly in one batched GPU call (clicking quests afterwards is instant)."""
    raw = generate_quests_raw(message, theme)
    adv = validate_and_clamp(raw, theme)
    day_quests, adventure = adv["quests"], adv["adventure"]
    keep = [q for q in (quests_old or []) if q.get("campaign_id")]
    keep_ids = {q["id"] for q in keep}
    for q in day_quests:  # avoid id collisions with the kept campaign quests
        while q["id"] in keep_ids:
            q["id"] += "-d"
    quests = day_quests + keep
    images = {k: v for k, v in (images_old or {}).items() if k in keep_ids}
    frog_id = day_quests[0]["id"]
    images = _batch_initials(photo, adventure["art_style"], adventure["seed"],
                             day_quests, images, frog_id)
    # Cache is warm now (or photo absent) -> this just selects the frog; no extra GPU call.
    _sel, scene, desc, images2, _b = select_quest(frog_id, photo, adventure, images, quests,
                                                  campaigns, browser)
    browser2 = _merge(browser, theme=theme, adventure=adventure, quests=quests,
                      selected_id=frog_id, images=images2)
    return (quests, adventure, frog_id, images2, scene, desc,
            _stats_html(quests, (browser or {}).get("stats")), 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, campaigns, 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, quests, images,
                                                  campaigns, 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, (browser or {}).get("stats")), _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, campaigns,
                                                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)


# forge_campaign outputs (order): campaign_goal_box, quests_state, campaigns_state, stats_html,
#                                 images_state, selected_id_state, scene_image, desc_html, browser
def forge_campaign(goal, do_research, theme, photo, quests, campaigns, images, browser):
    """One long-term goal -> a campaign + its ordered quest chain, appended to the log. With
    do_research, ddgs+BS4 snippets ground the plan (research.py; degrades gracefully without it).
    With a photo present, ALL the campaign's scenes are generated in one batched GPU call and the
    first quest is shown."""
    goal = (goal or "").strip()
    nochange = (gr.update(),) * 8
    if not goal:
        gr.Info("Name your long-term goal first.")
        return (gr.update(),) + nochange
    if MODEL_IMPORT_ERROR:
        raise gr.Error(_import_error_message())
    campaigns = list(campaigns or [])
    if len(campaigns) >= MAX_CAMPAIGNS:
        raise gr.Error(f"You already have {MAX_CAMPAIGNS} campaigns — conquer or clear one first.")
    theme = theme if theme in THEMES else "fantasy"

    snippets, sources = "", []
    if do_research:
        try:
            from research import research_goal
            gr.Info("Scouting the web for your goal…")
            r = research_goal(goal)
            snippets = "\n".join(f"- {s}" for s in r.get("snippets") or [])
            sources = list(r.get("sources") or [])
        except Exception:
            gr.Info("Web research unavailable — forging from the model's own knowledge.")

    raw = generate_campaign_raw(goal, theme, snippets)
    result = validate_campaign(
        raw, theme, goal,
        existing_quest_ids={q.get("id") for q in (quests or [])},
        existing_campaign_ids={c.get("id") for c in campaigns},
    )
    camp = result["campaign"]
    camp["sources"] = sources
    campaigns2 = campaigns + [camp]
    new_quests = result["quests"]
    quests2 = list(quests or []) + new_quests
    first = new_quests[0]

    images2 = _batch_initials(photo, camp["art_style"], camp["seed"],
                              new_quests, dict(images or {}), first["id"])
    cache = images2.get(first["id"]) or {}
    scene = _dataurl_to_pil(cache["initial"]) if cache.get("initial") else None
    desc = _desc_html(first, _world_for(first, None, campaigns2),
                      None if scene else "Upload your photo (left) to draw this scene.")
    gr.Info(f"Campaign forged: {camp['title']} — {len(new_quests)} quests.")
    return ("", quests2, campaigns2, _stats_html(quests2, (browser or {}).get("stats")), images2,
            first["id"], scene, desc,
            _merge(browser, quests=quests2, campaigns=campaigns2, images=images2,
                   selected_id=first["id"]))


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(uid):
    d = _default_state()
    d["uid"] = uid  # reset clears the data, not the identity — the Hero Code stays
    return (d, None, None, [], [], None, None, {}, None,
            _desc_html(None, None), _stats_html([]), "fantasy")


def persist(uid, browser):
    """Write-through the full state to the user's durable server record. Bound via .then(...) on
    TOP-LEVEL events only — never on events created inside @gr.render: a render re-run re-registers
    those functions under new indices, so a queued .then step 500s with KeyError."""
    store.save(uid, browser)


def _persist_browser(b):
    """In-handler persistence for @gr.render-created events (quest select/remove), where a chained
    .then(persist) would break (see persist's docstring). The uid travels inside the browser dict."""
    store.save((b or {}).get("uid"), b)
    return b


def _hero_code_html(uid, hf_user=None) -> str:
    if hf_user:
        return (f'<p class="hint">Synced to your Hugging Face account (<b>{_esc(hf_user)}</b>) — '
                'no code needed. Your quests follow your login on any device. Paste an old Hero '
                'Code below to import that hero into your account.</p>')
    return (f'<p class="hint">Your Hero Code: <b>{_esc(uid)}</b><br>Write it down — paste it on '
            'any device (or after a cleared browser) to summon this hero back. Or just sign in '
            'with Hugging Face (top right) and skip the code.</p>')


def _hf_username(profile, request) -> str | None:
    """The logged-in HF username, or None. On Spaces with hf_oauth the identity arrives via the
    auto-injected gr.OAuthProfile — NOT request.username (that's only set by launch(auth=...));
    reading the wrong one is exactly how cross-device sync silently broke."""
    name = getattr(profile, "username", None)
    if not name and request is not None:
        name = getattr(request, "username", None)  # local-dev mock / password-auth fallback
    return name or None


def boot(browser, profile: gr.OAuthProfile | None = None, request: gr.Request = None):
    """Restore on page load. Identity precedence: a logged-in HF account ("hf-<username>", durable
    and cross-device) > the browser dict's random uid ("Hero Code") > a freshly generated one.
    The server-side JSON record for that identity is the source of truth and is mirrored back into
    BrowserState; on first sight of an identity the browser's current state migrates into it.
    Shows cached scenes only (never generates) so reloads are instant."""
    b = dict(browser or _default_state())
    hf_user = _hf_username(profile, request)
    if hf_user:
        uid = f"hf-{hf_user}"
    else:
        uid = b.get("uid") or uuid.uuid4().hex[:12]
        if str(uid).startswith("hf-"):
            # Leftover account uid after logout: never write anonymous data into the account
            # record — mint a fresh anonymous identity carrying the local state.
            uid = uuid.uuid4().hex[:12]
    saved = store.load(uid)
    if saved:
        b = saved
    b["uid"] = uid
    if not saved:
        store.save(uid, b)  # first visit under this identity: adopt the browser's current state
    return _rehydrate(uid, b, hf_user)


def restore_hero(code, browser, profile: gr.OAuthProfile | None = None, request: gr.Request = None):
    """Adopt a pasted Hero Code: load that uid's server record and take over this browser's
    identity/state with it. When LOGGED IN, the code's data is imported INTO the HF account record
    instead — tying it to the login so it follows the user to every device."""
    code = (code or "").strip()
    saved = store.load(code)
    if not saved:
        raise gr.Error("No hero found for that code. Check it and try again.")
    hf_user = _hf_username(profile, request)
    uid = f"hf-{hf_user}" if hf_user else code
    saved["uid"] = uid
    if uid != code:
        store.save(uid, saved)
        gr.Info(f"Hero imported into your HF account ({hf_user}) — it now follows your login.")
    return _rehydrate(uid, saved, hf_user)


def _rehydrate(uid, b, hf_user=None):
    """Build the full page state tuple (shared by boot and restore_hero)."""
    photo = b.get("photo")
    quests = b.get("quests") or []
    campaigns = b.get("campaigns") or []
    if "stats" not in b:  # migrate pre-stats records once: bank what's currently visible
        b["stats"] = _seed_stats(quests)
    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
    world = _world_for(quest, adventure, campaigns) or adventure
    return (uid, _hero_code_html(uid, hf_user), b, photo_pil, photo, quests, campaigns, adventure,
            selected, images, scene, _desc_html(quest, world), _stats_html(quests, b.get("stats")),
            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}}
    campaigns_state = gr.State([])        # long-term goals; each owns a chain of campaign quests
    campaign_filter_state = gr.State(None)  # campaign id to filter the quest log by (None = all)
    uid_state = gr.State(None)            # the "Hero Code" uid; generated by boot, lives in BrowserState

    with gr.Row(elem_classes=["fq-topbar"]):
        gr.HTML('<h1 class="fq-logo">FROG<b>QUEST</b></h1>')
        gr.LoginButton(elem_classes=["fq-login"])

    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">Upload your photo — used only to draw you into your scenes.</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"])
            with gr.Accordion("⚿ HERO CODE", open=False, elem_classes=["fq-herocode"]):
                herocode_html = gr.HTML("")  # filled by boot
                restore_box = gr.Textbox(show_label=False, lines=1, elem_classes=["compose"],
                                         placeholder="paste a Hero Code...")
                restore_btn = gr.Button("SUMMON HERO", elem_classes=["pix-btn", "small"])

        # ---------- CENTER: selected scene + actions + Frog Master chat ----------
        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"])
            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 class="fq-chat-hint">Where to Start, Where to Break and let Go,<br>'
                    'Where to change <br>Where to look and to grow </p>')

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

            @gr.render(inputs=[campaigns_state, quests_state, campaign_filter_state])
            def render_campaigns(campaigns, quests, cfilter):
                if not campaigns:
                    gr.HTML('<p class="hint">No campaigns yet — forge a long-term goal below.</p>')
                    return
                with gr.Column(elem_classes=["fq-tasklist"]):
                    for c in campaigns:
                        cid = c.get("id")
                        cq = [q for q in (quests or []) if q.get("campaign_id") == cid]
                        done = sum(1 for q in cq if q.get("status") == "success")
                        classes = ["fq-task-item"] + (["selected"] if cfilter == cid else [])
                        cb = gr.Button(f"⚑ {c.get('title', '')} — {done}/{len(cq)}",
                                       elem_classes=classes)
                        # Click filters the quest log to this campaign; click again to clear.
                        cb.click((lambda _cid: (lambda cur: None if cur == _cid else _cid))(cid),
                                 inputs=[campaign_filter_state], outputs=[campaign_filter_state])
                        if cfilter == cid and c.get("sources"):
                            links = " · ".join(
                                f'<a href="{_esc(u)}" target="_blank" rel="noopener">[{i + 1}]</a>'
                                for i, u in enumerate(c["sources"][:6]))
                            gr.HTML(f'<p class="hint">researched: {links}</p>')

            with gr.Accordion("⚒ NEW CAMPAIGN", open=False):
                campaign_goal_box = gr.Textbox(
                    show_label=False, lines=2, elem_classes=["compose"],
                    placeholder='A long-term goal, e.g. "score 95th percentile on the GRE"...',
                )
                research_check = gr.Checkbox(label="research the web first", value=False)
                campaign_btn = gr.Button("FORGE CAMPAIGN", elem_classes=["pix-btn", "primary"])

            gr.HTML('<h2 class="panel-title">QUEST LOG</h2>')

            @gr.render(inputs=[quests_state, selected_id_state, campaigns_state, campaign_filter_state])
            def render_tasklist(quests, selected, campaigns, cfilter):
                quests = quests or []
                if cfilter:
                    quests = [q for q in quests if q.get("campaign_id") == cfilter]
                if not quests:
                    gr.HTML('<p class="hint">No quests yet.<br>Forge them with the Frog '
                            'Master chat below.</p>')
                    return
                # Standalone day quests first, then each campaign's chain under its banner.
                titles = {c.get("id"): c.get("title", "") for c in (campaigns or [])}
                groups = [(None, [q for q in quests if not q.get("campaign_id")])]
                groups += [(c.get("id"), [q for q in quests if q.get("campaign_id") == c.get("id")])
                           for c in (campaigns or [])]
                with gr.Column(elem_classes=["fq-tasklist"]):
                    for gid, gquests in groups:
                        if not gquests:
                            continue
                        if gid:
                            gr.HTML(f'<p class="hint">⚑ {_esc(titles.get(gid, ""))}</p>')
                        for q in gquests:
                            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, cs, br:
                                                  select_quest(qid, photo, adv, imgs, qs, cs, br)))(q["id"]),
                                    inputs=[photo_state, adventure_state, images_state, quests_state,
                                            campaigns_state, browser],
                                    outputs=[selected_id_state, scene_image, desc_html, images_state, browser],
                                )  # NO .then here: render-internal chains 500 (see persist docstring)
                                # 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],
                                    )  # NO .then here: render-internal chains 500 (see persist docstring)

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

    campaign_btn.click(
        forge_campaign, [campaign_goal_box, research_check, theme_radio, photo_state, quests_state,
                         campaigns_state, images_state, browser],
        [campaign_goal_box, quests_state, campaigns_state, stats_html, images_state,
         selected_id_state, scene_image, desc_html, browser],
    ).then(persist, [uid_state, browser])

    _chat_inputs = [chat_input, quests_state, adventure_state, theme_radio, photo_state,
                    selected_id_state, images_state, campaigns_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).then(persist, [uid_state, browser])
    chat_input.submit(chat_send, _chat_inputs, _chat_outputs).then(persist, [uid_state, browser])

    _hydrate_outputs = [uid_state, herocode_html, browser, photo_image, photo_state, quests_state,
                        campaigns_state, adventure_state, selected_id_state, images_state,
                        scene_image, desc_html, stats_html, theme_radio]
    demo.load(boot, [browser], _hydrate_outputs)
    restore_btn.click(restore_hero, [restore_box, browser], _hydrate_outputs)


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