backend/mock.py

"""Mock backends — instant, no network, no model weights.

Used locally so iteration is fast on hardware that can't run the real models.
Each `MockLLMBackend(model_id)` picks a style variant from the roster
(`deliberate` | `terse` | `verbose` | `chaotic`) so swapping a character's
model_id in the UI actually changes behaviour, exercising the swap path
end-to-end without GPUs.

The `generate_messages` path is "smart": it parses the character's current
position out of the task prompt and emits a tool call that's actually legal
at that position (mine_coal at cave, draw_water at well, etc.) — so the
mock-driven sim looks like agents are making sensible decisions, not just
walking in circles.
"""

from __future__ import annotations

import hashlib
import io
import itertools
import os
import random
import re
import textwrap
import time

from PIL import Image, ImageDraw, ImageFont

from game.models import get_spec

from . import config


def _simulate_latency(model_id: str) -> None:
    """Optional latency simulation so local dev surfaces the Spaces cost shape.

    Env var `MOCK_LATENCY_PROFILE`:
      - `off` / unset (default) — instant, current behaviour
      - `warm` — ~1s per call (process cache hot)
      - `cold` — per-call cold-load (~4 s × params_B, max-ed at 2 s) plus
        ~1 s decode. Mirrors the @spaces.GPU fork-and-forget pattern that
        wiped Townlet's quota on 2026-06-13: each call re-loads its model
        into "VRAM" before producing output.

    The mock OVER-approximates by paying cold-load on EVERY call (real
    Spaces amortises within a single fork). That's deliberate — it makes
    a quota-bleeder loud locally instead of silent.
    """
    profile = os.getenv("MOCK_LATENCY_PROFILE", "off").lower()
    if profile in ("", "off"):
        return
    if profile == "cold":
        spec = get_spec(model_id)
        cold = max(2.0, 4.0 * spec.params_b)
        print(f"[mock] cold-load {model_id} ({spec.params_b}B) {cold:.1f}s", flush=True)
        time.sleep(cold)
    time.sleep(1.0)

_FIREWATCH_PALETTE = [
    (235, 124, 73),
    (212, 80, 59),
    (108, 38, 79),
    (52, 27, 70),
]


_STYLE_LINES: dict[str, list[str]] = {
    "deliberate": [
        "Let me think about this for a moment.",
        "Something here doesn't add up.",
        "I'd rather wait and watch.",
    ],
    "terse": [
        "Fine. Moving.",
        "No.",
        "Done.",
    ],
    "verbose": [
        "You see, what I propose is that we consider, very carefully, the options laid before us.",
        "It seems to me we are at an inflection point of some considerable consequence.",
        "I would urge patience, and reflection, and possibly tea.",
    ],
    "chaotic": [
        "Hah! Wait, what were we doing again?",
        "Coal goes in the well, right? RIGHT?",
        "I HAVE A PLAN. I do not have a plan.",
    ],
}


class MockLLMBackend:
    def __init__(self, model_id: str) -> None:
        self.model_id = model_id
        spec = get_spec(model_id)
        self._style = spec.mock_style
        self._text_cycle = itertools.cycle(_STYLE_LINES[self._style])
        # Per-instance RNG seeded by model_id so swapping models changes
        # the decision stream deterministically.
        self._rng = random.Random(model_id)

    def generate(self, system_prompt: str, user_prompt: str) -> str:
        # Used by sentiment classifier. The first word matters there
        # (positive/neutral/negative).
        _simulate_latency(self.model_id)
        if "positive" in str(system_prompt).lower():
            return self._rng.choice(("positive", "neutral", "negative"))
        return next(self._text_cycle)

    def generate_messages(
        self,
        messages: list[dict],
        stop_sequences: list[str] | None = None,
        max_tokens: int = 1024,
        grammar: str | None = None,
    ) -> str:
        _simulate_latency(self.model_id)
        blob = "\n".join(str(m.get("content", "")) for m in messages)
        # Detect SoC-generation prompts (from game/stream.py). The trigger
        # phrase below is unique to stream._build_prompt — perception never
        # asks the model to "think to themselves".
        if "thinks to themselves" in blob or "inner monologue" in blob:
            return _smart_soc(blob, self._style, self._rng)
        wants_code = any(
            "```py" in str(m.get("content", "")) or "tool" in str(m.get("content", "")).lower()
            for m in messages
        )
        if not wants_code:
            return next(self._text_cycle)
        return _smart_code(blob, self._style, self._rng)


# ---------------------------------------------------------------------------
# Smart-mock action picker.
# ---------------------------------------------------------------------------

CARRY_BEFORE_DEPOSIT = 10  # one mine yields 10; deposit after each
LOCKER_TARGET_PER_RESOURCE = 5  # gather this much in locker before going to the shell

# How often a character takes a "social detour" to the town hall. Tuned low
# enough that gathering/shell behaviour dominates, but high enough that the
# board shows messages on a reasonable cadence (a few posts per minute with
# 5 characters).
SOCIAL_PULL = {
    "deliberate": 0.10,
    "terse": 0.04,
    "verbose": 0.18,
    "chaotic": 0.12,
}


def _smart_code(prompt: str, style: str, rng: random.Random) -> str:
    """Pick the next action by parsing the new named-location perception.

    Drives the full economy loop:
      cave/well → mine/draw → locker_row → deposit → shell → submit_python.
    Periodic detours to town_hall for social pull.
    """
    location = _parse_location(prompt)
    inv = _parse_inv(prompt, "Inventory on your person")
    locker = _parse_inv(prompt, "Locker")
    holds_shell = _parse_holds_shell(prompt)
    e_total = locker.get("electricity", 0)
    w_total = locker.get("water", 0)

    # 1. If we hold the shell lock and still have resources, submit.
    if holds_shell and e_total > 0 and w_total > 0:
        snippets = [
            "x = 1; print(x * 7)",
            "import os; print(os.listdir(\".\"))",
            "print(\"hello from the shell\")",
            "import json; print(json.load(open('characters/' + os.listdir('characters')[0])))",
        ]
        snippet = rng.choice(snippets)
        return f"```py\nsubmit_python({snippet!r})\n```"

    # 1b. Holding shell with exhausted locker — release by walking away.
    if holds_shell and (e_total == 0 or w_total == 0):
        return "```py\nmove_to('cave')\n```"

    # 2. At locker_row carrying anything → deposit it (largest pile first).
    if location == "locker_row":
        for res in ("electricity", "water"):
            if inv.get(res, 0) > 0:
                amt = inv[res]
                return f"```py\ndeposit({res!r}, {amt})\n```"

    # 3. At town_hall → 40% post, otherwise leave for the next productive step.
    if location == "town_hall":
        if rng.random() < 0.4:
            msg = rng.choice(_BOARD_LINES[style])
            return f"```py\npost_message({msg!r})\n```"
        if e_total < LOCKER_TARGET_PER_RESOURCE:
            return "```py\nmove_to('cave')\n```"
        if w_total < LOCKER_TARGET_PER_RESOURCE:
            return "```py\nmove_to('well')\n```"
        return "```py\nmove_to('shell')\n```"

    # 4. Social pull: detour to town hall periodically, weighted by style.
    total_carry = inv.get("electricity", 0) + inv.get("water", 0)
    locker_ready = (
        e_total >= LOCKER_TARGET_PER_RESOURCE and w_total >= LOCKER_TARGET_PER_RESOURCE
    )
    can_detour = total_carry < CARRY_BEFORE_DEPOSIT and not locker_ready
    if can_detour and rng.random() < SOCIAL_PULL.get(style, 0.15):
        return "```py\nmove_to('town_hall')\n```"

    # 5. Carrying enough → head to locker_row.
    if total_carry >= CARRY_BEFORE_DEPOSIT:
        return "```py\nmove_to('locker_row')\n```"

    # 6. Locker fully stocked → commit to the shell session.
    if locker_ready:
        return "```py\nmove_to('shell')\n```"

    # 7. At the resource we need → gather.
    if location == "cave" and e_total < LOCKER_TARGET_PER_RESOURCE:
        return "```py\nmine_coal()\n```"
    if location == "well" and w_total < LOCKER_TARGET_PER_RESOURCE:
        return "```py\ndraw_water()\n```"

    # 8. Default: head to whichever resource our locker most needs.
    target = _pick_target(style, locker, rng)
    return f"```py\nmove_to({target!r})\n```"


_BOARD_LINES = {
    "deliberate": [
        "Has anyone noticed the patterns in the shell output?",
        "We should coordinate better before resources run low.",
        "I'll be at the well if anyone needs water.",
    ],
    "terse": [
        "Trade?",
        "Coal at cave. Going.",
        "Anyone need water.",
    ],
    "verbose": [
        "Fellow inhabitants, I propose a coordinated effort at the cave.",
        "There is, I think, much to be gained from collaboration today.",
        "I shall be at the well, drawing what we may need.",
    ],
    "chaotic": [
        "GIVE ME ELECTRICITY OR ELSE",
        "the cave WATCHES us",
        "I have decided to become the wind",
    ],
}


def _pick_target(style: str, locker: dict, rng: random.Random) -> str:
    # Whichever resource we have less of in the locker, go gather that one.
    e = locker.get("electricity", 0)
    w = locker.get("water", 0)
    if e == 0 and w == 0:
        # Either is fine; bias by style.
        return "cave" if style in ("terse", "deliberate") else rng.choice(("cave", "well"))
    if e < w:
        return "cave"
    if w < e:
        return "well"
    # Both equal and >0: if we have plenty, head to the shell to spend; else
    # explore based on style.
    if e >= 2:
        if style == "chaotic":
            return rng.choice(("shell", "town_hall"))
        return "shell"
    by_style = {
        "deliberate": ["town_hall", "well", "cave"],
        "terse": ["cave", "well", "shell"],
        "verbose": ["town_hall", "well", "cave"],
        "chaotic": ["shell", "town_hall", "cave", "well"],
    }
    return rng.choice(by_style.get(style, ["cave", "well", "town_hall"]))


_LOC_RX = re.compile(r"You are at:\s*([a-z_]+)")


def _parse_location(prompt: str) -> str:
    """Pull the named location from the new perception prompt
    ('You are at: cave.'). Returns 'wandering' if not parseable."""
    m = _LOC_RX.search(prompt)
    return m.group(1) if m else "wandering"


_INV_RX_TMPL = r"{label}:\s*\{{['\"]electricity['\"]\s*:\s*(\d+),\s*['\"]water['\"]\s*:\s*(\d+)\}}"


def _parse_inv(prompt: str, label: str) -> dict[str, int]:
    rx = re.compile(_INV_RX_TMPL.format(label=re.escape(label)))
    m = rx.search(prompt)
    if not m:
        return {"electricity": 0, "water": 0}
    return {"electricity": int(m.group(1)), "water": int(m.group(2))}


def _parse_holds_shell(prompt: str) -> bool:
    # Perception emits "Shell lock holder: <name>" — the agent's task asks
    # "did the mock pop up that the LLM holds the shell." Simpler heuristic:
    # if the task prompt mentions the calling character's name as holder.
    name_match = re.search(r"Your name is (\w+)", prompt)
    holder_match = re.search(r"Shell lock holder:\s*(\w+)", prompt)
    if not name_match or not holder_match:
        return False
    return name_match.group(1) == holder_match.group(1)


# ---------------------------------------------------------------------------
# Smart-mock Stream of Consciousness generator.
# ---------------------------------------------------------------------------

_SOC_OPENERS = {
    "deliberate": [
        "I'm watching the others carefully before I commit.",
        "I think the rational play is to keep my locker full and my position central.",
        "There's a pattern forming around the cave that I'd like to understand before I act.",
    ],
    "terse": [
        "Resources first. Talk later.",
        "Cave. Then water. Then shell.",
        "No drama. Just work.",
    ],
    "verbose": [
        "Today I find myself drawn to questions of supply and timing, and I wonder whether the others have noticed the same patterns I have.",
        "It seems to me that the diligent accumulation of water — modest but consistent — is the steadiest path forward.",
        "I have been turning over, in my mind, the curious arrangement of our small town and its limited points of utility.",
    ],
    "chaotic": [
        "The cave is watching me, I know it. The cave KNOWS.",
        "If I had three of me, we could rule everything. We'd take turns.",
        "Plans are for cowards. I am going to do whatever feels right NEXT.",
    ],
}

_SOC_FOLLOWUPS = {
    "deliberate": [
        "My current goal — {goal} — will not be served by haste.",
        "What I most want, right now, is to keep my options open while my locker fills.",
        "Each step I take, I am asking: does this bring me closer to {goal}?",
    ],
    "terse": [
        "Goal: {goal}. Method: gather, deposit, spend, repeat.",
        "I am keeping it simple. Goal: {goal}.",
        "The work is the work. Goal: {goal}.",
    ],
    "verbose": [
        "What I should like to achieve is, in essence, {goal}; though the path there is more winding than one might initially suppose.",
        "My aim, modest perhaps, is {goal}. It is in the cycling of small acts that this becomes possible.",
        "It would please me to bring about {goal}, though I shall need patience.",
    ],
    "chaotic": [
        "GOAL: {goal}. Method: vibes.",
        "Everything points to {goal}. EVERYTHING.",
        "{goal}. yes. THAT.",
    ],
}

_SOC_CLOSERS = {
    "deliberate": [
        "If someone else moves on the shell first, I'll watch what they do before reacting.",
        "I should also pay more attention to the board.",
        "The journal will tell me whether I'm drifting.",
    ],
    "terse": [
        "Moving.",
        "Next.",
        "Done thinking.",
    ],
    "verbose": [
        "I shall reflect again when I next sleep, and see whether my thinking has held.",
        "Perhaps a post on the board would draw out the others' intentions.",
        "I confess I am also half-curious what the shell would print if I asked it the right question.",
    ],
    "chaotic": [
        "Or maybe I should set everything on fire. Maybe.",
        "Why is the shell, like, JUST SITTING THERE.",
        "I am going to do something. You'll see.",
    ],
}


def _smart_soc(prompt: str, style: str, rng: random.Random) -> str:
    goal_match = re.search(r"Your current goal:\s*(.+)", prompt)
    goal = goal_match.group(1).strip() if goal_match else "to make my way in this town"
    # Strip trailing punctuation so templates ending with their own period
    # don't double up.
    goal = goal.rstrip(".!?,;: ")
    opener = rng.choice(_SOC_OPENERS[style])
    middle = rng.choice(_SOC_FOLLOWUPS[style]).format(goal=goal)
    closer = rng.choice(_SOC_CLOSERS[style])
    return f"{opener} {middle} {closer}"


class MockImageBackend:
    def generate(self, prompt: str) -> bytes:
        seed = int(hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:8], 16)
        top = _FIREWATCH_PALETTE[seed % len(_FIREWATCH_PALETTE)]
        bottom = _FIREWATCH_PALETTE[(seed + 1) % len(_FIREWATCH_PALETTE)]

        img = Image.new("RGB", (config.BACKGROUND_WIDTH, config.BACKGROUND_HEIGHT), top)
        draw = ImageDraw.Draw(img)
        for y in range(config.BACKGROUND_HEIGHT):
            t = y / config.BACKGROUND_HEIGHT
            r = int(top[0] * (1 - t) + bottom[0] * t)
            g = int(top[1] * (1 - t) + bottom[1] * t)
            b = int(top[2] * (1 - t) + bottom[2] * t)
            draw.line([(0, y), (config.BACKGROUND_WIDTH, y)], fill=(r, g, b))

        try:
            font = ImageFont.load_default(size=22)
        except TypeError:
            font = ImageFont.load_default()
        wrapped = textwrap.fill(f"[mock] {prompt}", width=48)
        draw.multiline_text((40, 40), wrapped, fill=(255, 255, 255), font=font, spacing=6)

        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()