mind/planner.py

"""The Planner — GODSEED's plan-act-observe loop (Agent protocol, ARCHITECTURE.md).

Flow per wish:
  1. Reading — free-text generation streamed token-by-token via
     emit({"type": "thought_token", ...}), capped at <=80 tokens.
  2. Turn loop — at most 7 grammar-constrained JSON turns. Each turn is
     either a tool call (executed through `act`, observation appended to
     context) or a done turn carrying the epitaph. Malformed output gets one
     re-ask with the parse error in context; if that also fails the turn is
     skipped and recorded. Two consecutive skipped turns abort the loop (a
     dead backend should not burn 7 turns).

Division of labor for the WishTrace (load-bearing — see ARCHITECTURE.md):
  - The planner returns ONLY the fields it can know:
        reading, turns, epitaph, ms_total, model, backend
  - The queue worker (engine/queue_worker.py, Agent A) merges the rest:
        wish_id, text, submitted_at, moderation, feature_ids
    Feature ids are assigned by the engine when `act` executes a call; `act`
    returns just the observation string, so the worker — not the planner —
    tracks which features each wish produced.
  - Likewise, emitted events carry no wish_id here; the worker wraps `emit`
    and injects it before broadcast.

Events emitted (in order):
    thought_token*               (the reading, chunk by chunk, then "\n\n")
    [per call turn] thought_token*
    [done turn]     thought_token*
The planner emits ONLY thought_token. tool_call and world_delta are emitted by
the queue worker after the engine accepts a call — the worker's copy carries
the wish_id and the engine's canonical (clamped) args, so a second planner-side
emission would double the god's-gaze camera move in the renderer.
"""

from __future__ import annotations

import time
import zlib
from typing import Awaitable, Callable

from . import prompts
from .backends import make_backend
from .validate import TURN_GRAMMAR, parse_plan, parse_turn

Act = Callable[[dict], Awaitable[str]]
Emit = Callable[[dict], Awaitable[None]]

_FALLBACK_READING = "The god listens, says nothing for a long moment, and begins."

_FALLBACK_EPITAPHS = (
    "The god has done what it set out to do. The rest is weather.",
    "Enough. The world will finish the sentence on its own.",
    "It is made, as far as making goes tonight.",
)

_READING_MAX_CHARS = 700
_TURN_RAW_MAX_CHARS = 4000
_OBSERVATION_MAX_CHARS = 500
_ONESHOT_RAW_MAX_CHARS = 6000


def _fallback_epitaph(wish: str) -> str:
    index = zlib.crc32((wish or "").encode("utf-8")) % len(_FALLBACK_EPITAPHS)
    return _FALLBACK_EPITAPHS[index]


def _epitaph_tail(reading: str) -> str:
    """The reading's last sentence — a varied epitaph when the model omits one."""
    import re as _re
    parts = [p.strip() for p in _re.split(r"(?<=[.!?])\s+", str(reading or "")) if p.strip()]
    return parts[-1][:120] if parts else ""


def _thought_for(call: dict) -> str:
    """A short liturgical thought for a planned act (the one-shot model gives a
    plan, not per-call thoughts, so the god still 'narrates' as each lands)."""
    tool = (call or {}).get("tool", "")
    a = (call or {}).get("args", {}) or {}
    kind = a.get("kind")
    phrases = {
        "build_district": "A quarter of homes rises.",
        "place_structure": f"A {kind} takes its place." if kind else "A landmark rises.",
        "place_road": "A road to bind them.",
        "place_water": "Water, to soften the stone.",
        "spawn_flora": "Green, where there was none.",
        "spawn_life": "And life begins to move through it.",
        "raise_terrain": "First, ground to stand on.",
        "lower_terrain": "The land sinks here.",
        "set_sky": "The sky turns.",
        "set_weather": "The weather answers.",
        "inscribe_wish": "Words for those who come after.",
    }
    return phrases.get(tool, "It is made.")


def _fallback_town_plan(world_summary: str) -> list[dict]:
    """A deterministic small town near the world's center — the safety net so a
    wish NEVER builds nothing, even if the model's JSON was unusable."""
    import re as _re
    lat, lon = 14.0, 38.0  # genesis monolith / default town seat
    m = _re.search(r"near \((-?\d+(?:\.\d+)?),\s*(-?\d+(?:\.\d+)?)\)", str(world_summary or ""))
    if m:
        lat, lon = float(m.group(1)), float(m.group(2))
    return [
        {"tool": "build_district", "args": {"lat": lat, "lon": lon, "radius_deg": 7, "density": 0.8, "hue": 40}},
        {"tool": "place_structure", "args": {"lat": lat, "lon": lon + 1, "kind": "market", "scale": 1.2, "hue": 45}},
        {"tool": "place_structure", "args": {"lat": lat + 1, "lon": lon, "kind": "bank", "scale": 1.1, "hue": 50}},
        {"tool": "place_road", "args": {"path": [[lat, lon], [lat, lon + 1], [lat + 1, lon]]}},
        {"tool": "spawn_life", "args": {"lat": lat, "lon": lon, "radius_deg": 6, "kind": "carts", "count": 5, "hue": 40}},
    ]


class Planner:
    """Grants one wish at a time against an injected backend.

    The planner owns prompting, parsing, and pacing. It never touches world
    state: every effect goes through `act`, every visible word through `emit`.
    """

    MAX_TURNS = 7
    READING_MAX_TOKENS = 80
    TURN_MAX_TOKENS = 220
    ONESHOT_MAX_TOKENS = 700  # reading + a full plan JSON + epitaph, one call
    MAX_CONSECUTIVE_SKIPS = 2

    def __init__(self, backend=None, *, max_turns: int | None = None):
        self.backend = backend if backend is not None else make_backend()
        self.max_turns = max_turns or self.MAX_TURNS

    # ------------------------------------------------------------------ emit

    @staticmethod
    async def _safe_emit(emit: Emit, event: dict) -> None:
        """Broadcast failures must never kill a grant."""
        try:
            await emit(event)
        except Exception:
            pass

    async def _emit_text(self, emit: Emit, text: str, *, end: str = "\n") -> None:
        """Stream already-known text (turn thoughts) as small chunks."""
        words = text.split()
        step = 3
        for i in range(0, len(words), step):
            chunk = " ".join(words[i : i + step])
            if i + step < len(words):
                chunk += " "
            await self._safe_emit(emit, {"type": "thought_token", "text": chunk})
        if end:
            await self._safe_emit(emit, {"type": "thought_token", "text": end})

    # --------------------------------------------------------------- reading

    async def _reading(self, wish: str, world_summary: str, emit: Emit) -> str:
        prompt = prompts.build_reading_prompt(wish, world_summary)
        parts: list[str] = []
        total_chars = 0
        token_count = 0
        stream = self.backend.generate_stream(
            prompt, None, self.READING_MAX_TOKENS
        )
        try:
            async for chunk in stream:
                if not chunk:
                    continue
                parts.append(chunk)
                total_chars += len(chunk)
                token_count += max(1, len(chunk.split()))
                await self._safe_emit(emit, {"type": "thought_token", "text": chunk})
                if token_count >= self.READING_MAX_TOKENS or total_chars >= _READING_MAX_CHARS:
                    break
        except Exception:
            pass  # fall through to the fallback reading
        finally:
            try:
                await stream.aclose()
            except Exception:
                pass

        reading = "".join(parts).strip()
        if not reading:
            reading = _FALLBACK_READING
            # The backend gave us nothing; keep the stream alive for watchers.
            await self._emit_text(emit, reading, end="")
        await self._safe_emit(emit, {"type": "thought_token", "text": "\n\n"})
        return reading

    # ----------------------------------------------------------------- turns

    async def _collect(self, prompt: str, max_tokens: int) -> str:
        parts: list[str] = []
        total = 0
        stream = self.backend.generate_stream(prompt, TURN_GRAMMAR, max_tokens)
        try:
            async for chunk in stream:
                if not chunk:
                    continue
                parts.append(chunk)
                total += len(chunk)
                if total >= _TURN_RAW_MAX_CHARS:
                    break
        finally:
            try:
                await stream.aclose()
            except Exception:
                pass
        return "".join(parts)

    async def _attempt_turn(self, prompt: str) -> tuple[dict | None, str | None]:
        try:
            raw = await self._collect(prompt, self.TURN_MAX_TOKENS)
        except Exception as exc:
            return None, f"backend failure: {type(exc).__name__}"
        return parse_turn(raw)

    async def _next_turn(
        self, wish: str, world_summary: str, reading: str, turns: list[dict]
    ) -> tuple[dict | None, str | None]:
        """One turn with the single re-ask on malformed output."""
        prompt = prompts.build_turn_prompt(
            wish, world_summary, reading, turns, self.max_turns
        )
        obj, err = await self._attempt_turn(prompt)
        if err is None:
            return obj, None
        reask = prompts.build_turn_prompt(
            wish, world_summary, reading, turns, self.max_turns, error=err
        )
        return await self._attempt_turn(reask)

    # ------------------------------------------------------------- one-shot

    async def _grant_oneshot(
        self, wish: str, world_summary: str, act: Act, emit: Emit, started: float
    ) -> dict:
        """Whole wish in ONE generation: reading + full plan + epitaph. Then the
        engine executes the calls (CPU). One GPU call total — ZeroGPU-safe."""
        prompt = prompts.build_oneshot_prompt(wish, world_summary)
        try:
            raw = await self._collect_free(prompt, self.ONESHOT_MAX_TOKENS)
        except Exception:
            raw = ""  # generation failed; the fallback town plan still builds

        reading, calls, epitaph = parse_plan(raw)
        reading = reading.strip() or _FALLBACK_READING
        # Stream the reading so the god still "speaks aloud".
        await self._emit_text(emit, reading, end="\n\n")

        # Safety net: a wish must never build nothing. If the model gave no
        # usable calls, found a small town near the world's center so the
        # wisher always sees something rise.
        if not calls:
            calls = _fallback_town_plan(world_summary)

        turns: list[dict] = []
        for call in calls[: self.max_turns]:
            thought = _thought_for(call)
            if thought:
                await self._emit_text(emit, thought)
            try:
                observation = await act(call)
            except Exception as exc:
                observation = f"error: act failed ({type(exc).__name__})"
            turns.append({
                "thought": thought,
                "call": call,
                "observation": str(observation or "")[:_OBSERVATION_MAX_CHARS],
            })

        epitaph = epitaph.strip() or _epitaph_tail(reading) or _fallback_epitaph(wish)
        return {
            "reading": reading,
            "turns": turns,
            "epitaph": epitaph,
            "ms_total": int((time.perf_counter() - started) * 1000),
            "model": str(getattr(self.backend, "model_id", "unknown")),
            "backend": str(getattr(self.backend, "name", "unknown")),
        }

    async def _collect_free(self, prompt: str, max_tokens: int) -> str:
        """One free-text generation (no grammar), fully collected."""
        parts: list[str] = []
        total = 0
        stream = self.backend.generate_stream(prompt, None, max_tokens)
        try:
            async for chunk in stream:
                if not chunk:
                    continue
                parts.append(chunk)
                total += len(chunk)
                if total >= _ONESHOT_RAW_MAX_CHARS:
                    break
        finally:
            try:
                await stream.aclose()
            except Exception:
                pass
        return "".join(parts)

    # ----------------------------------------------------------------- grant

    async def grant(
        self, wish: str, world_summary: str, act: Act, emit: Emit
    ) -> dict:
        """Grant one wish. Returns the planner's share of the WishTrace.

        Returned dict (worker merges wish_id/text/submitted_at/moderation/
        feature_ids — see module docstring):

            {"reading": str,
             "turns": [{"thought": str, "call": dict | None,
                        "observation": str | None}, ...],
             "epitaph": str, "ms_total": int, "model": str, "backend": str}

        Turn records: executed calls carry their observation string; the
        final done turn has call=None and observation=None; a skipped
        (malformed) turn has call=None and an observation starting "error:".
        This method never raises for backend/act misbehavior — it always
        returns a complete trace.
        """
        started = time.perf_counter()
        wish = str(wish or "")
        world_summary = str(world_summary or "")

        # ZeroGPU: a wish must be ONE @spaces.GPU call — the per-turn loop makes
        # ~8 and ZeroGPU detaches the GPU between them (NVML crash). One-shot
        # backends generate reading + full plan in a single call.
        if getattr(self.backend, "oneshot", False):
            return await self._grant_oneshot(wish, world_summary, act, emit, started)

        reading = await self._reading(wish, world_summary, emit)

        turns: list[dict] = []
        epitaph: str | None = None
        calls_executed = 0
        consecutive_skips = 0

        for _ in range(self.max_turns):
            obj, err = await self._next_turn(wish, world_summary, reading, turns)
            if obj is None:
                turns.append({
                    "thought": "",
                    "call": None,
                    "observation": f"error: {err or 'malformed output'}; turn skipped",
                })
                consecutive_skips += 1
                if consecutive_skips >= self.MAX_CONSECUTIVE_SKIPS:
                    break
                continue
            consecutive_skips = 0

            await self._emit_text(emit, obj["thought"])

            if obj.get("done"):
                epitaph = obj["epitaph"].strip() or _fallback_epitaph(wish)
                turns.append({
                    "thought": obj["thought"],
                    "call": None,
                    "observation": None,
                })
                break

            call = obj["call"]
            try:
                observation = await act(call)
            except Exception as exc:
                observation = f"error: act failed ({type(exc).__name__})"
            observation = str(observation or "")[:_OBSERVATION_MAX_CHARS]
            calls_executed += 1
            turns.append({
                "thought": obj["thought"],
                "call": call,
                "observation": observation,
            })

        if epitaph is None:
            epitaph = _fallback_epitaph(wish)

        return {
            "reading": reading,
            "turns": turns,
            "epitaph": epitaph,
            "ms_total": int((time.perf_counter() - started) * 1000),
            "model": str(getattr(self.backend, "model_id", "unknown")),
            "backend": str(getattr(self.backend, "name", "unknown")),
        }