ai-rpg/backend/app/engine/tools/tools.py

"""Tool definitions and handlers for the orchestrator's tool-calling loop."""
from __future__ import annotations

import json
import random
import uuid
from typing import Any, Awaitable, Callable, Dict, List, Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.llm import build_tool_schema
from app.core.rag import get_rag
from app.core.state_validator import apply_patch, validate_state
from app.logging_setup import get_logger
from app.models import DeferredTrigger, GlossaryEntry, World

log = get_logger("tools")


# === Tool schemas (OpenAI function-calling format) ===

DICE_ROLL_SCHEMA = build_tool_schema(
    name="dice_roll",
    description="Roll dice. Use 'sides' (e.g. 20 for d20) and optional 'count' (default 1) and 'modifier'. Returns the rolls and total.",
    params={
        "type": "object",
        "properties": {
            "sides": {"type": "integer", "description": "Number of sides on the die, e.g. 20 for d20"},
            "count": {"type": "integer", "description": "Number of dice to roll", "default": 1},
            "modifier": {"type": "integer", "description": "Modifier to add to total", "default": 0},
            "label": {"type": "string", "description": "What this roll represents, e.g. 'attack' or 'perception'"},
        },
        "required": ["sides"],
    },
)


UPDATE_STATE_SCHEMA = build_tool_schema(
    name="update_state",
    description="Apply a patch to world state. Paths use dot notation. ops: set, unset, append, increment, remove.",
    params={
        "type": "object",
        "properties": {
            "patch": {
                "type": "object",
                "description": "JSON-patch object with optional keys: set, unset, append, increment, remove. Each is a dict of path->value (or list of paths for unset).",
                "properties": {
                    "set": {"type": "object"},
                    "unset": {"type": "array", "items": {"type": "string"}},
                    "append": {"type": "object"},
                    "increment": {"type": "object"},
                    "remove": {"type": "object"},
                },
            }
        },
        "required": ["patch"],
    },
)


RAG_QUERY_SCHEMA = build_tool_schema(
    name="rag_query",
    description="Search the glossary (NPCs, locations, items, lore) for relevant facts.",
    params={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "Free-text search query"},
            "limit": {"type": "integer", "description": "Max results", "default": 5},
        },
        "required": ["query"],
    },
)


RAG_ADD_SCHEMA = build_tool_schema(
    name="rag_add",
    description="Add a new entry to the glossary (NPC, location, item, lore, event).",
    params={
        "type": "object",
        "properties": {
            "kind": {"type": "string", "enum": ["npc", "location", "item", "lore", "event", "rule"]},
            "name": {"type": "string"},
            "description": {"type": "string"},
            "payload": {"type": "object", "description": "Optional extra fields"},
        },
        "required": ["kind", "name", "description"],
    },
)


SCHEDULE_TRIGGER_SCHEMA = build_tool_schema(
    name="schedule_trigger",
    description=(
        "Schedule a deferred event tied to in-world time. When the world's "
        "internal clock reaches fire_at, the engine fires the event (calls the "
        "LLM with the description to produce a narrative beat and optional "
        "state patch). fire_at must use the same format as world.current_time "
        "('day_N_hour_H' or 'day_N_hour_H_min_M'). The world's calendar "
        "(hours_per_day, minutes_per_hour) is honored when comparing times."
    ),
    params={
        "type": "object",
        "properties": {
            "fire_at": {
                "type": "string",
                "description": "In-world time when the trigger fires, e.g. 'day_3_hour_14' or 'day_3_hour_14_min_30'.",
            },
            "description": {
                "type": "string",
                "description": "What should happen when the trigger fires. Be specific — this is fed to the LLM at fire time.",
            },
            "payload": {
                "type": "object",
                "description": "Optional structured payload (e.g. who, conditions, parameters).",
            },
        },
        "required": ["fire_at", "description"],
    },
)


ADVANCE_TIME_SCHEMA = build_tool_schema(
    name="advance_time",
    description=(
        "Advance the world's internal clock by days / hours / minutes. Use "
        "this when the player's action takes measurable in-world time (travel, "
        "sleep, crafting, long rest). The world's calendar (hours_per_day, "
        "minutes_per_hour) is honored. After time advances, any scheduled "
        "triggers whose fire_at is now <= the new time will fire "
        "automatically — so this is also how you 'run out the clock' on a "
        "scheduled event."
    ),
    params={
        "type": "object",
        "properties": {
            "days": {"type": "integer", "default": 0},
            "hours": {"type": "integer", "default": 0},
            "minutes": {"type": "integer", "default": 0},
            "reason": {"type": "string", "description": "Why time advances (logged for debugging)."},
        },
    },
)


RUN_SUBAGENT_SCHEMA = build_tool_schema(
    name="run_subagent",
    description="Spawn a sub-agent with clean context for a focused sub-task (e.g. generate NPC backstory, room description).",
    params={
        "type": "object",
        "properties": {
            "task": {"type": "string", "description": "The specific task for the sub-agent"},
            "context": {"type": "string", "description": "Minimal context needed (max 200 words)"},
        },
        "required": ["task"],
    },
)


# === Submission tools (how the LLM returns structured results) ===
# These replace the old "return JSON in your text response" pattern, which
# conflicted with tool use and caused the model to dump raw JSON into chat.

SUBMIT_PLAN_SCHEMA = build_tool_schema(
    name="submit_plan",
    description=(
        "Submit the orchestrator's final plan for this iteration. This MUST be "
        "the last tool you call. After you call it, the iteration ends and the "
        "step-writer takes over to produce the cinematic scene."
    ),
    params={
        "type": "object",
        "properties": {
            "assessment": {
                "type": "string",
                "description": "Brief assessment of the player's action (1-2 sentences, English).",
            },
            "outcome": {
                "type": "string",
                "description": "What concretely happened (1-3 sentences, English). Fed to the step-writer as the raw outcome.",
            },
            "state_patch": {
                "type": "object",
                "description": "JSON-patch for world state. Keys: set, unset, append, increment, remove. Empty object if no change.",
                "properties": {
                    "set": {"type": "object"},
                    "unset": {"type": "array", "items": {"type": "string"}},
                    "append": {"type": "object"},
                    "increment": {"type": "object"},
                    "remove": {"type": "object"},
                },
            },
            "time_advance": {
                "type": "object",
                "description": "How much in-world time advances. null/omitted if no time passes.",
                "properties": {
                    "days": {"type": "integer", "default": 0},
                    "hours": {"type": "integer", "default": 0},
                    "minutes": {"type": "integer", "default": 0},
                },
            },
            "narrative_prompt": {
                "type": "string",
                "description": "Facts the step-writer should know to write the scene (English). Max ~100 words.",
            },
            "next_options": {
                "type": "array",
                "items": {"type": "string"},
                "description": "3 suggested next actions for the player (short, 5-12 words each).",
            },
            "rails_update": {
                "type": "object",
                "description": "Optional update to plot rails. Omit if no change.",
                "properties": {
                    "main_goal": {"type": "string"},
                    "new_subgoals": {"type": "array", "items": {"type": "string"}},
                    "completed_subgoals": {"type": "array", "items": {"type": "string"}},
                },
            },
            "rag_facts": {
                "type": "array",
                "description": "New persistent facts to add to the glossary. Empty array if none.",
                "items": {
                    "type": "object",
                    "properties": {
                        "kind": {"type": "string", "enum": ["npc", "location", "item", "lore", "event"]},
                        "name": {"type": "string"},
                        "description": {"type": "string"},
                    },
                    "required": ["kind", "name", "description"],
                },
            },
        },
        "required": ["assessment", "outcome", "narrative_prompt", "next_options"],
    },
)


SUBMIT_SCENE_SCHEMA = build_tool_schema(
    name="submit_scene",
    description=(
        "Submit the narrative scene for this step. This is the ONLY way to "
        "return the scene — your text response is ignored. The narrative "
        "should be 200-400 words, cinematic, second-person ('You...'), in the "
        "world's player-facing language."
    ),
    params={
        "type": "object",
        "properties": {
            "narrative": {
                "type": "string",
                "description": "200-400 words of cinematic prose describing the scene. Second-person ('You...').",
            },
            "options": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Exactly 3 short (5-12 words) options for the player's next action.",
            },
        },
        "required": ["narrative", "options"],
    },
)


SUBMIT_WORLD_DEFINITION_SCHEMA = build_tool_schema(
    name="submit_world_definition",
    description=(
        "Submit a proposed world definition. Call this when you have enough "
        "information to build the world. Your text response will be shown to "
        "the player as your conversational reply (use it to summarize the "
        "proposed world in 2-4 sentences)."
    ),
    params={
        "type": "object",
        "properties": {
            "setting_description": {"type": "string", "description": "Expanded setting, 1-2 paragraphs."},
            "rules": {
                "type": "object",
                "description": "Object with keys like stats, combat, magic, time, inventory, death (whichever apply).",
            },
            "world_schema": {
                "type": "object",
                "description": "JSON Schema describing the shape of the world state.",
            },
            "plot_rails": {
                "type": "object",
                "description": "{main_goal, subgoals, hooks}.",
                "properties": {
                    "main_goal": {"type": "string"},
                    "subgoals": {"type": "array", "items": {"type": "string"}},
                    "hooks": {"type": "array", "items": {"type": "string"}},
                },
            },
            "initial_state": {
                "type": "object",
                "description": "Initial world state matching world_schema.",
            },
            "initial_time": {
                "type": "string",
                "description": "World time string e.g. 'day_1_hour_8'.",
            },
            "calendar": {
                "type": "object",
                "description": "Optional. Custom calendar. Include only if non-standard.",
                "properties": {
                    "hours_per_day": {"type": "integer"},
                    "minutes_per_hour": {"type": "integer"},
                    "days_per_week": {"type": "integer"},
                },
            },
            "is_final": {
                "type": "boolean",
                "description": "True ONLY when the player has explicitly accepted the world.",
            },
        },
        "required": ["setting_description", "rules", "world_schema", "initial_state", "initial_time"],
    },
)


SUBMIT_SUMMARY_SCHEMA = build_tool_schema(
    name="submit_summary",
    description="Submit the compressed summary of older session messages.",
    params={
        "type": "object",
        "properties": {
            "summary": {"type": "string", "description": "3-6 sentences, max 150 words."},
            "facts": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "kind": {"type": "string", "enum": ["npc", "location", "item", "lore", "event"]},
                        "name": {"type": "string"},
                        "description": {"type": "string"},
                    },
                    "required": ["kind", "name", "description"],
                },
            },
        },
        "required": ["summary", "facts"],
    },
)


SUBMIT_TRIGGER_RESULT_SCHEMA = build_tool_schema(
    name="submit_trigger_result",
    description="Submit the result of firing a deferred trigger.",
    params={
        "type": "object",
        "properties": {
            "outcome": {"type": "string", "description": "1-2 sentences, English. For logs."},
            "state_patch": {
                "type": "object",
                "description": "JSON-patch for world state. Empty object if no change.",
                "properties": {
                    "set": {"type": "object"},
                    "unset": {"type": "array", "items": {"type": "string"}},
                    "append": {"type": "object"},
                    "increment": {"type": "object"},
                    "remove": {"type": "object"},
                },
            },
            "narrative": {"type": "string", "description": "1-paragraph scene description for the player, in world.language. Empty string if offscreen."},
            "should_notify_player": {"type": "boolean", "description": "True if the player should see the narrative."},
        },
        "required": ["outcome", "narrative", "should_notify_player"],
    },
)


# Tools available to the orchestrator (game-loop tools + submit_plan)
ALL_TOOL_SCHEMAS = [
    DICE_ROLL_SCHEMA,
    UPDATE_STATE_SCHEMA,
    RAG_QUERY_SCHEMA,
    RAG_ADD_SCHEMA,
    SCHEDULE_TRIGGER_SCHEMA,
    ADVANCE_TIME_SCHEMA,
    RUN_SUBAGENT_SCHEMA,
    SUBMIT_PLAN_SCHEMA,
]

# Tools for the step writer (only submit_scene)
STEP_WRITER_TOOL_SCHEMAS = [SUBMIT_SCENE_SCHEMA]

# Tools for the world builder (only submit_world_definition)
WORLD_BUILDER_TOOL_SCHEMAS = [SUBMIT_WORLD_DEFINITION_SCHEMA]

# Tools for the summarizer
SUMMARIZER_TOOL_SCHEMAS = [SUBMIT_SUMMARY_SCHEMA]

# Tools for the trigger runner
TRIGGER_RUNNER_TOOL_SCHEMAS = [SUBMIT_TRIGGER_RESULT_SCHEMA]


# === Tool handlers ===

class ToolContext:
    """Holds everything tools need to execute."""
    def __init__(
        self,
        db: AsyncSession,
        world: World,
        session_id: uuid.UUID,
        user_id: uuid.UUID,
        subagent_runner: Optional[Callable[[str, str], Awaitable[str]]] = None,
        settings_map: Optional[Dict[str, Any]] = None,
    ):
        self.db = db
        self.world = world
        self.session_id = session_id
        self.user_id = user_id
        self.subagent_runner = subagent_runner
        self.settings_map = settings_map or {}
        # Track time advancement during this iteration
        self.time_advance: Dict[str, int] = {"days": 0, "hours": 0, "minutes": 0}
        # Track scheduled triggers
        self.scheduled_triggers: List[Dict[str, Any]] = []
        # Track rag facts added
        self.rag_added: List[Dict[str, Any]] = []


async def handle_tool_call(name: str, args: Dict[str, Any], ctx: ToolContext) -> Dict[str, Any]:
    if name == "dice_roll":
        sides = int(args.get("sides", 20))
        count = int(args.get("count", 1))
        modifier = int(args.get("modifier", 0))
        label = args.get("label", "")
        rolls = [random.randint(1, sides) for _ in range(max(1, count))]
        total = sum(rolls) + modifier
        return {"rolls": rolls, "modifier": modifier, "total": total, "label": label}

    if name == "update_state":
        patch = args.get("patch", {})
        new_state = apply_patch(ctx.world.state, patch)
        schema = ctx.world.definition.get("world_schema", {})
        ok, errors = validate_state(new_state, schema)
        if not ok:
            return {"ok": False, "errors": errors, "state_unchanged": True}
        ctx.world.state = new_state
        return {"ok": True, "new_state_summary": _summarize_state(new_state)}

    if name == "rag_query":
        query = args.get("query", "")
        limit = int(args.get("limit", 5))
        rag = await get_rag(ctx.settings_map)
        results = await rag.search_glossary(ctx.world.id, query, limit=limit, settings_map=ctx.settings_map)
        return {"results": results}

    if name == "rag_add":
        kind = args.get("kind", "lore")
        entry_name = args.get("name", "")
        desc = args.get("description", "")
        extra = args.get("payload", {}) or {}
        entry = GlossaryEntry(
            world_id=ctx.world.id,
            session_id=ctx.session_id,
            kind=kind,
            name=entry_name,
            description=desc,
            payload=extra,
        )
        ctx.db.add(entry)
        await ctx.db.flush()
        rag = await get_rag(ctx.settings_map)
        await rag.upsert_glossary(
            world_id=ctx.world.id,
            entry_id=entry.id,
            kind=kind,
            name=entry_name,
            description=desc,
            payload=extra,
            settings_map=ctx.settings_map,
        )
        ctx.rag_added.append({"kind": kind, "name": entry_name, "description": desc})
        return {"ok": True, "entry_id": str(entry.id)}

    if name == "schedule_trigger":
        fire_at = args.get("fire_at", "")
        description = args.get("description", "")
        payload = args.get("payload", {}) or {}
        trigger = DeferredTrigger(
            session_id=ctx.session_id,
            fire_at=fire_at,
            description=description,
            payload=payload,
        )
        ctx.db.add(trigger)
        await ctx.db.flush()
        ctx.scheduled_triggers.append({
            "id": str(trigger.id),
            "fire_at": fire_at,
            "description": description,
        })
        return {"ok": True, "trigger_id": str(trigger.id)}

    if name == "advance_time":
        days = int(args.get("days", 0))
        hours = int(args.get("hours", 0))
        minutes = int(args.get("minutes", 0))
        ctx.time_advance["days"] += days
        ctx.time_advance["hours"] += hours
        ctx.time_advance["minutes"] += minutes
        return {
            "ok": True,
            "advance": {"days": days, "hours": hours, "minutes": minutes},
            "reason": args.get("reason", ""),
        }

    if name == "run_subagent":
        if ctx.subagent_runner is None:
            return {"error": "subagent_runner_not_available"}
        task = args.get("task", "")
        context = args.get("context", "")
        try:
            result = await ctx.subagent_runner(task, context)
            return {"result": result}
        except Exception as e:
            return {"error": str(e)}

    return {"error": f"unknown_tool: {name}"}


def _summarize_state(state: Dict[str, Any]) -> str:
    """Quick human-readable summary of state for the LLM."""
    if not state:
        return "(empty)"
    parts: List[str] = []
    player = state.get("player", {})
    if player:
        name = player.get("name", "?")
        stats = player.get("stats", {})
        location = player.get("location", "?")
        hp = stats.get("health", "?")
        hp_max = stats.get("health_max", "?")
        mp = stats.get("mana", "?")
        parts.append(f"player={name} hp={hp}/{hp_max} mp={mp} loc={location}")
    inv = player.get("inventory", []) if isinstance(player, dict) else []
    if inv:
        parts.append("inv=" + ", ".join(f"{i.get('name','?')}x{i.get('qty',1)}" for i in inv[:8]))
    npcs = state.get("npcs", [])
    if npcs:
        parts.append(f"npcs={len(npcs)}")
    return " | ".join(parts)