pf1e-simulator/src/pf1e_simulator/prompt_builder.py

"""Prompt builder for the phase-2 LLM layer (P1.4).

Builds a deterministic prompt from an ``Observation`` plus the external
prompt pack: template + compact observation + plausible actions + JSON
action schema. The schema is derived from the ``Action`` dataclass via
``typing.get_type_hints`` so it stays in sync with the engine.
"""

from __future__ import annotations

import json
from dataclasses import MISSING, asdict, dataclass, fields
from types import UnionType
from typing import TYPE_CHECKING, Literal, Union, get_args, get_origin, get_type_hints

from pf1e_simulator.candidates import candidates
from pf1e_simulator.combat import UNSUPPORTED_ACTION_KINDS, Action, Observation
from pf1e_simulator.llm_view import compact_observation

if TYPE_CHECKING:
    from pf1e_simulator.prompt import PromptPack


def _supported_kinds() -> tuple[str, ...]:
    """Return the 27 action kinds supported by the LLM layer (declaration order)."""
    kind_literal = get_type_hints(Action)["kind"]
    all_kinds = get_args(kind_literal)
    return tuple(k for k in all_kinds if k not in UNSUPPORTED_ACTION_KINDS)


_EMPTY_LITERAL_MSG = "empty Literal"
_UNSUPPORTED_LITERAL_MSG = "unsupported Literal values: {}"
_MULTI_UNION_MSG = "multi-type union not supported in Action: {}"
_UNSUPPORTED_TYPE_MSG = "unsupported type in Action schema: {}"
_TUPLE_ELLIPSIS_LEN = 2


def _json_schema_base(tp: object) -> dict[str, object] | None:
    """Return the JSON Schema fragment for a base type, else None."""
    if tp is type(None):
        return {"type": "null"}
    if tp is str:
        return {"type": "string"}
    if tp is int:
        return {"type": "integer"}
    if tp is bool:
        return {"type": "boolean"}
    return None


def _json_schema_literal(args: tuple[object, ...]) -> dict[str, object]:
    """Return the JSON Schema fragment for a Literal type."""
    if not args:
        raise ValueError(_EMPTY_LITERAL_MSG)
    first = args[0]
    if isinstance(first, str):
        base_type = "string"
    elif isinstance(first, int):
        base_type = "integer"
    else:
        msg = _UNSUPPORTED_LITERAL_MSG.format(args)
        raise TypeError(msg)
    return {"type": base_type, "enum": list(args)}


def _json_schema_tuple(args: tuple[object, ...]) -> dict[str, object]:
    """Return the JSON Schema fragment for a tuple type."""
    if not args:
        return {"type": "array"}
    if len(args) == _TUPLE_ELLIPSIS_LEN and args[1] is Ellipsis:
        # tuple[T, ...] → homogeneous array
        item_schema = _json_schema_for(args[0])
        return {"type": "array", "items": item_schema}
    # Fixed-length tuple → prefixItems + minItems/maxItems
    prefix = [_json_schema_for(arg) for arg in args]
    return {
        "type": "array",
        "prefixItems": prefix,
        "minItems": len(args),
        "maxItems": len(args),
    }


def _json_schema_union(args: tuple[object, ...]) -> dict[str, object]:
    """Return the JSON Schema fragment for a nullable single-member union."""
    non_none = [a for a in args if a is not type(None)]
    if len(non_none) == 1:
        member_schema = _json_schema_for(non_none[0])
        # Merge "null" into the type array
        member_type = member_schema.get("type")
        if isinstance(member_type, str):
            member_schema["type"] = [member_type, "null"]
        # If member already has an enum (e.g. nullable Literal), keep it.
        return member_schema
    msg = _MULTI_UNION_MSG.format(args)
    raise ValueError(msg)


def _json_schema_for(tp: object) -> dict[str, object]:
    """Convert a Python type annotation to a JSON Schema fragment."""
    base_schema = _json_schema_base(tp)
    if base_schema is not None:
        return base_schema

    origin = get_origin(tp)
    args = get_args(tp)

    if origin is Literal:
        return _json_schema_literal(args)

    if origin is tuple:
        return _json_schema_tuple(args)

    if origin is Union or origin is UnionType:
        return _json_schema_union(args)

    msg = _UNSUPPORTED_TYPE_MSG.format(tp)
    raise ValueError(msg)


def action_schema() -> dict[str, object]:
    """JSON Schema for a single Action object."""
    hints = get_type_hints(Action)
    properties = {name: _json_schema_for(tp) for name, tp in hints.items()}
    # Override kind with the filtered 27-kind enum (source Literal order)
    properties["kind"] = {"type": "string", "enum": list(_supported_kinds())}
    # Only 'kind' lacks a default
    required = [
        f.name
        for f in fields(Action)
        if f.default is MISSING and f.default_factory is MISSING
    ]
    return {
        "type": "object",
        "properties": properties,
        "required": required,
        "additionalProperties": False,
    }


def render_candidates(actions: tuple[Action, ...]) -> tuple[str, ...]:
    """Render the plausible-action menu as numbered lines with compact JSON."""
    lines: list[str] = []
    for index, action in enumerate(actions):
        payload = {
            key: value
            for key, value in asdict(action).items()
            if key != "kind" and value is not None and value != () and value is not False
        }
        suffix = f" {json.dumps(payload, sort_keys=True, separators=(',', ':'))}" if payload else ""
        lines.append(f"{index}. {action.kind}{suffix}")
    return tuple(lines)


@dataclass(frozen=True)
class PromptBuilder:
    """Builds the complete LLM prompt from a prompt pack and an observation."""

    pack: PromptPack

    def prompt_version(self, name: str) -> str:
        """Pack version plus per-prompt content hash (cache-key component, D6)."""
        return self.pack.prompt_version(name)

    def build(self, observation: object, *, prompt_name: str = "nl_default") -> str:
        """Assemble the full prompt for the given observation."""
        if not isinstance(observation, Observation):
            msg = "observation must be an Observation instance"
            raise TypeError(msg)

        template = self.pack.get(prompt_name)
        obs_text = compact_observation(observation).rstrip("\n")
        menu = candidates(observation)
        actions_text = "\n".join(render_candidates(menu))
        schema_text = json.dumps(action_schema(), sort_keys=True, indent=2)
        return "\n\n".join([
            template,
            "## observation\n" + obs_text,
            "## plausible actions\n" + actions_text,
            "## action schema\n" + schema_text,
        ])