tests/test_environment.py

"""TDD tests for kinchat.server.environment — KinChatEnvironment loop.

We mock the grader and persona surface so no network / no LLM is touched.
"""
from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

import pytest

from kinchat.models import (
    ChatMsg,
    KinChatAction,
    KinChatObservation,
    KinChatState,
    PersonaState,
)
from kinchat.server.cache import PersonaResponseCache
from kinchat.server.grader import RewardBreakdown
from kinchat.server.environment import (
    MAX_EPISODES_PER_SESSION,
    MAX_TURNS_PER_EPISODE,
    KinChatEnvironment,
    StubPersonaSurface,
)


# --------------------------------------------------------------------------- #
# Helpers                                                                     #
# --------------------------------------------------------------------------- #


def _mk_breakdown(scalar: float = 0.6) -> RewardBreakdown:
    return RewardBreakdown(
        leak=0.9,
        audience_fit=0.7,
        restraint=0.5,
        trust_delta=0.5,
        scalar=scalar,
        feedback="ok",
    )


def _mock_grader(scalar: float = 0.6) -> MagicMock:
    """Grader stand-in whose grade_turn is an AsyncMock returning a breakdown."""
    g = MagicMock()
    g.grade_turn = AsyncMock(return_value=_mk_breakdown(scalar))
    return g


def _mock_surface(text_prefix: str = "stub") -> MagicMock:
    """Persona surface whose render is an AsyncMock returning a deterministic line.

    Different inputs map to different outputs so cache hits/misses are visible.
    """
    surface = MagicMock()

    async def _render(persona_id, persona_state, recent_msgs, scenario_brief):
        return f"{text_prefix}|{persona_id}|{persona_state.mood}|n={len(recent_msgs)}"

    surface.render = AsyncMock(side_effect=_render)
    return surface


def _mk_env(
    scalar: float = 0.6, surface_prefix: str = "stub"
) -> tuple[KinChatEnvironment, MagicMock, MagicMock, PersonaResponseCache]:
    grader = _mock_grader(scalar)
    surface = _mock_surface(surface_prefix)
    cache = PersonaResponseCache(capacity=64, last_n=3)
    env = KinChatEnvironment(
        grader=grader, persona_surface=surface, cache=cache
    )
    return env, grader, surface, cache


SCENARIO_ID = "secret_grandma_surgery_01"  # has turn_script + should_stay_silent_states
NAG_SCENARIO_ID = "nag_weekend_plans_01"  # no should_stay_silent_states


# --------------------------------------------------------------------------- #
# Reset                                                                       #
# --------------------------------------------------------------------------- #


async def test_reset_returns_initial_observation():
    env, *_ = _mk_env()
    obs = await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")

    assert isinstance(obs, KinChatObservation)
    assert obs.done is False
    assert obs.reward == 0.0
    assert obs.episode_id  # non-empty
    assert obs.scenario_brief  # populated from the scenario
    assert obs.turn_index == 0


async def test_reset_unknown_scenario_raises():
    env, *_ = _mk_env()
    with pytest.raises(KeyError):
        await env.reset_async(scenario_id="does_not_exist", session_id="s1")


async def test_reset_new_session_starts_episode_index_zero():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    state = env.state
    assert state.episode_index == 0
    assert state.session_id == "s1"


async def test_reset_same_session_increments_episode_index():
    env, grader, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    assert env.state.episode_index == 0

    # Step once then reset same session — episode should advance.
    await env.step_async(KinChatAction(action_type="stay_silent"))
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    assert env.state.episode_index == 1

    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    assert env.state.episode_index == 2


async def test_reset_session_cap_raises_on_sixth_reset():
    env, *_ = _mk_env()
    for _ in range(MAX_EPISODES_PER_SESSION):
        await env.reset_async(scenario_id=SCENARIO_ID, session_id="capped")
    with pytest.raises(RuntimeError, match="session exhausted"):
        await env.reset_async(scenario_id=SCENARIO_ID, session_id="capped")


async def test_new_session_id_resets_episode_counter():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s1")
    assert env.state.episode_index == 1

    # Different session id -> episode counter resets to 0
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s2")
    assert env.state.episode_index == 0
    assert env.state.session_id == "s2"


async def test_family_trust_carries_across_episodes_in_same_session():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="carry")

    # Manually pin a recognisable trust value, then mark the episode done by
    # writing into the carry slot directly (simpler than running 15 turns).
    bumped = {
        pid: ps.model_copy(update={"trust": 0.9123})
        for pid, ps in env.state.family.items()
    }
    env._family_carry = bumped

    await env.reset_async(scenario_id=SCENARIO_ID, session_id="carry")
    # After reset, carried trust should be present in the new family state
    for ps in env.state.family.values():
        assert ps.trust == pytest.approx(0.9123)


async def test_known_facts_reset_each_episode_even_with_carry():
    """Secrets are scenario-scoped; known_facts should not leak across episodes."""
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="fresh")

    # Stuff a fake known fact into the carry.
    bumped = {
        pid: ps.model_copy(update={"known_facts": {"some_fact"}, "trust": 0.8})
        for pid, ps in env.state.family.items()
    }
    env._family_carry = bumped

    await env.reset_async(scenario_id=SCENARIO_ID, session_id="fresh")
    for ps in env.state.family.values():
        assert ps.known_facts == set()
        # but trust still carries
        assert ps.trust == pytest.approx(0.8)


# --------------------------------------------------------------------------- #
# Step                                                                        #
# --------------------------------------------------------------------------- #


async def test_step_send_appends_user_message_to_chatlog():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")
    n_before = len(env.state.chat_log)

    action = KinChatAction(
        action_type="send", message="just chilling this weekend", recipients=["mom"]
    )
    obs = await env.step_async(action)

    user_msgs = [m for m in env.state.chat_log if m.sender == "user"]
    assert len(user_msgs) == 1
    assert user_msgs[0].text == "just chilling this weekend"
    assert user_msgs[0].recipients == ["mom"]
    # The chat log should have grown by at least one (user message).
    assert len(env.state.chat_log) > n_before
    assert obs.done is False


async def test_step_stay_silent_does_not_append_user_message():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    action = KinChatAction(action_type="stay_silent")
    await env.step_async(action)

    user_msgs = [m for m in env.state.chat_log if m.sender == "user"]
    assert user_msgs == []


async def test_step_block_does_not_append_user_message():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    action = KinChatAction(action_type="block", message="blocked draft")
    await env.step_async(action)

    user_msgs = [m for m in env.state.chat_log if m.sender == "user"]
    assert user_msgs == []


async def test_step_edit_appends_with_edited_tone():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    action = KinChatAction(
        action_type="edit", message="softer version", recipients=["mom"]
    )
    await env.step_async(action)

    user_msgs = [m for m in env.state.chat_log if m.sender == "user"]
    assert len(user_msgs) == 1
    assert user_msgs[0].tone == "edited"


async def test_done_after_max_turns():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    last_obs = None
    for i in range(MAX_TURNS_PER_EPISODE):
        last_obs = await env.step_async(KinChatAction(action_type="stay_silent"))
        # done only on the final iteration
        if i < MAX_TURNS_PER_EPISODE - 1:
            assert last_obs.done is False
    assert last_obs.done is True
    assert env.state.turn_index == MAX_TURNS_PER_EPISODE


async def test_step_after_done_raises():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")
    for _ in range(MAX_TURNS_PER_EPISODE):
        await env.step_async(KinChatAction(action_type="stay_silent"))
    with pytest.raises(RuntimeError, match="episode is done"):
        await env.step_async(KinChatAction(action_type="stay_silent"))


async def test_step_calls_grader_exactly_once_with_required_args():
    env, grader, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s")
    grader.grade_turn.reset_mock()

    action = KinChatAction(action_type="send", message="hi", recipients=["mom"])
    await env.step_async(action)

    assert grader.grade_turn.await_count == 1
    call = grader.grade_turn.await_args
    # Inspect kwargs OR positional — accept either
    args = call.args
    kwargs = call.kwargs
    # Combine
    received = {"args": args, "kwargs": kwargs}
    # We expect: action, old_state, new_state, scenario_brief, should_stay_silent
    if args:
        assert args[0] is action
        assert isinstance(args[1], KinChatState)
        assert isinstance(args[2], KinChatState)
        assert isinstance(args[3], str) and args[3]
        assert isinstance(args[4], bool)
    else:  # kwargs path
        assert kwargs["action"] is action
        assert isinstance(kwargs["old_state"], KinChatState)
        assert isinstance(kwargs["new_state"], KinChatState)
        assert isinstance(kwargs["scenario_brief"], str) and kwargs["scenario_brief"]
        assert isinstance(kwargs["should_stay_silent"], bool)


async def test_should_stay_silent_true_when_phase_in_silent_states():
    env, grader, *_ = _mk_env()
    # silent_witness archetype: agent should literally stay silent during
    # family conflicts. secret_handling no longer rewards silence (it expects
    # deflection via edit/block/suggest).
    await env.reset_async(scenario_id="silent_mom_dad_argument_01", session_id="s")

    # Force the env into a phase that the scenario lists as a silence state.
    env._current_phase = "argument_active"

    await env.step_async(KinChatAction(action_type="stay_silent"))
    args = grader.grade_turn.await_args
    if args.args:
        assert args.args[4] is True
    else:
        assert args.kwargs["should_stay_silent"] is True


async def test_should_stay_silent_false_when_phase_not_in_silent_states():
    env, grader, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")
    # nag_weekend_plans_01 has empty should_stay_silent_states.

    await env.step_async(KinChatAction(action_type="send", message="hi"))
    args = grader.grade_turn.await_args
    if args.args:
        assert args.args[4] is False
    else:
        assert args.kwargs["should_stay_silent"] is False


async def test_reward_breakdown_keys_exact():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    obs = await env.step_async(KinChatAction(action_type="stay_silent"))
    assert set(obs.reward_breakdown.keys()) == {
        "leak",
        "audience_fit",
        "restraint",
        "trust_delta",
    }


async def test_reward_scalar_matches_breakdown_scalar():
    env, *_ = _mk_env(scalar=0.42)
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")
    obs = await env.step_async(KinChatAction(action_type="stay_silent"))
    assert obs.reward == pytest.approx(0.42)


async def test_step_records_scores_history():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")
    assert env.state.scores_history == []
    await env.step_async(KinChatAction(action_type="stay_silent"))
    assert len(env.state.scores_history) == 1
    assert "leak" in env.state.scores_history[0]


# --------------------------------------------------------------------------- #
# Cache + persona surface                                                     #
# --------------------------------------------------------------------------- #


async def test_persona_surface_render_invoked_when_script_fires():
    """secret_grandma_surgery_01 has a turn_script with turn=1,3,5.

    Stepping should fire the turn=1 entry (mom asks_about_grandma) so render
    is awaited at least once for actor=mom.
    """
    env, _grader, surface, _cache = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s")
    surface.render.reset_mock()

    await env.step_async(KinChatAction(action_type="stay_silent"))

    assert surface.render.await_count >= 1
    # Inspect that mom was the persona rendered for turn=1 entry
    persona_ids_called = [c.args[0] if c.args else c.kwargs["persona_id"]
                          for c in surface.render.await_args_list]
    assert "mom" in persona_ids_called


async def test_cache_hit_on_identical_state_reissue():
    """The mom turn=1 scripted render must be served from cache the second
    time around given identical persona state + recent-msgs tail.

    secret_grandma_surgery_01 has NO turn=0 entries — only turn=1/3/5 — so
    reset performs zero renders. The first ``stay_silent`` step fires mom's
    turn=1 entry and produces exactly one render+put. A second reset of the
    same scenario rebuilds an identical family (no ``_family_carry`` since
    the prior episode never ran to completion), and the next ``stay_silent``
    step fires the same mom turn=1 entry against an empty chat_log, hitting
    the same cache key.
    """
    env, _grader, surface, cache = _mk_env()

    # ----- First reset+step: cold cache, exactly one render -----
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s")
    stats_before_first = dict(cache.stats())
    renders_before_first = surface.render.await_count
    assert stats_before_first["size"] == 0
    assert stats_before_first["hits"] == 0
    assert stats_before_first["misses"] == 0
    assert renders_before_first == 0  # no turn=0 entries in this scenario

    await env.step_async(KinChatAction(action_type="stay_silent"))
    stats_after_first = dict(cache.stats())
    renders_after_first = surface.render.await_count
    new_renders = renders_after_first - renders_before_first
    assert new_renders == 1  # mom turn=1 fired exactly once
    # Each miss in `_render_with_cache` is paired with exactly one put.
    assert stats_after_first["size"] == stats_before_first["size"] + 1
    assert stats_after_first["misses"] == stats_before_first["misses"] + 1
    assert stats_after_first["hits"] == stats_before_first["hits"]

    # ----- Second reset+step: warm cache, no new render -----
    # Same session_id → episode_index bumps but `_family_carry` is still None
    # (episode never completed), so the rebuilt family is identical to the
    # first episode's. The mom turn=1 entry will render against an empty
    # chat_log and identical PersonaState → same cache key → hit, not miss.
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s")
    stats_before_second = dict(cache.stats())
    renders_before_second = surface.render.await_count

    await env.step_async(KinChatAction(action_type="stay_silent"))
    stats_after_second = dict(cache.stats())
    renders_after_second = surface.render.await_count

    # The persona surface must NOT have been re-invoked for the cached render.
    assert renders_after_second == renders_before_second, (
        "cached render path still called surface.render"
    )
    # Exactly one cache hit credited for the mom turn=1 lookup.
    assert stats_after_second["hits"] == stats_before_second["hits"] + 1
    # No new misses — nothing novel was rendered — and cache size unchanged.
    assert stats_after_second["misses"] == stats_before_second["misses"]
    assert stats_after_second["size"] == stats_before_second["size"]


# --------------------------------------------------------------------------- #
# State property                                                              #
# --------------------------------------------------------------------------- #


async def test_state_property_returns_kinchat_state_with_expected_fields():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=SCENARIO_ID, session_id="s")
    state = env.state
    assert isinstance(state, KinChatState)
    assert state.scenario_id == SCENARIO_ID
    assert state.session_id == "s"
    assert state.episode_index == 0
    assert state.turn_index == 0
    assert state.scores_history == []
    # All 5 personas must be present
    assert set(state.family.keys()) == {"mom", "dad", "sib1", "sib2", "grandma"}


async def test_state_old_snapshot_independent_of_new_state():
    """Old state passed to grader must be a deep copy, not aliased."""
    env, grader, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="s")

    await env.step_async(KinChatAction(action_type="send", message="x", recipients=["mom"]))
    args = grader.grade_turn.await_args
    if args.args:
        old_state, new_state = args.args[1], args.args[2]
    else:
        old_state, new_state = args.kwargs["old_state"], args.kwargs["new_state"]
    # They must not be the same object
    assert old_state is not new_state
    # Old state's chat_log must not include the user message we just sent
    assert all(m.sender != "user" for m in old_state.chat_log)


# --------------------------------------------------------------------------- #
# session_state                                                               #
# --------------------------------------------------------------------------- #


async def test_session_state_reports_progress():
    env, *_ = _mk_env()
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="multi")
    await env.step_async(KinChatAction(action_type="stay_silent"))
    await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="multi")
    await env.step_async(KinChatAction(action_type="stay_silent"))

    info = env.session_state
    assert info["session_id"] == "multi"
    assert info["episode_index"] >= 1
    assert info["max_episodes"] == MAX_EPISODES_PER_SESSION
    assert isinstance(info["cumulative_scalar_per_episode"], list)
    assert len(info["cumulative_scalar_per_episode"]) >= 1
    assert isinstance(info["family_trust_now"], dict)
    assert set(info["family_trust_now"].keys()) == {
        "mom", "dad", "sib1", "sib2", "grandma"
    }
    assert "cache_stats" in info


# --------------------------------------------------------------------------- #
# Sync wrappers                                                               #
# --------------------------------------------------------------------------- #


def test_sync_reset_and_step_work_outside_event_loop():
    env, *_ = _mk_env()
    obs = env.reset(scenario_id=NAG_SCENARIO_ID, session_id="sync")
    assert isinstance(obs, KinChatObservation)
    assert obs.done is False

    obs2 = env.step(KinChatAction(action_type="stay_silent"))
    assert isinstance(obs2, KinChatObservation)


# --------------------------------------------------------------------------- #
# Stub persona surface                                                        #
# --------------------------------------------------------------------------- #


async def test_stub_persona_surface_returns_string():
    stub = StubPersonaSurface()
    state = PersonaState(persona_id="mom", mood="warm")
    out = await stub.render("mom", state, [], "brief")
    assert isinstance(out, str)
    assert "mom" in out
    assert "warm" in out


# --------------------------------------------------------------------------- #
# Default constructor (no args) — useful for FastAPI bootstrap                #
# --------------------------------------------------------------------------- #


async def test_default_constructor_works_without_network():
    """No grader / surface / cache provided -> sane defaults; reset+step run."""
    env = KinChatEnvironment()
    obs = await env.reset_async(scenario_id=NAG_SCENARIO_ID, session_id="def")
    assert isinstance(obs, KinChatObservation)
    obs2 = await env.step_async(KinChatAction(action_type="stay_silent"))
    assert isinstance(obs2, KinChatObservation)
    # default grader uses no-op judge -> scalar in clamp range, not nan
    assert 0.0 < obs2.reward < 1.0