"""Vassal agents — War, Economy, and Diplomacy.

Vassals are mid-tier agents responsible for a domain of the kingdom.
Each vassal:
- Listens to the King's subgoal queue
- Computes its domain reward at each tick
- Issues TaskMessages to companion workers
- Reports ResultMessages back up to the King

Model: Qwen3:14b (balanced capability vs. latency).
Frequency: up to 4× per campaign day.

Refs: #1097, #1099.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

from bannerlord.gabs_client import GABSClient, GABSUnavailable
from bannerlord.models import (
    DiplomacyReward,
    EconomyReward,
    KingSubgoal,
    ResultMessage,
    SubgoalMessage,
    TaskMessage,
    WarReward,
)

logger = logging.getLogger(__name__)

# Tokens each vassal responds to (all others are ignored)
_WAR_TOKENS = {"EXPAND_TERRITORY", "RAID_ECONOMY", "TRAIN"}
_ECON_TOKENS = {"FORTIFY", "CONSOLIDATE"}
_DIPLO_TOKENS = {"ALLY"}
_LOGISTICS_TOKENS = {"RECRUIT", "HEAL"}
_TRADE_TOKENS = {"TRADE"}
_SCOUT_TOKENS = {"SPY"}


class BaseVassal:
    """Shared vassal lifecycle — subscribes to subgoal queue, runs tick loop."""

    name: str = "base_vassal"

    def __init__(
        self,
        gabs_client: GABSClient,
        subgoal_queue: asyncio.Queue[SubgoalMessage],
        result_queue: asyncio.Queue[ResultMessage] | None = None,
        task_queue: asyncio.Queue[TaskMessage] | None = None,
    ) -> None:
        self._gabs = gabs_client
        self._subgoal_queue = subgoal_queue
        self._result_queue = result_queue or asyncio.Queue()
        self._task_queue = task_queue or asyncio.Queue()
        self._active_subgoal: KingSubgoal | None = None
        self._running = False

    @property
    def task_queue(self) -> asyncio.Queue[TaskMessage]:
        return self._task_queue

    async def run(self) -> None:
        """Vassal event loop — processes subgoals and emits tasks."""
        self._running = True
        logger.info("%s started", self.name)
        try:
            while self._running:
                # Drain all pending subgoals (keep the latest)
                try:
                    while True:
                        msg = self._subgoal_queue.get_nowait()
                        if msg.to_agent == self.name:
                            self._active_subgoal = msg.subgoal
                            logger.debug(
                                "%s received subgoal %s", self.name, msg.subgoal.token
                            )
                except asyncio.QueueEmpty:
                    pass

                if self._active_subgoal is not None:
                    await self._tick(self._active_subgoal)

                await asyncio.sleep(0.25)  # yield to event loop
        except asyncio.CancelledError:
            logger.info("%s cancelled", self.name)
            raise
        finally:
            self._running = False

    def stop(self) -> None:
        self._running = False

    async def _tick(self, subgoal: KingSubgoal) -> None:
        raise NotImplementedError

    async def _get_state(self) -> dict[str, Any]:
        try:
            return await self._gabs.get_state() or {}
        except GABSUnavailable:
            return {}


# ── War Vassal ────────────────────────────────────────────────────────────────


class WarVassal(BaseVassal):
    """Military operations — sieges, field battles, raids, defensive maneuvers.

    Reward function:
        R = 0.40*ΔTerritoryValue + 0.25*ΔArmyStrengthRatio
            - 0.20*CasualtyCost - 0.10*SupplyCost + 0.05*SubgoalBonus
    """

    name = "war_vassal"

    async def _tick(self, subgoal: KingSubgoal) -> None:
        if subgoal.token not in _WAR_TOKENS | _LOGISTICS_TOKENS:
            return

        state = await self._get_state()
        reward = self._compute_reward(state, subgoal)

        task = self._plan_action(state, subgoal)
        if task:
            await self._task_queue.put(task)

        logger.debug(
            "%s tick: subgoal=%s reward=%.3f action=%s",
            self.name,
            subgoal.token,
            reward.total,
            task.primitive if task else "none",
        )

    def _compute_reward(self, state: dict[str, Any], subgoal: KingSubgoal) -> WarReward:
        bonus = subgoal.priority * 0.05 if subgoal.token in _WAR_TOKENS else 0.0
        return WarReward(
            territory_delta=float(state.get("territory_delta", 0.0)),
            army_strength_ratio=float(state.get("army_strength_ratio", 1.0)),
            casualty_cost=float(state.get("casualty_cost", 0.0)),
            supply_cost=float(state.get("supply_cost", 0.0)),
            subgoal_bonus=bonus,
        )

    def _plan_action(
        self, state: dict[str, Any], subgoal: KingSubgoal
    ) -> TaskMessage | None:
        if subgoal.token == "EXPAND_TERRITORY" and subgoal.target:
            return TaskMessage(
                from_agent=self.name,
                to_agent="logistics_companion",
                primitive="move_party",
                args={"destination": subgoal.target},
                priority=subgoal.priority,
            )
        if subgoal.token == "RECRUIT":
            qty = subgoal.quantity or 20
            return TaskMessage(
                from_agent=self.name,
                to_agent="logistics_companion",
                primitive="recruit_troop",
                args={"troop_type": "infantry", "quantity": qty},
                priority=subgoal.priority,
            )
        if subgoal.token == "TRAIN":
            return TaskMessage(
                from_agent=self.name,
                to_agent="logistics_companion",
                primitive="upgrade_troops",
                args={},
                priority=subgoal.priority,
            )
        return None


# ── Economy Vassal ────────────────────────────────────────────────────────────


class EconomyVassal(BaseVassal):
    """Settlement management, tax collection, construction, food supply.

    Reward function:
        R = 0.35*DailyDenarsIncome + 0.25*FoodStockBuffer + 0.20*LoyaltyAverage
            - 0.15*ConstructionQueueLength + 0.05*SubgoalBonus
    """

    name = "economy_vassal"

    async def _tick(self, subgoal: KingSubgoal) -> None:
        if subgoal.token not in _ECON_TOKENS | _TRADE_TOKENS:
            return

        state = await self._get_state()
        reward = self._compute_reward(state, subgoal)

        task = self._plan_action(state, subgoal)
        if task:
            await self._task_queue.put(task)

        logger.debug(
            "%s tick: subgoal=%s reward=%.3f",
            self.name,
            subgoal.token,
            reward.total,
        )

    def _compute_reward(self, state: dict[str, Any], subgoal: KingSubgoal) -> EconomyReward:
        bonus = subgoal.priority * 0.05 if subgoal.token in _ECON_TOKENS else 0.0
        return EconomyReward(
            daily_denars_income=float(state.get("daily_income", 0.0)),
            food_stock_buffer=float(state.get("food_days_remaining", 0.0)),
            loyalty_average=float(state.get("avg_loyalty", 50.0)),
            construction_queue_length=int(state.get("construction_queue", 0)),
            subgoal_bonus=bonus,
        )

    def _plan_action(
        self, state: dict[str, Any], subgoal: KingSubgoal
    ) -> TaskMessage | None:
        if subgoal.token == "FORTIFY" and subgoal.target:
            return TaskMessage(
                from_agent=self.name,
                to_agent="logistics_companion",
                primitive="build_project",
                args={"settlement": subgoal.target},
                priority=subgoal.priority,
            )
        if subgoal.token == "TRADE":
            return TaskMessage(
                from_agent=self.name,
                to_agent="caravan_companion",
                primitive="assess_prices",
                args={"town": subgoal.target or "nearest"},
                priority=subgoal.priority,
            )
        return None


# ── Diplomacy Vassal ──────────────────────────────────────────────────────────


class DiplomacyVassal(BaseVassal):
    """Relations management — alliances, peace deals, tribute, marriage.

    Reward function:
        R = 0.30*AlliesCount + 0.25*TruceDurationValue + 0.25*RelationsScoreWeighted
            - 0.15*ActiveWarsFront + 0.05*SubgoalBonus
    """

    name = "diplomacy_vassal"

    async def _tick(self, subgoal: KingSubgoal) -> None:
        if subgoal.token not in _DIPLO_TOKENS | _SCOUT_TOKENS:
            return

        state = await self._get_state()
        reward = self._compute_reward(state, subgoal)

        task = self._plan_action(state, subgoal)
        if task:
            await self._task_queue.put(task)

        logger.debug(
            "%s tick: subgoal=%s reward=%.3f",
            self.name,
            subgoal.token,
            reward.total,
        )

    def _compute_reward(self, state: dict[str, Any], subgoal: KingSubgoal) -> DiplomacyReward:
        bonus = subgoal.priority * 0.05 if subgoal.token in _DIPLO_TOKENS else 0.0
        return DiplomacyReward(
            allies_count=int(state.get("allies_count", 0)),
            truce_duration_value=float(state.get("truce_value", 0.0)),
            relations_score_weighted=float(state.get("relations_weighted", 0.0)),
            active_wars_front=int(state.get("active_wars", 0)),
            subgoal_bonus=bonus,
        )

    def _plan_action(
        self, state: dict[str, Any], subgoal: KingSubgoal
    ) -> TaskMessage | None:
        if subgoal.token == "ALLY" and subgoal.target:
            return TaskMessage(
                from_agent=self.name,
                to_agent="scout_companion",
                primitive="track_lord",
                args={"name": subgoal.target},
                priority=subgoal.priority,
            )
        if subgoal.token == "SPY" and subgoal.target:
            return TaskMessage(
                from_agent=self.name,
                to_agent="scout_companion",
                primitive="assess_garrison",
                args={"settlement": subgoal.target},
                priority=subgoal.priority,
            )
        return None
