backend/agent/state.py

from __future__ import annotations
from dataclasses import dataclass, field


@dataclass
class AbilityInstance:
    name: str
    last_used_turn: int = -1
    uses_remaining: int | None = None

    def can_use(self, current_turn: int) -> bool:
        if self.last_used_turn == current_turn:
            return False
        if self.uses_remaining is not None and self.uses_remaining <= 0:
            return False
        return True

    def use(self, current_turn: int):
        self.last_used_turn = current_turn
        if self.uses_remaining is not None:
            self.uses_remaining -= 1


@dataclass
class Message:
    content: str
    sender: str | None = None


@dataclass
class AgentState:
    id: str
    name: str
    pos: list[int]
    hp: int = 100
    alive: bool = True
    score: int = 0
    kills: int = 0
    shielded: bool = False
    abilities: list[AbilityInstance] = field(default_factory=list)
    messages: list[dict] = field(default_factory=list)
    mailbox: list[Message] = field(default_factory=list)

    def get_ability(self, name: str) -> AbilityInstance | None:
        for ab in self.abilities:
            if ab.name == name:
                return ab
        return None

    def has_ability(self, name: str) -> bool:
        return self.get_ability(name) is not None