schema.py

"""FrogQuest data contract: the JSON schema the LLM emits + server-side validate/clamp.

The LLM only ever writes DATA to this contract. Nothing it returns is trusted: every
field is validated and clamped here before the frontend renders it, so a malformed or
missing field can never break a quest card. See CLAUDE.md "Data contract".
"""
from __future__ import annotations

import random
import re
from typing import Any

THEMES = ("cyberpunk", "fantasy", "space")
QUEST_TYPES = ("main", "bonus")
STATUSES = ("active", "success", "failure")
IMAGE_STATES = ("initial", "success", "failure")

MAX_QUESTS = 12
MAX_BONUS = 3

# JSON Schema passed to llama-cpp-python via response_format to constrain generation.
# Kept deliberately permissive (the strict guarantees are enforced by validate_and_clamp,
# not the grammar) because GBNF-from-schema chokes on overly nested constraints.
RESPONSE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "adventure": {
            "type": "object",
            "properties": {
                "title": {"type": "string"},
                "theme": {"type": "string", "enum": list(THEMES)},
                "art_style": {"type": "string"},
                "seed": {"type": "integer"},
            },
            "required": ["title", "theme", "art_style", "seed"],
        },
        "quests": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "id": {"type": "string"},
                    "task": {"type": "string"},
                    "quest_title": {"type": "string"},
                    "narrative": {"type": "string"},
                    "type": {"type": "string", "enum": list(QUEST_TYPES)},
                    "goal_group": {"type": "string"},
                    "is_frog": {"type": "boolean"},
                    "initial_image_prompt": {"type": "string"},
                    "success_edit": {"type": "string"},
                    "failure_edit": {"type": "string"},
                    "xp": {"type": "integer"},
                },
                "required": [
                    "task", "quest_title", "narrative", "type",
                    "is_frog", "initial_image_prompt", "success_edit",
                    "failure_edit", "xp",
                ],
            },
        },
    },
    "required": ["adventure", "quests"],
}


# JSON Schema for the Frog Master chat router. The LLM classifies each user message into one
# intent and (optionally) names a target task and/or a free-text reason. Kept tiny: the chat
# never sees images — only a short text context (see llm.route_intent).
INTENT_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "intent": {
            "type": "string",
            "enum": ["forge", "add_tasks", "mark_done", "mark_couldnt", "unknown"],
        },
        "target_task": {"type": "string"},
        "reason": {"type": "string"},
    },
    "required": ["intent"],
}


def _s(value: Any, default: str = "") -> str:
    """Coerce to a stripped string."""
    if value is None:
        return default
    if isinstance(value, str):
        return value.strip()
    return str(value).strip()


def _clamp_int(value: Any, lo: int, hi: int, default: int) -> int:
    try:
        n = int(value)
    except (TypeError, ValueError):
        return default
    return max(lo, min(hi, n))


def _slugify(text: str, fallback: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
    return slug or fallback


def _clean_quest(q: Any, idx: int, seen_ids: set[str], qtype: str) -> dict[str, Any]:
    """Coerce one raw quest dict into a render-ready quest with a unique non-empty id.

    `seen_ids` is mutated in place; `qtype` is the already-resolved type (so the caller owns
    the bonus-cap accounting). Runtime state (status/image_state) is always set fresh here and
    never trusted from the model.
    """
    title = _s(q.get("quest_title")) or _s(q.get("task")) or f"Quest {idx + 1}"
    qid = _s(q.get("id"))
    if not qid or qid in seen_ids:
        qid = _slugify(title, f"quest-{idx + 1}")
        base, n = qid, 2
        while qid in seen_ids:
            qid, n = f"{base}-{n}", n + 1
    seen_ids.add(qid)

    return {
        "id": qid,
        "task": _s(q.get("task")) or title,
        "quest_title": title,
        "narrative": _s(q.get("narrative")),
        "type": qtype,
        "goal_group": _s(q.get("goal_group")) or None,
        "is_frog": bool(q.get("is_frog")),
        "initial_image_prompt": _s(q.get("initial_image_prompt")),
        "success_edit": _s(q.get("success_edit")) or "Show the hero victorious.",
        "failure_edit": _s(q.get("failure_edit"))
        or "The hero retreats safely to rest and try again another day.",
        "xp": _clamp_int(q.get("xp"), 0, 100, 25),
        "status": "active",
        "image_state": "initial",
    }


def _resolve_type(q: Any, bonus_count: int) -> tuple[str, int]:
    """Resolve a quest's type, demoting bonus quests to main once MAX_BONUS is reached.
    Returns (type, new_bonus_count)."""
    qtype = q.get("type") if isinstance(q, dict) else None
    qtype = qtype if qtype in QUEST_TYPES else "main"
    if qtype == "bonus":
        if bonus_count >= MAX_BONUS:
            qtype = "main"  # demote excess bonus quests to main
        else:
            bonus_count += 1
    return qtype, bonus_count


def validate_and_clamp(raw: Any, theme: str) -> dict[str, Any]:
    """Turn whatever the LLM returned into a safe, render-ready adventure object.

    Guarantees on return:
      - adventure has title/theme/art_style/seed (int)
      - 1..MAX_QUESTS quests, each with all schema fields and unique non-empty id
      - exactly one is_frog quest, and it is ordered FIRST
      - <= MAX_BONUS bonus quests
      - type/status/image_state are valid enum values; xp in 0..100
    """
    theme = theme if theme in THEMES else THEMES[0]
    raw = raw if isinstance(raw, dict) else {}

    adv_in = raw.get("adventure") if isinstance(raw.get("adventure"), dict) else {}
    seed = adv_in.get("seed")
    try:
        seed = int(seed)
    except (TypeError, ValueError):
        seed = random.randint(0, 2**31 - 1)

    adv_theme = adv_in.get("theme")
    adv_theme = adv_theme if adv_theme in THEMES else theme

    adventure = {
        "title": _s(adv_in.get("title")) or "Your Quest Log",
        "theme": adv_theme,
        "art_style": _s(adv_in.get("art_style"))
        or f"8-bit / 16-bit retro pixel art, {adv_theme} palette, NES RPG style",
        "seed": seed,
    }

    quests_in = raw.get("quests")
    quests_in = quests_in if isinstance(quests_in, list) else []

    cleaned: list[dict[str, Any]] = []
    seen_ids: set[str] = set()
    bonus_count = 0

    for idx, q in enumerate(quests_in):
        if len(cleaned) >= MAX_QUESTS:
            break
        if not isinstance(q, dict):
            continue
        qtype, bonus_count = _resolve_type(q, bonus_count)
        cleaned.append(_clean_quest(q, idx, seen_ids, qtype))

    if not cleaned:
        # Degenerate output: synthesise a single placeholder so the UI still renders.
        cleaned.append({
            "id": "quest-1", "task": "Add your first task", "quest_title": "The Journey Begins",
            "narrative": "Tell the oracle your goals to fill this log.", "type": "main",
            "goal_group": None, "is_frog": True, "initial_image_prompt": "",
            "success_edit": "Show the hero victorious.",
            "failure_edit": "The hero retreats to try again.", "xp": 10,
            "status": "active", "image_state": "initial",
        })

    # Enforce exactly one frog, ordered first. Prefer a main quest flagged by the model.
    frog_idx = next((i for i, q in enumerate(cleaned) if q["is_frog"] and q["type"] == "main"), None)
    if frog_idx is None:
        frog_idx = next((i for i, q in enumerate(cleaned) if q["is_frog"]), None)
    if frog_idx is None:
        frog_idx = next((i for i, q in enumerate(cleaned) if q["type"] == "main"), 0)
    for i, q in enumerate(cleaned):
        q["is_frog"] = (i == frog_idx)
    frog = cleaned.pop(frog_idx)
    cleaned.insert(0, frog)

    return {"adventure": adventure, "quests": cleaned}


def merge_quests(existing: list[dict[str, Any]], raw_new: Any, theme: str) -> list[dict[str, Any]]:
    """Append the LLM's freshly generated quests to an existing log (chat "add_tasks").

    New quests are cleaned through the same id-uniqueness / bonus-cap / MAX_QUESTS rules as
    validate_and_clamp, but the existing log's structure is preserved: the original frog stays
    first (new quests are forced is_frog=False) and existing quests are untouched. Returns the
    combined list (capped at MAX_QUESTS).
    """
    existing = list(existing or [])
    if len(existing) >= MAX_QUESTS:
        return existing

    raw = raw_new if isinstance(raw_new, dict) else {}
    quests_in = raw.get("quests")
    quests_in = quests_in if isinstance(quests_in, list) else []

    seen_ids = {q.get("id") for q in existing if q.get("id")}
    bonus_count = sum(1 for q in existing if q.get("type") == "bonus")
    base_idx = len(existing)

    added: list[dict[str, Any]] = []
    for i, q in enumerate(quests_in):
        if len(existing) + len(added) >= MAX_QUESTS:
            break
        if not isinstance(q, dict):
            continue
        qtype, bonus_count = _resolve_type(q, bonus_count)
        cleaned = _clean_quest(q, base_idx + i, seen_ids, qtype)
        cleaned["is_frog"] = False  # the original frog keeps its place at the front
        added.append(cleaned)

    return existing + added