from __future__ import annotations import logging from collections import defaultdict from dataclasses import dataclass from typing import Any, Optional, TYPE_CHECKING, Union from game.ato.flighttype import FlightType from game.theater.controlpoint import ControlPoint if TYPE_CHECKING: from game.theater import ConflictTheater DEFAULT_SQUADRON_SIZE = 12 @dataclass(frozen=True) class SquadronConfig: primary: FlightType secondary: list[FlightType] aircraft: list[str] max_size: int name: Optional[str] nickname: Optional[str] female_pilot_percentage: Optional[int] aircraft_type: Optional[str] @property def auto_assignable(self) -> set[FlightType]: return set(self.secondary) | {self.primary} @classmethod def from_data(cls, data: dict[str, Any]) -> SquadronConfig: secondary_raw = data.get("secondary") if secondary_raw is None: secondary = [] elif isinstance(secondary_raw, str): secondary = cls.expand_secondary_alias(secondary_raw) else: secondary = [FlightType(s) for s in secondary_raw] max_size = data.get("size", DEFAULT_SQUADRON_SIZE) return SquadronConfig( FlightType(data["primary"]), secondary, data.get("aircraft", []), max_size, data.get("name", None), data.get("nickname", None), data.get("female_pilot_percentage", None), data.get("aircraft_type", None), ) @staticmethod def expand_secondary_alias(alias: str) -> list[FlightType]: if alias == "any": return list(FlightType) elif alias == "air-to-air": return [t for t in FlightType if t.is_air_to_air] elif alias == "air-to-ground": return [t for t in FlightType if t.is_air_to_ground] raise KeyError(f"Unknown secondary mission type: {alias}") @dataclass(frozen=True) class CampaignAirWingConfig: by_location: dict[ControlPoint, list[SquadronConfig]] @classmethod def from_campaign_data( cls, data: dict[Union[str, int], Any], theater: ConflictTheater ) -> CampaignAirWingConfig: by_location: dict[ControlPoint, list[SquadronConfig]] = defaultdict(list) carriers = theater.find_carriers() lhas = theater.find_lhas() for base_id, squadron_configs in data.items(): base: Optional[ControlPoint] = None if isinstance(base_id, int): base = theater.find_control_point_by_airport_id(base_id) else: try: base = theater.control_point_named(base_id) except: logging.warning( f"Control point {base_id} not found, trying to match by full name" ) if not base: try: base = theater.control_point_by_full_name(base_id) except KeyError: logging.error(f"Control point {base_id} not found, skipping") for squadron_data in squadron_configs: if base is None: logging.warning( f"Skipping squadron config for unknown base: {base_id}" ) else: by_location[base].append(SquadronConfig.from_data(squadron_data)) return CampaignAirWingConfig(by_location)