Timmy-time-dashboard/tests/unit/test_bannerlord/test_agents.py

"""Unit tests for bannerlord agents — King, Vassals, Companions."""

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from bannerlord.agents.companions import (
    CaravanCompanion,
    LogisticsCompanion,
    ScoutCompanion,
)
from bannerlord.agents.king import KingAgent
from bannerlord.agents.vassals import DiplomacyVassal, EconomyVassal, WarVassal
from bannerlord.gabs_client import GABSClient, GABSUnavailable
from bannerlord.ledger import Ledger
from bannerlord.models import (
    KingSubgoal,
    ResultMessage,
    SubgoalMessage,
    TaskMessage,
    VictoryCondition,
)


# ── Helpers ───────────────────────────────────────────────────────────────────


def _mock_gabs(state: dict | None = None) -> GABSClient:
    """Return a disconnected GABS stub that returns *state* from get_state."""
    gabs = MagicMock(spec=GABSClient)
    gabs.connected = False
    if state is not None:
        gabs.get_state = AsyncMock(return_value=state)
    else:
        gabs.get_state = AsyncMock(side_effect=GABSUnavailable("no game"))
    gabs.call = AsyncMock(return_value={})
    gabs.recruit_troops = AsyncMock(return_value={"recruited": 10})
    gabs.move_party = AsyncMock(return_value={"moving": True})
    return gabs


def _mock_ledger(tmp_path) -> Ledger:
    ledger = Ledger(db_path=tmp_path / "ledger.db")
    ledger.initialize()
    return ledger


# ── King agent ────────────────────────────────────────────────────────────────


class TestKingAgent:
    async def test_victory_detected(self, tmp_path):
        """Campaign stops immediately when victory condition is met."""
        gabs = _mock_gabs({"player_title": "King", "territory_control_pct": 55.0})
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0)
        victory = await king.run_campaign(max_ticks=10)
        assert victory.achieved

    async def test_max_ticks_respected(self, tmp_path):
        """Campaign stops after max_ticks when victory not yet achieved."""
        gabs = _mock_gabs({"player_title": "Lord", "territory_control_pct": 10.0})
        ledger = _mock_ledger(tmp_path)

        # Patch LLM to return a valid subgoal without calling Ollama
        king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0)
        with patch.object(king, "_decide", AsyncMock(return_value=KingSubgoal(token="RECRUIT"))):
            victory = await king.run_campaign(max_ticks=3)

        assert not victory.achieved
        assert king._tick == 3

    async def test_llm_failure_falls_back_to_recruit(self, tmp_path):
        """If LLM fails, King defaults to RECRUIT subgoal."""
        gabs = _mock_gabs({"player_title": "Lord", "territory_control_pct": 5.0})
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0)

        with patch.object(king, "_llm_decide", side_effect=RuntimeError("Ollama down")):
            subgoal = await king._decide({})

        assert subgoal.token == "RECRUIT"

    async def test_subgoal_broadcast_to_all_vassals(self, tmp_path):
        """King broadcasts subgoal to all three vassals."""
        gabs = _mock_gabs({})
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger)
        subgoal = KingSubgoal(token="EXPAND_TERRITORY", target="Epicrotea")
        await king._broadcast_subgoal(subgoal)

        messages = []
        while not king.subgoal_queue.empty():
            messages.append(king.subgoal_queue.get_nowait())

        assert len(messages) == 3
        recipients = {m.to_agent for m in messages}
        assert recipients == {"war_vassal", "economy_vassal", "diplomacy_vassal"}

    async def test_gabs_unavailable_uses_empty_state(self, tmp_path):
        """King handles GABS being offline gracefully."""
        gabs = _mock_gabs()  # raises GABSUnavailable
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger)
        state = await king._fetch_state()
        assert state == {}

    def test_evaluate_victory_king_with_majority(self, tmp_path):
        gabs = _mock_gabs()
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger)
        v = king._evaluate_victory({"player_title": "King", "territory_control_pct": 60.0})
        assert v.achieved

    def test_evaluate_victory_not_king(self, tmp_path):
        gabs = _mock_gabs()
        ledger = _mock_ledger(tmp_path)
        king = KingAgent(gabs_client=gabs, ledger=ledger)
        v = king._evaluate_victory({"player_title": "Lord", "territory_control_pct": 80.0})
        assert not v.achieved


# ── Vassals ───────────────────────────────────────────────────────────────────


class TestWarVassal:
    async def test_expand_territory_emits_move_task(self):
        gabs = _mock_gabs({"territory_delta": 1.0, "army_strength_ratio": 1.5})
        queue = asyncio.Queue()
        vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="EXPAND_TERRITORY", target="Seonon")
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "move_party"
        assert task.args["destination"] == "Seonon"

    async def test_recruit_emits_recruit_task(self):
        gabs = _mock_gabs({})
        queue = asyncio.Queue()
        vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="RECRUIT", quantity=15)
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "recruit_troop"
        assert task.args["quantity"] == 15

    async def test_irrelevant_token_emits_no_task(self):
        gabs = _mock_gabs({})
        queue = asyncio.Queue()
        vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="ALLY")
        await vassal._tick(subgoal)
        assert vassal.task_queue.empty()


class TestEconomyVassal:
    async def test_fortify_emits_build_task(self):
        gabs = _mock_gabs({"daily_income": 200.0})
        queue = asyncio.Queue()
        vassal = EconomyVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="FORTIFY", target="Epicrotea")
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "build_project"
        assert task.args["settlement"] == "Epicrotea"

    async def test_trade_emits_assess_prices(self):
        gabs = _mock_gabs({})
        queue = asyncio.Queue()
        vassal = EconomyVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="TRADE", target="Pravend")
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "assess_prices"


class TestDiplomacyVassal:
    async def test_ally_emits_track_lord(self):
        gabs = _mock_gabs({"allies_count": 1})
        queue = asyncio.Queue()
        vassal = DiplomacyVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="ALLY", target="Derthert")
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "track_lord"
        assert task.args["name"] == "Derthert"

    async def test_spy_emits_assess_garrison(self):
        gabs = _mock_gabs({})
        queue = asyncio.Queue()
        vassal = DiplomacyVassal(gabs_client=gabs, subgoal_queue=queue)
        subgoal = KingSubgoal(token="SPY", target="Marunath")
        await vassal._tick(subgoal)
        task: TaskMessage = vassal.task_queue.get_nowait()
        assert task.primitive == "assess_garrison"
        assert task.args["settlement"] == "Marunath"


# ── Companions ────────────────────────────────────────────────────────────────


class TestLogisticsCompanion:
    async def test_recruit_troop(self):
        gabs = _mock_gabs()
        gabs.recruit_troops = AsyncMock(return_value={"recruited": 10, "type": "infantry"})
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = LogisticsCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="war_vassal",
            to_agent="logistics_companion",
            primitive="recruit_troop",
            args={"troop_type": "infantry", "quantity": 10},
        )
        result = await comp._execute(task)
        assert result.success is True
        assert result.outcome["recruited"] == 10

    async def test_unknown_primitive_fails_gracefully(self):
        gabs = _mock_gabs()
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = LogisticsCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="war_vassal",
            to_agent="logistics_companion",
            primitive="launch_nukes",
            args={},
        )
        result = await comp._execute(task)
        assert result.success is False
        assert "Unknown primitive" in result.outcome["error"]

    async def test_gabs_unavailable_returns_failure(self):
        gabs = _mock_gabs()
        gabs.recruit_troops = AsyncMock(side_effect=GABSUnavailable("offline"))
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = LogisticsCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="war_vassal",
            to_agent="logistics_companion",
            primitive="recruit_troop",
            args={"troop_type": "infantry", "quantity": 5},
        )
        result = await comp._execute(task)
        assert result.success is False


class TestCaravanCompanion:
    async def test_assess_prices(self):
        gabs = _mock_gabs()
        gabs.call = AsyncMock(return_value={"grain": 12, "linen": 45})
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = CaravanCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="economy_vassal",
            to_agent="caravan_companion",
            primitive="assess_prices",
            args={"town": "Pravend"},
        )
        result = await comp._execute(task)
        assert result.success is True

    async def test_abandon_route(self):
        gabs = _mock_gabs()
        gabs.call = AsyncMock(return_value={"abandoned": True})
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = CaravanCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="economy_vassal",
            to_agent="caravan_companion",
            primitive="abandon_route",
            args={},
        )
        result = await comp._execute(task)
        assert result.success is True
        assert result.outcome["abandoned"] is True


class TestScoutCompanion:
    async def test_assess_garrison(self):
        gabs = _mock_gabs()
        gabs.call = AsyncMock(return_value={"garrison_size": 120, "settlement": "Marunath"})
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = ScoutCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="diplomacy_vassal",
            to_agent="scout_companion",
            primitive="assess_garrison",
            args={"settlement": "Marunath"},
        )
        result = await comp._execute(task)
        assert result.success is True
        assert result.outcome["garrison_size"] == 120

    async def test_report_intel(self):
        gabs = _mock_gabs()
        gabs.call = AsyncMock(return_value={"intel": ["Derthert at Epicrotea"]})
        q: asyncio.Queue[TaskMessage] = asyncio.Queue()
        comp = ScoutCompanion(gabs_client=gabs, task_queue=q)
        task = TaskMessage(
            from_agent="diplomacy_vassal",
            to_agent="scout_companion",
            primitive="report_intel",
            args={},
        )
        result = await comp._execute(task)
        assert result.success is True