backend/game/tech_tree.py

from __future__ import annotations

from typing import TYPE_CHECKING

from .buildings import BuildingType, BuildingStatus
from .units import UnitType

if TYPE_CHECKING:
    from .state import PlayerState


# unit_type -> required active buildings
UNIT_REQUIREMENTS: dict[UnitType, list[BuildingType]] = {
    UnitType.SCV:     [BuildingType.COMMAND_CENTER],
    UnitType.MARINE:  [BuildingType.BARRACKS],
    UnitType.MEDIC:   [BuildingType.BARRACKS, BuildingType.ENGINEERING_BAY],
    UnitType.GOLIATH: [BuildingType.FACTORY],
    UnitType.TANK:    [BuildingType.FACTORY, BuildingType.ARMORY],
    UnitType.WRAITH:  [BuildingType.STARPORT],
}

# building_type -> required active buildings (to be able to construct it)
BUILDING_REQUIREMENTS: dict[BuildingType, list[BuildingType]] = {
    BuildingType.COMMAND_CENTER:   [],
    BuildingType.SUPPLY_DEPOT:     [],
    BuildingType.REFINERY:         [],
    BuildingType.BARRACKS:         [BuildingType.SUPPLY_DEPOT],
    BuildingType.ENGINEERING_BAY:  [BuildingType.BARRACKS],
    BuildingType.FACTORY:          [BuildingType.BARRACKS],
    BuildingType.ARMORY:           [BuildingType.FACTORY],
    BuildingType.STARPORT:         [BuildingType.FACTORY],
}

# unit_type -> building that produces it
PRODUCTION_SOURCES: dict[UnitType, BuildingType] = {
    UnitType.SCV:     BuildingType.COMMAND_CENTER,
    UnitType.MARINE:  BuildingType.BARRACKS,
    UnitType.MEDIC:   BuildingType.BARRACKS,
    UnitType.GOLIATH: BuildingType.FACTORY,
    UnitType.TANK:    BuildingType.FACTORY,
    UnitType.WRAITH:  BuildingType.STARPORT,
}


def _active(player: "PlayerState", bt: BuildingType) -> bool:
    return any(
        b.building_type == bt
        and b.status not in (BuildingStatus.CONSTRUCTING, BuildingStatus.DESTROYED)
        for b in player.buildings.values()
    )


def can_build(bt: BuildingType, player: "PlayerState") -> bool:
    return all(_active(player, req) for req in BUILDING_REQUIREMENTS[bt])


def can_train(ut: UnitType, player: "PlayerState") -> bool:
    return all(_active(player, req) for req in UNIT_REQUIREMENTS[ut])


def missing_for_build(bt: BuildingType, player: "PlayerState") -> list[BuildingType]:
    return [req for req in BUILDING_REQUIREMENTS[bt] if not _active(player, req)]


def missing_for_train(ut: UnitType, player: "PlayerState") -> list[BuildingType]:
    return [req for req in UNIT_REQUIREMENTS[ut] if not _active(player, req)]


def get_producer(ut: UnitType) -> BuildingType:
    return PRODUCTION_SOURCES[ut]