state_handler.py

"""Dispatch game states to prompt builders and execute the chosen action."""

from __future__ import annotations

import asyncio
import logging
from typing import Any

from game_client import GameClient, GameActionError
from decision_engine import DecisionEngine
from memory.knowledge_base import KnowledgeBase
from memory.run_logger import RunLogger
from prompts.combat import build_combat_prompt, build_hand_select_prompt
from prompts.map import build_map_prompt
from prompts.rewards import build_rewards_prompt, build_card_reward_prompt
from prompts.event import (
    build_event_prompt,
    build_rest_site_prompt,
    build_shop_prompt,
    build_card_select_prompt,
    build_relic_select_prompt,
    build_treasure_prompt,
)

logger = logging.getLogger(__name__)

COMBAT_STATES = {"monster", "elite", "boss"}

_card_select_count = 0
_skipped_card_reward = False


class CombatLog:
    """Tracks actions and state changes within a single combat encounter."""

    def __init__(self) -> None:
        self._active = False
        self._last_round = 0
        self._round_actions: list[dict[str, Any]] = []
        self._history: list[dict[str, Any]] = []
        self._prev_player_hp = 0
        self._prev_enemies: dict[str, int] = {}
        self._ended_turn = False

    def reset(self) -> None:
        self._active = False
        self._last_round = 0
        self._round_actions.clear()
        self._history.clear()
        self._prev_player_hp = 0
        self._prev_enemies.clear()
        self._ended_turn = False

    def update(self, state: dict[str, Any]) -> None:
        """Call before each combat decision to track round transitions."""
        battle = state.get("battle", {})
        round_num = battle.get("round", 0)
        player = battle.get("player", {})
        enemies = battle.get("enemies", [])

        if not self._active:
            self._active = True
            self._last_round = round_num
            self._prev_player_hp = player.get("hp", 0)
            self._prev_enemies = {
                e.get("entity_id", "?"): e.get("hp", 0) for e in enemies
            }
            return

        if round_num != self._last_round:
            cur_hp = player.get("hp", 0)
            enemy_changes: list[str] = []
            for e in enemies:
                eid = e.get("entity_id", "?")
                ename = e.get("name", "?")
                ehp = e.get("hp", 0)
                prev_hp = self._prev_enemies.get(eid, ehp)
                if prev_hp != ehp:
                    enemy_changes.append(f"{ename}: {prev_hp}→{ehp}")

            self._history.append({
                "round": self._last_round,
                "actions": list(self._round_actions),
                "hp_change": f"{self._prev_player_hp}→{cur_hp}",
                "enemy_changes": enemy_changes,
            })

            self._round_actions.clear()
            self._last_round = round_num
            self._prev_player_hp = cur_hp
            self._prev_enemies = {
                e.get("entity_id", "?"): e.get("hp", 0) for e in enemies
            }
            self._ended_turn = False

    def record_action(self, action: dict[str, Any], hand: list[dict]) -> None:
        """Record an action taken during the current round."""
        act_name = action.get("action", "")
        if act_name == "play_card":
            idx = action.get("card_index", -1)
            card_name = "?"
            for c in hand:
                if c.get("index") == idx:
                    card_name = c.get("name", "?")
                    break
            self._round_actions.append({"type": "play", "card": card_name})
        elif act_name == "use_potion":
            self._round_actions.append({"type": "potion", "slot": action.get("slot")})
        elif act_name == "end_turn":
            self._round_actions.append({"type": "end_turn"})

    def format_history(self) -> str:
        """Format combat history for prompt injection."""
        if not self._history:
            return ""

        lines: list[str] = []
        for entry in self._history:
            rnd = entry["round"]
            acts = entry["actions"]
            cards = [a["card"] for a in acts if a["type"] == "play"]
            potions = [str(a.get("slot", "?")) for a in acts if a["type"] == "potion"]
            parts: list[str] = []
            if cards:
                parts.append(f"出牌: {', '.join(cards)}")
            if potions:
                parts.append(f"用药水: 槽位{', '.join(potions)}")
            action_str = " | ".join(parts) if parts else "无操作"
            hp_str = entry["hp_change"]
            enemy_str = ", ".join(entry["enemy_changes"]) if entry["enemy_changes"] else "无变化"
            lines.append(
                f"  回合{rnd}: {action_str} → 我方HP: {hp_str}, 敌方: {enemy_str}"
            )

        return "## 本场战斗历史\n\n" + "\n".join(lines)


_combat_log = CombatLog()


async def handle_state(
    state: dict[str, Any],
    md_state: str,
    client: GameClient,
    engine: DecisionEngine,
    kb: KnowledgeBase,
    run_logger: RunLogger,
) -> bool:
    """Process one game state cycle.

    Returns True if the run is still active, False if it ended (state_type == 'menu').
    """
    global _card_select_count, _skipped_card_reward, _combat_log
    state_type: str = state.get("state_type", "menu")

    if state_type == "menu":
        _card_select_count = 0
        _skipped_card_reward = False
        _combat_log.reset()
        return False

    if state_type not in COMBAT_STATES and _combat_log._active:
        _combat_log.reset()

    if state_type != "card_select":
        _card_select_count = 0

    # After skipping a card reward we return to combat_rewards where the
    # card entry is still listed.  Auto-proceed to avoid a claim loop.
    if state_type == "combat_rewards" and _skipped_card_reward:
        _skipped_card_reward = False
        inner = state.get("rewards", {})
        can_proceed = inner.get("can_proceed", False) if isinstance(inner, dict) else False
        if can_proceed:
            logger.info("Card reward was skipped — auto-proceeding from combat_rewards")
            run_logger.log_decision(state_type, {"action": "proceed"}, "已跳过卡牌奖励，自动前进", state)
            await client.proceed()
            return True

    # During enemy turns the API returns hand=[] and is_play_phase=false.
    # Sending end_turn here would queue it and instantly skip the NEXT player turn.
    if state_type in COMBAT_STATES:
        battle = state.get("battle", {})
        turn = battle.get("turn")
        is_play_phase = battle.get("is_play_phase", True)

        if turn != "player":
            # Genuine enemy turn — just wait.
            _combat_log._ended_turn = False
            logger.debug("Enemy turn — waiting")
            return True

        if not is_play_phase:
            if _combat_log._ended_turn:
                # We already sent end_turn this round; the game is
                # transitioning to the enemy turn.  Just wait.
                logger.debug("Waiting for enemy turn to start (end_turn already sent)")
                return True
            # Player turn but card-play phase ended (energy exhausted / no playable
            # cards).  The game is waiting for an explicit end_turn; send it now.
            try:
                logger.info("Play phase ended — auto end_turn")
                run_logger.log_decision(state_type, {"action": "end_turn"}, "出牌阶段结束，自动结束回合", state)
                await client.end_turn()
                _combat_log._ended_turn = True
            except GameActionError:
                # end_turn was rejected — enemy turn already in progress.
                _combat_log._ended_turn = True
                logger.debug("End turn rejected (enemy turn in progress) — waiting")
            return True

        # New player turn with play phase active — reset the flag.
        _combat_log._ended_turn = False
        _combat_log.update(state)

    # Event states: the API may initially return options=[] while the game
    # is still loading (e.g. Neow's reward screen).  Wait until options or
    # dialogue are available before asking the LLM to decide.
    if state_type == "event":
        inner = state.get("event", {})
        if isinstance(inner, dict):
            options = inner.get("options", [])
            in_dialogue = inner.get("in_dialogue", False)
            if not in_dialogue and (not isinstance(options, list) or not options):
                logger.debug("Event options not loaded yet — waiting")
                return True

    # Rest site / combat_rewards / treasure: if no actionable options remain
    # but can_proceed is True, auto-proceed without calling the LLM.
    if state_type in ("rest_site", "combat_rewards", "treasure"):
        inner_key = {
            "rest_site": "rest_site",
            "combat_rewards": "rewards",
            "treasure": "treasure",
        }[state_type]
        inner = state.get(inner_key, {})
        if isinstance(inner, dict):
            options = inner.get("options", inner.get("items", inner.get("relics", [])))
            can_proceed = inner.get("can_proceed", False)
            if (not options or not isinstance(options, list)) and can_proceed:
                logger.info("No options remaining in %s — auto-proceeding", state_type)
                run_logger.log_decision(state_type, {"action": "proceed"}, "无剩余选项，自动前进", state)
                await client.proceed()
                return True

    prompt, allowed, experience_category = _build_prompt_for_state(
        state_type, state, md_state
    )

    if prompt is None:
        logger.warning("No prompt built for state_type=%s, attempting proceed", state_type)
        try:
            await client.proceed()
        except Exception:
            pass
        return True

    experience = kb.retrieve(experience_category, top_k=3)

    action = await engine.decide(
        user_prompt=prompt,
        experience_context=experience,
        allowed_actions=allowed,
    )

    reasoning = action.pop("reasoning", "")
    logger.info("Decision: %s | %s", action, reasoning)

    run_logger.log_decision(state_type, action, reasoning, state)

    await _execute_action(action, client)

    if state_type in COMBAT_STATES:
        battle = state.get("battle", {})
        hand = battle.get("player", {}).get("hand", [])
        _combat_log.record_action(action, hand)
        if action.get("action") == "end_turn":
            _combat_log._ended_turn = True

    if action.get("action") == "skip_card_reward":
        _skipped_card_reward = True
    elif state_type != "card_reward":
        _skipped_card_reward = False

    # card_select two-step flow: after select_card, auto-confirm
    if state_type == "card_select" and action.get("action") == "select_card":
        _card_select_count += 1
        await asyncio.sleep(0.5)
        try:
            await client.confirm_selection()
            logger.info("Auto-confirmed card selection")
        except GameActionError:
            logger.debug("Auto-confirm not ready yet (may need more selections)")

    # hand_select two-step flow: after combat_select_card, auto-confirm
    if state_type == "hand_select" and action.get("action") == "combat_select_card":
        await asyncio.sleep(0.5)
        try:
            await client.combat_confirm_selection()
            logger.info("Auto-confirmed combat card selection")
        except GameActionError:
            logger.debug("Combat auto-confirm not ready yet (may need more selections)")

    return True


# ── prompt routing ──────────────────────────────────────────────────

def _build_prompt_for_state(
    state_type: str,
    state: dict[str, Any],
    md_state: str,
) -> tuple[str | None, list[str] | None, str]:
    """Return (prompt, allowed_actions, experience_category)."""

    if state_type in COMBAT_STATES:
        return (
            build_combat_prompt(state, md_state, _combat_log.format_history()),
            ["play_card", "use_potion", "end_turn"],
            "combat",
        )

    if state_type == "hand_select":
        return (
            build_hand_select_prompt(state, md_state),
            ["combat_select_card", "combat_confirm_selection"],
            "combat",
        )

    if state_type == "combat_rewards":
        return (
            build_rewards_prompt(state, md_state),
            ["claim_reward", "proceed"],
            "rewards",
        )

    if state_type == "card_reward":
        return (
            build_card_reward_prompt(state, md_state),
            ["select_card_reward", "skip_card_reward"],
            "deck_building",
        )

    if state_type == "map":
        return (
            build_map_prompt(state, md_state),
            ["choose_map_node"],
            "pathing",
        )

    if state_type == "rest_site":
        return (
            build_rest_site_prompt(state, md_state),
            ["choose_rest_option", "proceed"],
            "rest_site",
        )

    if state_type == "shop":
        return (
            build_shop_prompt(state, md_state),
            ["shop_purchase", "proceed"],
            "shop",
        )

    if state_type == "event":
        return (
            build_event_prompt(state, md_state),
            ["choose_event_option", "advance_dialogue"],
            "event",
        )

    if state_type == "card_select":
        return (
            build_card_select_prompt(state, md_state),
            ["select_card", "confirm_selection", "cancel_selection"],
            "deck_building",
        )

    if state_type == "relic_select":
        return (
            build_relic_select_prompt(state, md_state),
            ["select_relic", "skip_relic_selection"],
            "relic",
        )

    if state_type == "treasure":
        return (
            build_treasure_prompt(state, md_state),
            ["claim_treasure_relic", "proceed"],
            "treasure",
        )

    if state_type == "overlay":
        return None, None, ""

    logger.warning("Unhandled state_type: %s", state_type)
    return None, None, ""


# ── action execution ───────────────────────────────────────────────

async def _execute_action(action: dict[str, Any], client: GameClient) -> None:
    """Map an action dict to the corresponding GameClient method call."""
    name = action["action"]

    dispatch: dict[str, Any] = {
        "play_card":                lambda: client.play_card(action["card_index"], action.get("target")),
        "use_potion":               lambda: client.use_potion(action["slot"], action.get("target")),
        "end_turn":                 lambda: client.end_turn(),
        "combat_select_card":       lambda: client.combat_select_card(action["card_index"]),
        "combat_confirm_selection": lambda: client.combat_confirm_selection(),
        "claim_reward":             lambda: client.claim_reward(action["index"]),
        "select_card_reward":       lambda: client.select_card_reward(action["card_index"]),
        "skip_card_reward":         lambda: client.skip_card_reward(),
        "proceed":                  lambda: client.proceed(),
        "choose_rest_option":       lambda: client.choose_rest_option(action["index"]),
        "shop_purchase":            lambda: client.shop_purchase(action["index"]),
        "choose_event_option":      lambda: client.choose_event_option(action["index"]),
        "advance_dialogue":         lambda: client.advance_dialogue(),
        "choose_map_node":          lambda: client.choose_map_node(action["index"]),
        "select_card":              lambda: client.select_card(action["index"]),
        "confirm_selection":        lambda: client.confirm_selection(),
        "cancel_selection":         lambda: client.cancel_selection(),
        "select_relic":             lambda: client.select_relic(action["index"]),
        "skip_relic_selection":     lambda: client.skip_relic_selection(),
        "claim_treasure_relic":     lambda: client.claim_treasure_relic(action["index"]),
    }

    handler = dispatch.get(name)
    if handler is None:
        logger.error("No executor for action: %s", name)
        return

    await handler()