freeciv_env/models.py

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field, model_validator

from openenv.core.env_server.types import Action, Observation, State


class UnitSummary(BaseModel):
    unit_id: int = Field(..., description="Freeciv unit id")
    unit_type: str = Field(..., description="Ruleset unit type name")
    health: int = Field(0, description="Current health")
    moves_left: int = Field(0, description="Movement points remaining")
    home_city_id: int | None = Field(None, description="Home city id, if any")
    veteran_level: int = Field(0, description="Veteran level")
    can_build_city: bool = Field(False, description="Whether the unit can found a city now")
    move_directions: list[int] = Field(default_factory=list, description="Legal move direction indexes")


class CitySummary(BaseModel):
    city_id: int = Field(..., description="Freeciv city id")
    size: int = Field(..., description="Population size")
    prod_food: int = Field(0, description="Gross food output")
    prod_shield: int = Field(0, description="Gross shield output")
    prod_trade: int = Field(0, description="Gross trade output")
    surplus_food: int = Field(0, description="Net food surplus")
    surplus_shield: int = Field(0, description="Net shield surplus")
    surplus_trade: int = Field(0, description="Net trade surplus")
    production_kind: int | None = Field(None, description="Current production kind enum from Freeciv")
    production_value: int | None = Field(None, description="Current production value id from Freeciv")
    turns_to_complete: float | None = Field(None, description="Turns until current production completes")
    production_options: list[str] = Field(default_factory=list, description="Legal production targets")


class LegalAction(BaseModel):
    action_type: Literal[
        "end_turn",
        "move_unit",
        "build_city",
        "set_city_production",
        "set_research",
    ]
    label: str = Field(..., description="Human-readable action label")
    unit_id: int | None = Field(None, description="Target unit id")
    city_id: int | None = Field(None, description="Target city id")
    direction: int | None = Field(None, description="Freeciv direction index 0..7")
    target: str | None = Field(None, description="Production or tech target name")
    raw_action_key: str | None = Field(None, description="Underlying freeciv-bot action key")


class FreecivAction(Action):
    action_type: Literal[
        "end_turn",
        "move_unit",
        "build_city",
        "set_city_production",
        "set_research",
    ]
    unit_id: int | None = None
    city_id: int | None = None
    direction: int | None = None
    target: str | None = None

    @model_validator(mode="after")
    def validate_shape(self) -> "FreecivAction":
        if self.action_type == "end_turn":
            return self
        if self.action_type == "move_unit":
            if self.unit_id is None or self.direction is None:
                raise ValueError("move_unit requires unit_id and direction")
            return self
        if self.action_type == "build_city":
            if self.unit_id is None:
                raise ValueError("build_city requires unit_id")
            return self
        if self.action_type == "set_city_production":
            if self.city_id is None or not self.target:
                raise ValueError("set_city_production requires city_id and target")
            return self
        if self.action_type == "set_research":
            if not self.target:
                raise ValueError("set_research requires target")
            return self
        raise ValueError(f"unsupported action_type: {self.action_type}")


class FreecivObservation(Observation):
    turn: int = Field(..., description="Current game turn")
    score: float = Field(..., description="Current player score")
    known_tiles: int = Field(..., description="Tiles known to the player")
    visible_tiles: int = Field(..., description="Tiles currently visible to the player")
    city_count: int = Field(..., description="Number of owned cities")
    unit_count: int = Field(..., description="Number of owned units")
    techs_researched: int = Field(..., description="Number of researched techs")
    status: str = Field("ok", description="High-level environment status")
    summary: str = Field(..., description="Compact text summary for LLMs")
    units: list[UnitSummary] = Field(default_factory=list, description="Compact unit summaries")
    cities: list[CitySummary] = Field(default_factory=list, description="Compact city summaries")
    legal_actions: list[LegalAction] = Field(default_factory=list, description="Legal actions exposed by the environment")
    reward: float = Field(0.0, description="Reward from the last action")
    done: bool = Field(False, description="Whether the episode is done")


class FreecivState(State):
    turn: int = Field(0, description="Current game turn")
    score: float = Field(0.0, description="Current player score")
    known_tiles: int = Field(0, description="Known tiles")
    visible_tiles: int = Field(0, description="Visible tiles")
    city_count: int = Field(0, description="Owned city count")
    unit_count: int = Field(0, description="Owned unit count")
    techs_researched: int = Field(0, description="Researched tech count")