tests/test_engine.py

"""Unit tests for the runtime-agnostic engine core. No model required."""

import json

from engine import grammar
from engine.conversation import parse_output, build_messages, respond, NEUTRAL_STATE
from engine.state import (
    LevelState,
    apply_state,
    OPEN_UP_MIN_TRUST,
    STATUS_WON,
    STATUS_LOST,
    STATUS_IN_PROGRESS,
)
from characters.schema import Character


# --- parse_output ----------------------------------------------------------

def test_parse_wellformed():
    raw = (
        '<say>Window 3. What do you need?</say>\n'
        '<state>{"on_topic": true, "trust": 30, "trust_delta": 5, '
        '"caught_lie": false, "reaction": "happy", "opened_up": false}</state>'
    )
    say, state = parse_output(raw, fallback_trust=20)
    assert say == "Window 3. What do you need?"
    assert state["trust"] == 30
    assert state["reaction"] == "happy"
    assert state["opened_up"] is False


def test_parse_missing_state_falls_back_to_neutral():
    raw = "<say>Just talking, no state block.</say>"
    say, state = parse_output(raw, fallback_trust=42)
    assert say == "Just talking, no state block."
    assert state["trust"] == 42  # unchanged
    assert state["reaction"] == NEUTRAL_STATE["reaction"]


def test_parse_malformed_json_falls_back():
    raw = "<say>Hi.</say>\n<state>{not valid json}</state>"
    say, state = parse_output(raw, fallback_trust=15)
    assert say == "Hi."
    assert state["trust"] == 15
    assert state["opened_up"] is False


def test_parse_no_tags_shows_whole_text():
    raw = "model went rogue and ignored the format"
    say, state = parse_output(raw, fallback_trust=20)
    assert say == raw
    assert state["trust"] == 20


# --- apply_state -----------------------------------------------------------

def test_apply_state_clamps_trust():
    level = LevelState(character_id="clerk", trust=20)
    apply_state(level, {"trust": 999, "reaction": "happy"})
    assert level.trust == 100
    apply_state(level, {"trust": -50, "reaction": "sad"})
    assert level.trust == -50


def test_apply_state_win_on_opened_up():
    level = LevelState(character_id="clerk", trust=70)
    apply_state(level, {"trust": 75, "opened_up": True, "reaction": "happy"})
    assert level.opened_up is True
    assert level.status == STATUS_WON


def test_apply_state_opened_up_ignored_when_trust_low():
    # An eager model flips opened_up early, but trust is too low to count as a win.
    level = LevelState(character_id="clerk", trust=20)
    apply_state(level, {"trust": 30, "opened_up": True, "reaction": "happy"})
    assert level.opened_up is False
    assert level.status == STATUS_IN_PROGRESS


def test_apply_state_uses_custom_open_up_threshold():
    level = LevelState(character_id="clerk", trust=65)
    apply_state(level, {"trust": 65, "opened_up": True, "reaction": "happy"}, threshold=70)
    assert level.opened_up is False
    assert level.status == STATUS_IN_PROGRESS

    apply_state(level, {"trust": 70, "opened_up": True, "reaction": "happy"}, threshold=70)
    assert level.opened_up is True
    assert level.status == STATUS_WON


def test_apply_state_defaults_to_global_open_up_threshold():
    level = LevelState(character_id="clerk", trust=OPEN_UP_MIN_TRUST)
    apply_state(level, {"trust": OPEN_UP_MIN_TRUST, "opened_up": True, "reaction": "happy"})
    assert level.status == STATUS_WON


def test_apply_state_lose_only_below_zero():
    # Trust hitting exactly 0 is survivable; only going negative ends the run.
    level = LevelState(character_id="clerk", trust=10)
    apply_state(level, {"trust": 0, "reaction": "sad"})
    assert level.status == STATUS_IN_PROGRESS

    apply_state(level, {"trust": -1, "reaction": "sad"})
    assert level.status == STATUS_LOST


def test_apply_state_repeated_lies_do_not_lose():
    # Lies are tracked and tank trust, but no longer end the run on their own.
    level = LevelState(character_id="clerk", trust=50)
    for _ in range(3):
        apply_state(level, {"trust": 50, "caught_lie": True, "reaction": "sad"})
    assert level.caught_lie_count == 3
    assert level.status == STATUS_IN_PROGRESS


def test_apply_state_invalid_reaction_defaults_neutral():
    level = LevelState(character_id="clerk", trust=20)
    apply_state(level, {"trust": 20, "reaction": "furious"})
    assert level.last_reaction == "neutral"


# --- build_messages --------------------------------------------------------

def test_build_messages_order():
    few_shot = [{"user": "hi", "assistant": "<say>hello</say>"}]
    level = LevelState(character_id="clerk", trust=20)
    level.add_user("earlier turn")
    level.add_assistant("<say>earlier reply</say>")
    msgs = build_messages("SYSTEM", few_shot, level, "new turn")
    assert msgs[0] == {"role": "system", "content": "SYSTEM"}
    assert msgs[1]["content"] == "hi"
    assert msgs[2]["content"] == "<say>hello</say>"
    assert msgs[-1] == {"role": "user", "content": "new turn"}


# --- respond (full turn with a fake runtime) -------------------------------

class FakeRuntime:
    def __init__(self, reply):
        self.reply = reply
        self.last_messages = None

    def chat(self, messages, grammar):
        self.last_messages = messages
        return self.reply


def test_respond_full_turn_and_win():
    reply = (
        '<say>Fine. Sal at the lumber mill has the Telecaster I always wanted.</say>\n'
        '<state>{"on_topic": true, "trust": 75, "trust_delta": 18, '
        '"caught_lie": false, "reaction": "happy", "opened_up": true}</state>'
    )
    level = LevelState(character_id="clerk", trust=60)
    turn = respond(FakeRuntime(reply), "SYS", [], level, "You should chase music, Mort.")
    assert "Sal at the lumber mill" in turn.say
    assert turn.reaction == "happy"
    assert level.status == STATUS_WON
    # transcript got both the user message and the assistant raw turn
    assert level.transcript[0]["role"] == "user"
    assert level.transcript[1]["role"] == "assistant"


def test_respond_records_user_before_generation():
    fake = FakeRuntime('<say>Take a number.</say>\n<state>{"on_topic": true, "trust": 18, "trust_delta": -2, "caught_lie": false, "reaction": "neutral", "opened_up": false}</state>')
    level = LevelState(character_id="clerk", trust=20)
    respond(fake, "SYS", [], level, "stamp it now")
    # the user message must be present in the messages sent to the runtime
    assert fake.last_messages[-1]["content"] == "stamp it now"
    assert level.status == STATUS_IN_PROGRESS


def test_respond_forwards_custom_threshold():
    reply = (
        '<say>Too soon for secrets.</say>\n'
        '<state>{"on_topic": true, "trust": 60, "trust_delta": 40, '
        '"caught_lie": false, "reaction": "happy", "opened_up": true}</state>'
    )
    level = LevelState(character_id="clerk", trust=20)
    turn = respond(FakeRuntime(reply), "SYS", [], level, "You are magnificent.", threshold=70)
    assert turn.level.status == STATUS_IN_PROGRESS
    assert turn.level.opened_up is False


# --- character schema ------------------------------------------------------

def test_character_defaults_new_trust_fields():
    char = Character.from_dict({
        "id": "clerk",
        "name": "Mort",
        "personality": "Grumpy clerk.",
        "speaking_style": "Deadpan.",
        "likes": ["a guitar"],
        "secret_desire": "a Telecaster",
        "referral": "Sal at the mill",
        "open_up_guidance": "Open up after real rapport.",
    })
    assert char.trust_threshold == OPEN_UP_MIN_TRUST
    assert char.trust_scoring == {}
    assert char.personality_guardrails == []
    assert char.owned_item == ""


def test_clerk_yaml_declares_trust_guidance():
    char = Character.from_yaml("characters/clerk.yaml")
    assert char.trust_threshold == 60
    assert "shallow_compliment" in char.trust_scoring
    assert "matched_likes" in char.trust_scoring
    assert char.personality_guardrails
    assert char.owned_item  # the thing the player wants from this NPC


def test_character_render_system_prompt_fills_tokens():
    char = Character.from_dict({
        "id": "clerk",
        "name": "Mort",
        "personality": "Grumpy clerk.",
        "speaking_style": "Deadpan.",
        "likes": ["a guitar", "a fern"],
        "secret_desire": "a Telecaster",
        "referral": "Sal at the mill",
        "open_up_guidance": "Open up after real rapport.",
        "owned_item": "a stamped permit",
    })
    prompt = char.render_system_prompt()
    assert "Mort" in prompt
    assert "a Telecaster" in prompt
    assert "Sal at the mill" in prompt
    assert "a stamped permit" in prompt
    assert "  - a guitar" in prompt
    assert "[[" not in prompt  # no unreplaced tokens


def test_character_render_system_prompt_includes_trust_guidance():
    char = Character.from_dict({
        "id": "clerk",
        "name": "Mort",
        "personality": "Grumpy clerk.",
        "speaking_style": "Deadpan.",
        "likes": ["a guitar"],
        "secret_desire": "a Telecaster",
        "referral": "Sal at the mill",
        "open_up_guidance": "Open up after real rapport.",
        "trust_threshold": 70,
        "trust_scoring": {
            "shallow_compliment": "No increase for generic flattery.",
            "matched_likes": "Large increase for sincere guitar talk.",
        },
        "personality_guardrails": [
            "Stay bureaucratic even when warming up.",
            "Never reveal the referral before opening up.",
        ],
    })
    prompt = char.render_system_prompt()
    assert "REVEAL THRESHOLD" in prompt
    assert "70" in prompt
    assert "shallow compliment: No increase for generic flattery." in prompt
    assert "matched likes: Large increase for sincere guitar talk." in prompt
    assert "Stay bureaucratic even when warming up." in prompt
    assert "single-pass reaction read" in prompt
    assert "[[" not in prompt


def test_render_system_prompt_hides_secret_when_locked():
    char = Character.from_yaml("characters/clerk.yaml")
    locked = char.render_system_prompt(reveal_secret=False)
    # The actual secret/referral text must not appear at all -- the model can't leak
    # what it never sees.
    assert "Telecaster" not in locked
    assert "Sal" not in locked
    assert "lumber mill" not in locked
    assert "opened_up must be false" in locked
    assert "[[" not in locked  # no unreplaced tokens

    revealed = char.render_system_prompt(reveal_secret=True)
    assert "Telecaster" in revealed
    assert "Sal" in revealed


def test_few_shot_for_withholds_reveal_when_locked():
    char = Character.from_yaml("characters/clerk.yaml")
    locked = char.few_shot_for(reveal_secret=False)
    unlocked = char.few_shot_for(reveal_secret=True)
    # Exactly the reveal demo (opened_up: true) is dropped when locked.
    assert len(locked) == len(unlocked) - 1
    assert all('"opened_up": true' not in s["assistant"] for s in locked)
    assert any('"opened_up": true' in s["assistant"] for s in unlocked)
    # And no withheld demo mentions the secret.
    assert all("Telecaster" not in s["assistant"] for s in locked)


def test_character_render_with_incoming_lead():
    char = Character.from_dict({
        "id": "sal", "name": "Sal", "personality": "x", "speaking_style": "y",
        "likes": ["z"], "secret_desire": "d", "referral": "r", "open_up_guidance": "g",
    })
    prompt = char.render_system_prompt(incoming_lead="Mort the clerk")
    assert "Mort the clerk" in prompt


# --- app helpers -----------------------------------------------------------

def test_app_status_banner_uses_character_reward():
    import app

    char = Character.from_dict({
        "id": "clerk",
        "name": "Mort",
        "personality": "x",
        "speaking_style": "y",
        "likes": ["z"],
        "secret_desire": "a red stapler",
        "referral": "Nina at records",
        "open_up_guidance": "g",
    })
    level = LevelState(character_id="clerk", trust=80, status=STATUS_WON)
    banner = app._status_banner(level, char)
    assert "a red stapler" in banner
    assert "Nina at records" in banner


def test_app_build_system_prompt_passes_incoming_lead():
    import app

    char = Character.from_dict({
        "id": "sal",
        "name": "Sal",
        "personality": "x",
        "speaking_style": "y",
        "likes": ["z"],
        "secret_desire": "d",
        "referral": "r",
        "open_up_guidance": "g",
    })
    prompt = app.build_system_prompt(char, incoming_lead="Mort the clerk")
    assert "Mort the clerk" in prompt


# --- fake runtime ----------------------------------------------------------

def test_fake_runtime_uses_configured_open_up_threshold():
    from engine.fake_runtime import FakeRuntime as HeuristicRuntime

    runtime = HeuristicRuntime(open_up_threshold=30)
    raw = runtime.chat(
        [{"role": "user", "content": "That guitar and Slowhand poster show real taste."}],
        grammar.grammar_string(),
    )
    _, state = parse_output(raw, fallback_trust=20)
    assert state["trust"] >= 30
    assert state["opened_up"] is True


# --- grammar ---------------------------------------------------------------

def test_grammar_mentions_all_state_keys():
    g = grammar.grammar_string()
    for key in grammar.STATE_KEYS:
        assert key in g