tests/test_survival_loop.py

from __future__ import annotations

import json
from typing import Any

from world_simulator.config import (
    ConnectorConfig,
    GameConfig,
    NpcConfig,
    ServerConfig,
    SimulationConfig,
    WorldConfig,
)
from world_simulator.domain import Beast, Npc, ResourceNode, Terrain, Vec3, WorldState
from world_simulator.simulation.connectors.base import NpcDirective, TickPlan
from world_simulator.simulation.connectors.deterministic import DeterministicWorldSimulator
from world_simulator.simulation.connectors.openai_compatible import (
    TOOL_DEFINITIONS,
    OpenAICompatibleWorldSimulator,
    primitive_tools,
    survival_tools_for_role,
)
from world_simulator.simulation.spawning import create_world
import pytest

from world_simulator.simulation.survival import (
    GOAL_SUGGESTED_ACTIONS,
    MAX_BEASTS,
    _natural_beast_spawn_interval,
    apply_action_effects,
    apply_survival_tick,
    select_goal,
    validate_survival_action,
)
from world_simulator.simulation.tick import advance_world, apply_tick_plan


def _npc(npc_id: str = "npc-001", name: str = "Ada", *, x: float = 0.0, z: float = 0.0,
         role: str = "farmer", **state: object) -> Npc:
    return Npc(id=npc_id, name=name, role=role, position=Vec3(x=x, y=0.0, z=z), **state)  # type: ignore[arg-type]


def _world(
    npcs: list[Npc],
    *,
    beasts: list[Beast] | None = None,
    resource_nodes: list[ResourceNode] | None = None,
    tick: int = 0,
) -> WorldState:
    return WorldState(
        tick=tick,
        seed=42,
        terrain=Terrain(kind="plain_green", width=80, depth=80),
        npcs=npcs,
        resource_nodes=resource_nodes or [],
        beasts=beasts or [],
    )


def _survival_config(*, npc_count: int = 6) -> GameConfig:
    return GameConfig(
        world=WorldConfig(width=80, depth=80, terrain="plain_green", seed=42, survival=True),
        npcs=NpcConfig(count=npc_count),
        simulation=SimulationConfig(tick_ms=500),
        server=ServerConfig(host="127.0.0.1", port=8000),
    )


# 1. hunger grows every tick (GAME_DESIGN: +0.8/tick)
def test_hunger_increases_each_tick() -> None:
    npc = _npc(hunger=25.0)
    world = _world([npc], resource_nodes=[ResourceNode("r", "food", Vec3(30.0, 0.0, 30.0), 5)])

    apply_survival_tick(world)
    assert npc.hunger == pytest.approx(25.8)

    apply_survival_tick(world)
    assert npc.hunger == pytest.approx(26.6)


# 2. hunger > 80 + food in inventory -> goal = eat_food
def test_high_hunger_with_food_selects_eat_food() -> None:
    npc = _npc(hunger=85.0, inventory_food=1)
    assert select_goal(npc, _world([npc])) == "eat_food"


# 3. hunger > 80 + no food -> goal = find_food
def test_high_hunger_without_food_selects_find_food() -> None:
    npc = _npc(hunger=85.0, inventory_food=0)
    assert select_goal(npc, _world([npc])) == "find_food"


# 4. health < 40 + herbs -> goal = heal_self
def test_low_health_with_herbs_selects_heal_self() -> None:
    npc = _npc(health=30, inventory_herbs=1)
    assert select_goal(npc, _world([npc])) == "heal_self"


# 5. beast picks the nearest living NPC
def test_beast_targets_nearest_alive_npc() -> None:
    near = _npc("npc-001", "Ada", x=8.0, z=0.0)
    far = _npc("npc-002", "Boris", x=30.0, z=0.0)
    beast = Beast("beast_1", Vec3(0.0, 0.0, 0.0))
    world = _world([far, near], beasts=[beast])

    apply_survival_tick(world)

    assert beast.target_npc_id == "npc-001"


# 6. beast in range attacks and lowers health
def test_beast_in_range_attacks_and_reduces_health() -> None:
    victim = _npc(x=1.5, z=0.0, health=100)
    beast = Beast("beast_1", Vec3(0.0, 0.0, 0.0))  # damage 12, reach half a tile == 2 units
    world = _world([victim], beasts=[beast])

    apply_survival_tick(world)

    assert victim.health == 88
    assert beast.state == "attacking"


# 7. NPC attacked by beast gains fear >= 40
def test_attacked_npc_gains_fear() -> None:
    victim = _npc(x=1.5, z=0.0)
    beast = Beast("beast_1", Vec3(0.0, 0.0, 0.0))
    world = _world([victim], beasts=[beast])

    apply_survival_tick(world)

    assert victim.fear >= 40


# 8. a nearby beast no longer forces a survive/engage goal - the LLM decides.
def test_nearby_beast_does_not_force_threat_goal() -> None:
    npc = _npc(x=8.0, z=0.0)
    beast = Beast("beast_1", Vec3(0.0, 0.0, 0.0))

    goal = select_goal(npc, _world([npc], beasts=[beast]))

    assert not goal.startswith("survive_threat")
    assert goal != "engage_threat"


# 11. use (gather) requires a nearby resource node with amount > 0
def test_gather_requires_nearby_nonempty_resource() -> None:
    npc = _npc(x=0.0, z=0.0)
    node = ResourceNode("r1", "food", Vec3(2.0, 0.0, 0.0), amount=3)
    world = _world([npc], resource_nodes=[node])
    assert validate_survival_action(npc, "use", world) == "gather"

    empty = ResourceNode("r2", "food", Vec3(2.0, 0.0, 0.0), amount=0)
    npc_empty = _npc(x=0.0, z=0.0)
    empty_world = _world([npc_empty], resource_nodes=[empty])
    assert validate_survival_action(npc_empty, "use", empty_world) == "move_to_resource"

    far = ResourceNode("r3", "food", Vec3(40.0, 0.0, 0.0), amount=3)
    npc_far = _npc(x=0.0, z=0.0)
    assert validate_survival_action(npc_far, "use", _world([npc_far], resource_nodes=[far])) == (
        "move_to_resource"
    )

    # The gather effect consumes one unit and fills the right inventory slot.
    apply_action_effects(npc, "use", world)
    assert npc.inventory_food == 1
    assert node.amount == 2


# 12. invalid action from the model does not break the tick (safe fallback)
def test_invalid_action_does_not_break_tick() -> None:
    npc = _npc(x=0.0, z=0.0)
    world = _world([npc], resource_nodes=[ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), 5)])
    plan = TickPlan(
        source="survival",
        directives=[NpcDirective(npc_id=npc.id, action="banana_dance")],  # type: ignore[arg-type]
    )

    apply_tick_plan(world, 1, plan)  # must not raise

    assert world.tick == 1
    assert npc.intention == "idle"
    assert validate_survival_action(npc, "banana_dance", world) == "idle"


# Extra: end-to-end deterministic survival loop runs through the public tick path.
def test_survival_world_advances_and_produces_events() -> None:
    world = create_world(_survival_config(npc_count=6))
    assert world.beasts and world.resource_nodes

    starting_hunger = [npc.hunger for npc in world.npcs]
    for _ in range(30):
        advance_world(world)

    assert world.tick == 30
    assert world.last_tick_source == "survival"
    # Every survival goal is a recognised key for the suggestion table.
    assert all(npc.survival_goal in GOAL_SUGGESTED_ACTIONS for npc in world.npcs)
    # Hunger economy and beast pursuit both did something observable.
    assert any(npc.hunger != start for npc, start in zip(world.npcs, starting_hunger, strict=True))
    assert world.beasts[0].position != Vec3(6.0, 0.0, 6.0) or world.beasts[0].state != "hunting"


# Sustainability: resource nodes regrow toward their cap so the settlement does
# not starve to death once nodes are first depleted.
def test_resource_nodes_regrow_over_time() -> None:
    node = ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), amount=1, max_amount=5)
    world = _world([_npc(x=0.0, z=0.0)], resource_nodes=[node], tick=0)

    apply_survival_tick(world)  # food regrows every 8 ticks; tick 0 % 8 == 0 -> +1
    assert node.amount == 2

    world.tick = 8
    apply_survival_tick(world)  # -> +1
    assert node.amount == 3

    node.amount = 5
    world.tick = 16
    apply_survival_tick(world)  # never exceeds the cap
    assert node.amount == 5


def test_resources_spawn_randomly_but_deterministically_by_seed() -> None:
    node_a = ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), amount=1, max_amount=5)
    node_b = ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), amount=1, max_amount=5)
    first = _world([_npc(x=0.0, z=0.0)], resource_nodes=[node_a], tick=7)
    second = _world([_npc(x=0.0, z=0.0)], resource_nodes=[node_b], tick=7)

    apply_survival_tick(first)
    apply_survival_tick(second)

    assert len(first.resource_nodes) == 2
    assert first.resource_nodes[-1] == second.resource_nodes[-1]
    assert first.resource_nodes[-1].id.startswith("res_")


def test_beast_random_spawn_respects_max_beasts_and_cooldown() -> None:
    def villagers() -> list[Npc]:
        # Natural beast spawns require population >= 4.
        return [_npc(f"npc-{index}", f"Villager{index}", x=float(index)) for index in range(4)]

    def find_spawn_tick() -> int:
        probe = _world(villagers(), tick=0)
        for tick in range(60, 2000):
            probe.tick = tick
            if tick % _natural_beast_spawn_interval(probe) == 0:
                return tick
        raise AssertionError("no natural beast spawn tick found for this seed")

    spawn_tick = find_spawn_tick()
    node = ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), amount=1, max_amount=5)

    cooldown_world = _world(villagers(), resource_nodes=[node], beasts=[], tick=1)
    apply_survival_tick(cooldown_world)
    assert cooldown_world.beasts == []

    spawn_world = _world(villagers(), resource_nodes=[node], beasts=[], tick=spawn_tick)
    apply_survival_tick(spawn_world)
    assert len(spawn_world.beasts) == 1
    assert spawn_world.beasts[0].id == "beast_1"

    capped_world = _world(
        villagers(),
        resource_nodes=[node],
        beasts=[
            Beast(f"beast_{index}", Vec3(float(index), 0.0, 0.0))
            for index in range(1, MAX_BEASTS + 1)
        ],
        tick=spawn_tick,
    )
    apply_survival_tick(capped_world)

    assert len(capped_world.beasts) == MAX_BEASTS


# Interaction: speak records a real memory on BOTH participants.
def test_speak_without_target_is_heard_by_everyone_nearby() -> None:
    ada = _npc("npc-001", "Ada", x=0.0, z=0.0)
    boris = _npc("npc-002", "Boris", x=2.0, z=0.0)
    world = _world([ada, boris], resource_nodes=[ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), 5)])

    summary = apply_action_effects(ada, "speak", world)

    assert summary.startswith("shouting")
    assert any("You said" in m.text for m in ada.memory)
    assert any("Ada said:" in m.text for m in boris.memory)


def test_speak_with_target_exchanges_memory_between_npcs() -> None:
    ada = _npc("npc-001", "Ada", x=0.0, z=0.0)
    boris = _npc("npc-002", "Boris", x=2.0, z=0.0)
    world = _world([ada, boris], resource_nodes=[ResourceNode("r", "food", Vec3(40.0, 0.0, 40.0), 5)])

    directive = NpcDirective(npc_id=ada.id, action="speak", target_npc_id=boris.id, message="hi Boris")
    summary = apply_action_effects(ada, "speak", world, directive=directive)

    assert summary == "talking to Boris"
    assert any("talked" in m.text for m in ada.memory)
    assert any("Ada talked with you" in m.text for m in boris.memory)


# Cooperation: a surplus holder shares food with a needier neighbour.
def test_transfer_shares_food_with_needier_neighbor() -> None:
    giver = _npc("npc-001", "Ada", x=0.0, z=0.0, inventory_food=3)
    taker = _npc("npc-002", "Boris", x=2.0, z=0.0, inventory_food=0)
    world = _world([giver, taker])

    summary = apply_action_effects(giver, "transfer", world)

    assert summary == "sharing food with Boris"
    assert giver.inventory_food == 2
    assert taker.inventory_food == 1
    assert any("Gave food to Boris" in m.text for m in giver.memory)
    assert any("Received food from Ada" in m.text for m in taker.memory)


# Extra: the LLM connector gets a compact survival prompt and its chosen survival
# action resolves through the engine without breaking the tick.
def test_llm_connector_drives_survival_actions() -> None:
    captured: list[str] = []

    def fake_chat_completer(request: dict[str, Any]) -> dict[str, Any]:
        # The survival prompt is a readable text briefing; pick the first tool the
        # connector offered (always valid) instead of parsing a JSON payload.
        captured.append(request["messages"][1]["content"])
        action = request["tools"][0]["function"]["name"]
        return {
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "type": "function",
                                "function": {
                                    "name": action,
                                    "arguments": json.dumps({}),
                                },
                            }
                        ]
                    }
                }
            ]
        }

    world = create_world(_survival_config(npc_count=4))
    simulator = OpenAICompatibleWorldSimulator(
        ConnectorConfig(type="openai_compatible", base_url="http://llm.local/v1", model="test"),
        fallback=DeterministicWorldSimulator(),
        chat_completer=fake_chat_completer,
    )

    advance_world(world, simulator)

    assert world.tick == 1
    assert captured  # the survival briefing was built and sent
    # Every captured prompt is the readable survival briefing, not a JSON dump.
    assert all("YOU are" in briefing for briefing in captured)
    assert all("ENEMY NATION" in briefing for briefing in captured)
    assert all("ACTIONS YOU CAN TAKE:" in briefing for briefing in captured)
    assert all("max_age" not in briefing for briefing in captured)
    assert "openai_compatible" in world.last_tick_source


def test_llm_always_sees_all_public_primitive_tools() -> None:
    names = [tool["function"]["name"] for tool in primitive_tools()]

    assert names == ["move", "speak", "attack", "use", "transfer"]
    assert set(names).issubset(set(TOOL_DEFINITIONS))
    assert "vote" in TOOL_DEFINITIONS
    # Every suggested action references a real tool.
    for suggested in GOAL_SUGGESTED_ACTIONS.values():
        assert set(suggested).issubset(set(TOOL_DEFINITIONS))


def test_survival_tools_are_not_role_gated() -> None:
    guard_names = {tool["function"]["name"] for tool in survival_tools_for_role("guard")}
    builder_names = {tool["function"]["name"] for tool in survival_tools_for_role("builder")}
    gatherer_names = {tool["function"]["name"] for tool in survival_tools_for_role("gatherer")}

    assert guard_names == builder_names == gatherer_names
    for names in (guard_names, builder_names, gatherer_names):
        assert names == {"move", "speak", "attack", "use", "transfer"}