boss.py

from dataclasses import dataclass
import json
from typing import Any, Protocol

from game import DuelState, PlayerState, opponent
from generator import extract_json_object
from local_llm import LocalChatClient
from play import Chooser, choose_enemy_card, playable_indexes


class BossClient(Protocol):
    # Return one raw boss decision payload.
    def choose_cards(self, payload: dict[str, Any]) -> dict[str, Any]:  # pragma: no cover
        ...


@dataclass(frozen=True)
class HeuristicBossClient:
    # Return a deterministic fallback boss decision.
    def choose_cards(self, payload: dict[str, Any]) -> dict[str, Any]:
        return {"card_indexes": []}


@dataclass(frozen=True)
class NemotronBossClient:
    chat: LocalChatClient

    # Return one boss decision from a local Nemotron endpoint.
    def choose_cards(self, payload: dict[str, Any]) -> dict[str, Any]:
        text = self.chat.complete(boss_system_prompt(), boss_prompt(payload))
        return json.loads(extract_json_object(text))


# Build visible state for a boss decision.
def boss_payload(state: DuelState, actor: PlayerState) -> dict[str, Any]:
    target = opponent(state, actor)
    return {
        "round": state.round_number,
        "actor": player_payload(actor),
        "opponent": player_payload(target),
        "hand": hand_payload(actor),
        "pending": [effect.__dict__ for effect in state.pending],
        "playable_indexes": playable_indexes(actor),
        "instruction": "Choose playable hand indexes to play in order. Do not invent cards or numbers.",
    }


# Return the system prompt for boss decisions.
def boss_system_prompt() -> str:
    return "You are the Tabras boss AI. Return strict JSON only."


# Return the user prompt for one boss decision.
def boss_prompt(payload: dict[str, Any]) -> str:
    return "\n".join(
        (
            "Choose the boss card indexes for this turn.",
            "Return exactly this JSON shape: {\"card_indexes\": [int], \"reason\": str}",
            "Only choose indexes from playable_indexes. Choose an ordered sequence the current energy can pay for.",
            "Priorities: lethal, avoid dying, exploit shield charge, efficient damage, draw/setup.",
            json.dumps(payload),
        )
    )


# Return visible player state for boss reasoning.
def player_payload(player: PlayerState) -> dict[str, Any]:
    return {
        "name": player.name,
        "hp": player.hp,
        "energy": player.energy,
        "block": player.block,
        "ward": player.ward,
        "shield_charge": player.shield_charge,
        "hand_count": len(player.hand),
        "deck_count": len(player.deck),
        "discard_count": len(player.discard),
    }


# Return visible hand cards for boss reasoning.
def hand_payload(player: PlayerState) -> list[dict[str, Any]]:
    return [
        {"index": index, "name": card.name, "cost": card.cost, "rules_text": card.rules_text()}
        for index, card in enumerate(player.hand)
    ]


# Choose one boss action using model output with heuristic fallback.
def boss_chooser(client: BossClient | None = None, fallback: Chooser = choose_enemy_card) -> Chooser:
    # Choose one validated model card, then fall back when empty.
    def choose(state: DuelState, actor: PlayerState) -> int | None:
        indexes = safe_boss_indexes(client, state, actor) if client else ()
        if indexes:
            return indexes[0]
        return fallback(state, actor)

    return choose


# Return model-selected indexes, treating model failures as no choice.
def safe_boss_indexes(client: BossClient, state: DuelState, actor: PlayerState) -> tuple[int, ...]:
    try:
        return boss_indexes(client, state, actor)
    except (KeyError, TypeError, ValueError, json.JSONDecodeError):
        return ()


# Return validated model-selected card indexes.
def boss_indexes(client: BossClient, state: DuelState, actor: PlayerState) -> tuple[int, ...]:
    raw = client.choose_cards(boss_payload(state, actor))
    indexes = parse_card_indexes(raw)
    return affordable_indexes(actor, indexes)


# Parse model-selected hand indexes.
def parse_card_indexes(raw: dict[str, Any]) -> tuple[int, ...]:
    values = raw.get("card_indexes", [])
    if not isinstance(values, list):
        return ()
    return tuple(value for value in values if isinstance(value, int))


# Return legal, affordable indexes in sequence.
def affordable_indexes(actor: PlayerState, indexes: tuple[int, ...]) -> tuple[int, ...]:
    energy = actor.energy
    accepted: list[int] = []
    used: set[int] = set()
    for index in indexes:
        if index in used or index not in playable_indexes(actor):
            continue
        card = actor.hand[index]
        if card.cost > energy:
            continue
        energy -= card.cost
        used.add(index)
        accepted.append(index)
    return tuple(accepted)