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 datetime import datetime, timezone
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
MAX_CAMPAIGNS = 5         # concurrent long-term goals
MAX_CAMPAIGN_QUESTS = 8   # steps per campaign chain

# JSON Schemas passed to llama-cpp-python via response_format to constrain generation.
# Kept deliberately permissive (the strict guarantees are enforced by validate_and_clamp /
# validate_campaign, not the grammar) because GBNF-from-schema chokes on overly nested constraints.
_QUEST_ITEM_SCHEMA: dict[str, Any] = {
    "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",
    ],
}

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": _QUEST_ITEM_SCHEMA},
    },
    "required": ["adventure", "quests"],
}

# Campaign forging: one long-term goal -> a themed campaign + an ORDERED chain of concrete steps.
CAMPAIGN_RESPONSE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "campaign": {
            "type": "object",
            "properties": {
                "title": {"type": "string"},
                "art_style": {"type": "string"},
                "seed": {"type": "integer"},
            },
            "required": ["title", "art_style", "seed"],
        },
        "quests": {"type": "array", "items": _QUEST_ITEM_SCHEMA},
    },
    "required": ["campaign", "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",
        "campaign_id": None,  # set by validate_campaign for campaign chains; None = standalone
    }


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 validate_campaign(
    raw: Any,
    theme: str,
    goal: str,
    existing_quest_ids: set[str] | None = None,
    existing_campaign_ids: set[str] | None = None,
) -> dict[str, Any]:
    """Turn the LLM's campaign JSON into a safe campaign entity + its ordered quest chain.

    Guarantees on return:
      - campaign has a unique id ("camp-..."), title, verbatim goal, status "active",
        theme/art_style/seed (its own cohesive world), empty sources, created_at
      - 1..MAX_CAMPAIGN_QUESTS quests, each fully cleaned (`_clean_quest` rules), tagged with the
        campaign id, forced type "main" and is_frog False (the day log owns the frog), and with
        ids unique against `existing_quest_ids`
    """
    theme = theme if theme in THEMES else THEMES[0]
    raw = raw if isinstance(raw, dict) else {}
    camp_in = raw.get("campaign") if isinstance(raw.get("campaign"), dict) else {}

    goal_s = _s(goal)
    title = _s(camp_in.get("title")) or (goal_s[:48] if goal_s else "New Campaign")
    existing_c = set(existing_campaign_ids or ())
    cid = "camp-" + _slugify(title, "campaign")
    base, n = cid, 2
    while cid in existing_c:
        cid, n = f"{base}-{n}", n + 1

    try:
        seed = int(camp_in.get("seed"))
    except (TypeError, ValueError):
        seed = random.randint(0, 2**31 - 1)
    art_style = (_s(camp_in.get("art_style"))
                 or f"8-bit / 16-bit retro pixel art, {theme} palette, NES RPG style")

    quests_in = raw.get("quests")
    quests_in = quests_in if isinstance(quests_in, list) else []
    seen_ids = set(existing_quest_ids or ())
    cleaned: list[dict[str, Any]] = []
    for idx, q in enumerate(quests_in):
        if len(cleaned) >= MAX_CAMPAIGN_QUESTS:
            break
        if not isinstance(q, dict):
            continue
        cq = _clean_quest(q, idx, seen_ids, "main")
        cq["is_frog"] = False
        cq["campaign_id"] = cid
        cleaned.append(cq)

    if not cleaned:
        # Degenerate output: one placeholder step so the campaign still renders.
        qid = f"{cid}-step-1"
        while qid in seen_ids:
            qid += "x"
        cleaned.append({
            "id": qid, "task": f"Plan the first step toward: {goal_s or title}",
            "quest_title": "Chart the First Step", "narrative":
            "Every campaign begins with a single decision: what to do first.", "type": "main",
            "goal_group": None, "is_frog": False, "initial_image_prompt": "",
            "success_edit": "Show the hero victorious.",
            "failure_edit": "The hero retreats to try again another day.", "xp": 15,
            "status": "active", "image_state": "initial", "campaign_id": cid,
        })

    campaign = {
        "id": cid,
        "title": title,
        "goal": goal_s,
        "status": "active",
        "theme": theme,
        "art_style": art_style,
        "seed": seed,
        "sources": [],
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    return {"campaign": campaign, "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