tests/fakes.py

from __future__ import annotations

from copy import deepcopy

from freeciv_env.adapter import ActionRef, RawSnapshot


class FakeFreecivSession:
    def __init__(self):
        self.current = None

    def reset(self, seed: int | None = None) -> RawSnapshot:
        del seed
        self.current = _initial_snapshot()
        return deepcopy(self.current)

    def apply_action(self, action_ref: ActionRef) -> RawSnapshot:
        if self.current is None:
            raise RuntimeError("session was not reset")
        raw_action_key = action_ref.raw_action_key
        if raw_action_key == "goto_0":
            self.current = _moved_snapshot(self.current.turn)
        elif raw_action_key == "build":
            self.current = _built_snapshot(self.current.turn)
        elif raw_action_key == "change_unit_prod_Settlers_0":
            self.current = _production_snapshot(self.current.turn)
        elif raw_action_key == "research_tech_Pottery_63":
            self.current = _research_snapshot(self.current.turn)
        else:
            raise ValueError(f"unsupported fake action: {raw_action_key}")
        return deepcopy(self.current)

    def end_turn(self) -> RawSnapshot:
        if self.current is None:
            raise RuntimeError("session was not reset")
        self.current = _advanced_turn_snapshot(self.current)
        return deepcopy(self.current)

    def close(self) -> None:
        self.current = None


def _base_state(*, score: float, techs: int, status: list[list[int]], cities: dict, units: dict) -> dict:
    return {
        "player": {
            "my_score": score,
            "my_gold": 20,
            "my_science": 60,
            "my_techs_researched": techs,
            "my_is_alive": True,
        },
        "map": {"status": status},
        "city": cities,
        "unit": units,
        "tech": {},
    }


def _base_actions(*, can_build: bool, can_move: bool, include_city_prod: bool, include_research: bool) -> dict:
    unit_actions = {}
    if can_move:
        unit_actions["goto_0"] = True
    if can_build:
        unit_actions["build"] = True
    city_actions = {"change_unit_prod_Settlers_0": True} if include_city_prod else {}
    tech_actions = {"research_tech_Pottery_63": True} if include_research else {}
    return {
        "unit": {"201": unit_actions},
        "city": {"101": city_actions},
        "tech": {"cur_player": tech_actions},
    }


def _initial_snapshot(turn: int = 1) -> RawSnapshot:
    return RawSnapshot(
        turn=turn,
        state=_base_state(
            score=10.0,
            techs=0,
            status=[[2, 2, 1], [1, 0, 0]],
            cities={
                "101": {
                    "id": 101,
                    "size": 1,
                    "prod_food": 4,
                    "prod_shield": 2,
                    "prod_trade": 1,
                    "surplus_food": 2,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 0,
                    "turns_to_prod_complete": 2.0,
                }
            },
            units={
                "201": {
                    "health": 10,
                    "moves_left": 1,
                    "home_city": 101,
                    "type_rule_name": "Settlers",
                    "veteran": 0,
                }
            },
        ),
        actions=_base_actions(can_build=True, can_move=True, include_city_prod=True, include_research=True),
    )


def _moved_snapshot(turn: int) -> RawSnapshot:
    return RawSnapshot(
        turn=turn,
        state=_base_state(
            score=11.0,
            techs=0,
            status=[[2, 2, 2], [1, 1, 0]],
            cities={
                "101": {
                    "id": 101,
                    "size": 1,
                    "prod_food": 4,
                    "prod_shield": 2,
                    "prod_trade": 1,
                    "surplus_food": 2,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 0,
                    "turns_to_prod_complete": 2.0,
                }
            },
            units={
                "201": {
                    "health": 10,
                    "moves_left": 0,
                    "home_city": 101,
                    "type_rule_name": "Settlers",
                    "veteran": 0,
                }
            },
        ),
        actions=_base_actions(can_build=False, can_move=False, include_city_prod=True, include_research=True),
    )


def _built_snapshot(turn: int) -> RawSnapshot:
    return RawSnapshot(
        turn=turn,
        state=_base_state(
            score=14.0,
            techs=0,
            status=[[2, 2, 1], [1, 0, 0]],
            cities={
                "101": {
                    "id": 101,
                    "size": 1,
                    "prod_food": 4,
                    "prod_shield": 2,
                    "prod_trade": 1,
                    "surplus_food": 2,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 0,
                    "turns_to_prod_complete": 2.0,
                },
                "102": {
                    "id": 102,
                    "size": 1,
                    "prod_food": 2,
                    "prod_shield": 1,
                    "prod_trade": 1,
                    "surplus_food": 1,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 3,
                    "turns_to_prod_complete": 4.0,
                },
            },
            units={},
        ),
        actions={
            "unit": {},
            "city": {
                "101": {"change_unit_prod_Settlers_0": True},
                "102": {"change_unit_prod_Settlers_0": True},
            },
            "tech": {"cur_player": {"research_tech_Pottery_63": True}},
        },
    )


def _production_snapshot(turn: int) -> RawSnapshot:
    return RawSnapshot(
        turn=turn,
        state=_base_state(
            score=10.0,
            techs=0,
            status=[[2, 2, 1], [1, 0, 0]],
            cities={
                "101": {
                    "id": 101,
                    "size": 1,
                    "prod_food": 4,
                    "prod_shield": 2,
                    "prod_trade": 1,
                    "surplus_food": 2,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 0,
                    "turns_to_prod_complete": 1.0,
                }
            },
            units={
                "201": {
                    "health": 10,
                    "moves_left": 1,
                    "home_city": 101,
                    "type_rule_name": "Settlers",
                    "veteran": 0,
                }
            },
        ),
        actions=_base_actions(can_build=True, can_move=True, include_city_prod=True, include_research=True),
    )


def _research_snapshot(turn: int) -> RawSnapshot:
    return RawSnapshot(
        turn=turn,
        state=_base_state(
            score=10.0,
            techs=1,
            status=[[2, 2, 1], [1, 0, 0]],
            cities={
                "101": {
                    "id": 101,
                    "size": 1,
                    "prod_food": 4,
                    "prod_shield": 2,
                    "prod_trade": 1,
                    "surplus_food": 2,
                    "surplus_shield": 1,
                    "surplus_trade": 1,
                    "production_kind": 6,
                    "production_value": 0,
                    "turns_to_prod_complete": 2.0,
                }
            },
            units={
                "201": {
                    "health": 10,
                    "moves_left": 1,
                    "home_city": 101,
                    "type_rule_name": "Settlers",
                    "veteran": 0,
                }
            },
        ),
        actions={
            "unit": {"201": {"goto_0": True, "build": True}},
            "city": {"101": {"change_unit_prod_Settlers_0": True}},
            "tech": {"cur_player": {}},
        },
    )


def _advanced_turn_snapshot(current: RawSnapshot) -> RawSnapshot:
    state = deepcopy(current.state)
    state["player"]["my_score"] = float(state["player"].get("my_score", 0.0)) + 2.0
    state["map"]["status"] = [[2, 2, 2], [2, 1, 0]]
    if "201" in state.get("unit", {}):
        state["unit"]["201"]["moves_left"] = 1
    actions = _base_actions(can_build=True, can_move=True, include_city_prod=True, include_research=True)
    return RawSnapshot(turn=current.turn + 1, state=state, actions=actions)