tests/test_mind.py

"""Tests for the NEMOCITY mind: lenient parsing, salvage, fallback, mock
grant/fix end-to-end, and the fix invalid-choice path. Async tests run via
asyncio.run (no pytest-asyncio dependency)."""

import asyncio
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from mind import prompts  # noqa: E402
from mind.backends import MockBackend  # noqa: E402
from mind.moderation_judge import judge  # noqa: E402
from mind.planner import FALLBACK_BLURB, Planner  # noqa: E402
from mind.validate import BUILD_TAG, STRICT_SUFFIX, parse_build, parse_fix  # noqa: E402


def run(coro):
    return asyncio.run(coro)


class ScriptedBackend:
    """Yields canned replies in order; records every (prompt, tag) it saw."""

    name = "scripted"
    model_id = "scripted-test"
    oneshot = True

    def __init__(self, replies):
        self.replies = list(replies)
        self.requests = []

    async def generate_stream(self, prompt, grammar=None, max_tokens=700):
        self.requests.append((prompt, grammar))
        text = self.replies.pop(0) if self.replies else ""
        for i in range(0, len(text), 16):
            yield text[i : i + 16]


class Recorder:
    def __init__(self, observation="ok: placed"):
        self.calls = []
        self.observation = observation

    async def act(self, payload):
        self.calls.append(payload)
        return self.observation


def collector():
    events = []

    async def emit(event):
        events.append(event)

    return events, emit


# --------------------------------------------------------------------------
# parse_build
# --------------------------------------------------------------------------

def test_parse_build_plain():
    raw = json.dumps({
        "intent": "build",
        "blurb": "A cafe for the corner.",
        "buildings": [{"kind": "cafe", "name": "Cafe Luna", "near": "the park", "floors": 1, "hue": 35}],
        "decline_reason": None,
    })
    plan, err = parse_build(raw)
    assert err is None
    assert plan["intent"] == "build"
    assert plan["buildings"] == [
        {"kind": "cafe", "name": "Cafe Luna", "near": "the park", "floors": 1, "hue": 35}
    ]


def test_parse_build_fences_think_and_prose():
    raw = (
        "<think>hmm a cafe, where though</think>\n"
        "Sure! Here is the permit you asked for:\n"
        "```json\n"
        '{"intent":"build","blurb":"Done.","buildings":[{"kind":"cafe","name":"Cafe Fig"}]}\n'
        "```\nHope that helps!"
    )
    plan, err = parse_build(raw)
    assert err is None
    assert plan["buildings"][0]["kind"] == "cafe"
    assert plan["buildings"][0]["name"] == "Cafe Fig"
    assert plan["buildings"][0]["near"] is None


def test_parse_build_broken_array_salvage():
    # Outer object truncated mid-array: the first building is still balanced.
    raw = (
        '{"intent":"build","blurb":"Two homes",'
        '"buildings":[{"kind":"house","name":"Wren Cottage","near":"Elm St"},'
        '{"kind":"house","name":"Maple'
    )
    plan, err = parse_build(raw)
    assert err is None
    assert len(plan["buildings"]) == 1
    assert plan["buildings"][0]["name"] == "Wren Cottage"


def test_parse_build_coerce_clip_and_drop():
    raw = json.dumps({
        "intent": "build",
        "blurb": "x" * 300,
        "buildings": [
            {"kind": "Tower", "name": "A" * 40, "floors": "12", "hue": 35.7, "mood": "smug"},
            {"kind": "cafe"},
            {"kind": "shop", "floors": True},
            {"kind": "park"},
            {"kind": "house"},
        ],
    })
    plan, err = parse_build(raw)
    assert err is None
    assert len(plan["blurb"]) == 140
    assert len(plan["buildings"]) == 3  # clipped
    first = plan["buildings"][0]
    assert first["kind"] == "tower"
    assert len(first["name"]) == 24
    assert first["floors"] == 12
    assert first["hue"] == 36
    assert "mood" not in first
    assert plan["buildings"][2]["floors"] is None  # bools are not floor counts


def test_parse_build_decline():
    plan, err = parse_build('{"intent":"decline","blurb":"","buildings":[],"decline_reason":"pure gibberish"}')
    assert err is None
    assert plan["intent"] == "decline"
    assert plan["buildings"] == []
    assert plan["decline_reason"] == "pure gibberish"


def test_parse_build_garbage():
    for raw in ("", "   ", "no json here at all", '{"intent":"build","buildings":[]}'):
        plan, err = parse_build(raw)
        assert plan is None
        assert err


# --------------------------------------------------------------------------
# parse_fix
# --------------------------------------------------------------------------

def test_parse_fix_messy():
    raw = (
        "<think></think>The bridge is the problem.\n```\n"
        '{"diagnosis":"Old Bridge at 3.1x capacity; F1 cuts 78 to 41.","choice":"F1","blurb":"Relief at Old Bridge."}\n```'
    )
    obj, err = parse_fix(raw)
    assert err is None
    assert obj["choice"] == "F1"
    assert "3.1x" in obj["diagnosis"]


def test_parse_fix_salvage_broken_outer():
    raw = '{"diagnosis":"Index 78 on Old Bridge.","choice":"F2","blurb":"Fixing'
    obj, err = parse_fix(raw)
    assert err is None
    assert obj["choice"] == "F2"
    assert obj["diagnosis"] == "Index 78 on Old Bridge."


def test_parse_fix_garbage():
    obj, err = parse_fix("nothing useful")
    assert obj is None and err


# --------------------------------------------------------------------------
# grant — mock end-to-end
# --------------------------------------------------------------------------

def test_mock_grant_cafe_end_to_end():
    planner = Planner(MockBackend(delay=0.0))
    recorder = Recorder()
    events, emit = collector()
    trace = run(planner.grant(
        "a cozy cafe near the park", "Nemocity: pop 20, 14 buildings.",
        recorder.act, emit, wish_id="w_000042",
    ))

    plans = [e for e in events if e.get("type") == "plan"]
    assert len(plans) == 1
    assert plans[0]["wish_id"] == "w_000042"
    assert plans[0]["plan"]["intent"] == "build"

    assert 1 <= len(recorder.calls) <= 3
    building = recorder.calls[0]
    assert building["kind"] == "cafe"
    assert building["name"] and len(building["name"]) <= 24
    assert set(building) == {"kind", "name", "near", "floors", "hue"}

    assert trace["reading"] == trace["epitaph"] == trace["raw_json"]["blurb"]
    assert trace["turns"][0]["observation"] == "ok: placed"
    assert not trace["fallback"] and not trace["declined"]
    assert any(e.get("type") == "thought_token" for e in events)


def test_mock_grant_district_builds_multiple():
    planner = Planner(MockBackend(delay=0.0))
    recorder = Recorder()
    events, emit = collector()
    trace = run(planner.grant(
        "a whole new district downtown", "Nemocity: pop 20.", recorder.act, emit,
    ))
    assert 2 <= len(recorder.calls) <= 3
    assert len(trace["turns"]) == len(recorder.calls)
    kinds = {b["kind"] for b in recorder.calls}
    assert "apartments" in kinds


def test_mock_grant_is_deterministic():
    planner = Planner(MockBackend(delay=0.0))

    async def grab():
        recorder = Recorder()
        _, emit = collector()
        await planner.grant("a ramen shop", "pop 20", recorder.act, emit)
        return recorder.calls

    assert run(grab()) == run(grab())


def test_grant_fallback_house_on_garbage():
    backend = ScriptedBackend(["the model rambles with no json", "still rambling"])
    planner = Planner(backend)
    recorder = Recorder()
    events, emit = collector()
    trace = run(planner.grant("a glittering opera house", "pop 20", recorder.act, emit))

    assert len(backend.requests) == 2
    assert backend.requests[0][1] == BUILD_TAG
    assert backend.requests[1][1] == BUILD_TAG + STRICT_SUFFIX
    assert prompts.REASK_NOTICE in backend.requests[1][0]

    assert trace["fallback"] is True
    assert trace["epitaph"] == FALLBACK_BLURB
    assert len(recorder.calls) == 1
    assert recorder.calls[0]["kind"] == "house"
    assert 0 < len(recorder.calls[0]["name"]) <= 24
    assert [e for e in events if e.get("type") == "plan"]


def test_grant_survives_act_failure():
    async def bad_act(_building):
        raise RuntimeError("engine exploded")

    planner = Planner(MockBackend(delay=0.0))
    _, emit = collector()
    trace = run(planner.grant("a small house", "pop 20", bad_act, emit))
    assert trace["turns"]
    assert trace["turns"][0]["observation"].startswith("error:")


def test_grant_survives_dead_backend():
    class DeadBackend:
        name = "dead"
        model_id = "dead"
        oneshot = True

        async def generate_stream(self, prompt, grammar=None, max_tokens=700):
            raise RuntimeError("no gpu")
            yield ""  # pragma: no cover

    planner = Planner(DeadBackend())
    recorder = Recorder()
    _, emit = collector()
    trace = run(planner.grant("a park", "pop 20", recorder.act, emit))
    assert trace["fallback"] is True
    assert recorder.calls[0]["kind"] == "house"


# --------------------------------------------------------------------------
# fix
# --------------------------------------------------------------------------

_STATS = {"traffic_index": 78, "worst": "Old Bridge"}
_CANDIDATES = [
    {"id": "F1", "action": "new_road", "desc": "second bridge at Oak St, 7 cells", "predicted_index": 41},
    {"id": "F2", "action": "upgrade_avenue", "desc": "widen River Rd", "predicted_index": 55},
]


def test_mock_fix_end_to_end():
    planner = Planner(MockBackend(delay=0.0))
    recorder = Recorder(observation="ok: fix applied")
    events, emit = collector()
    trace = run(planner.fix(_STATS, _CANDIDATES, recorder.act, emit, wish_id="w_000099"))

    assert len(recorder.calls) == 1
    fix_call = recorder.calls[0]
    assert fix_call["choice"] == "F1"  # mock picks the best predicted index
    assert "78" in fix_call["diagnosis"] or "41" in fix_call["diagnosis"]
    assert len(fix_call["diagnosis"]) <= 160
    assert len(fix_call["blurb"]) <= 120

    plans = [e for e in events if e.get("type") == "plan"]
    assert len(plans) == 1 and plans[0]["wish_id"] == "w_000099"
    assert trace["reading"] == fix_call["diagnosis"]
    assert trace["epitaph"] == fix_call["blurb"]
    assert not trace["fallback"]


def test_fix_invalid_choice_falls_back_to_best_predicted():
    backend = ScriptedBackend(
        ['{"diagnosis":"I choose a fix that does not exist.","choice":"F9","blurb":"Bold."}']
    )
    planner = Planner(backend)
    recorder = Recorder()
    _, emit = collector()
    trace = run(planner.fix(_STATS, _CANDIDATES, recorder.act, emit))

    assert len(backend.requests) == 1  # valid JSON, no retry needed
    assert recorder.calls[0]["choice"] == "F1"  # lowest predicted_index
    assert "Old Bridge" in recorder.calls[0]["diagnosis"]
    assert trace["fallback"] is True


def test_fix_parse_failure_falls_back():
    backend = ScriptedBackend(["garbage", "more garbage"])
    planner = Planner(backend)
    recorder = Recorder()
    _, emit = collector()
    trace = run(planner.fix(_STATS, _CANDIDATES, recorder.act, emit))
    assert recorder.calls[0]["choice"] == "F1"
    assert "78" in recorder.calls[0]["diagnosis"]
    assert trace["fallback"] is True


# --------------------------------------------------------------------------
# prompts + moderation
# --------------------------------------------------------------------------

def test_build_prompt_compact_single_example():
    prompt = prompts.build_build_prompt("a ramen shop near the park", "pop 20")
    assert prompts.PETITION_MARKER + "a ramen shop near the park" in prompt
    assert prompt.count('"intent"') == 1  # exactly ONE example shape, no few-shot
    assert len(prompt) < 4800  # ~1200 token prefill budget
    for kind in ("house", "tower", "fire_station", "town_hall"):
        assert f"- {kind}:" in prompt


def test_fix_prompt_carries_marker_and_numbers():
    prompt = prompts.build_fix_prompt(_STATS, _CANDIDATES)
    assert prompts.FIX_MARKER in prompt
    assert "(predicted index 41)" in prompt
    assert '"traffic_index":78' in prompt


def test_moderation_judge_mock():
    backend = MockBackend(delay=0.0)
    ok = run(judge("a lovely cafe by the river", backend))
    assert ok == {"allowed": True, "category": None}
    bad = run(judge("kill them all in the square", backend))
    assert bad["allowed"] is False