dcs-retribution/game/dcs/unitgroup.py

from __future__ import annotations

import copy
import itertools
import random
from dataclasses import dataclass, field
from pathlib import Path
from typing import ClassVar, TYPE_CHECKING, Any, Iterator

import yaml

from game.data.groups import GroupRole, GroupTask
from game.dcs.groundunittype import GroundUnitType
from game.dcs.shipunittype import ShipUnitType
from game.dcs.unittype import UnitType
from game.point_with_heading import PointWithHeading
from gen.templates import GroundObjectTemplate

if TYPE_CHECKING:
    from game import Game
    from game.factions.faction import Faction
    from game.theater import TheaterGroundObject, ControlPoint


@dataclass
class UnitGroup:
    name: str
    ground_units: list[GroundUnitType]
    ship_units: list[ShipUnitType]
    statics: list[str]
    role: GroupRole
    tasks: list[GroupTask] = field(default_factory=list)
    template_names: list[str] = field(default_factory=list)

    _by_name: ClassVar[dict[str, UnitGroup]] = {}
    _by_role: ClassVar[dict[GroupRole, list[UnitGroup]]] = {}
    _loaded: bool = False
    _templates: list[GroundObjectTemplate] = field(default_factory=list)

    def __str__(self) -> str:
        return self.name

    def update_from_unit_group(self, unit_group: UnitGroup) -> None:
        # Update tasking and templates
        self.tasks.extend([task for task in unit_group.tasks if task not in self.tasks])
        self._templates.extend(
            [
                template
                for template in unit_group.templates
                if template not in self.templates
            ]
        )

    @property
    def templates(self) -> list[GroundObjectTemplate]:
        return self._templates

    def add_template(self, faction_template: GroundObjectTemplate) -> None:
        template = copy.deepcopy(faction_template)
        updated_groups = []
        for group in template.groups:
            unit_types = list(
                itertools.chain(
                    [u.dcs_id for u in self.ground_units if group.can_use_unit(u)],
                    [s.dcs_id for s in self.ship_units if group.can_use_unit(s)],
                    [s for s in self.statics if group.can_use_unit_type(s)],
                )
            )
            if unit_types:
                group.set_possible_types(unit_types)
                updated_groups.append(group)
        template.groups = updated_groups
        self._templates.append(template)

    def load_templates(self, faction: Faction) -> None:
        self._templates = []
        if self.template_names:
            # Preferred templates
            for template_name in self.template_names:
                template = faction.templates.by_name(template_name)
                if template:
                    self.add_template(template)

        if not self._templates:
            # Find all matching templates if no preferred set or available
            for template in list(
                faction.templates.for_role_and_tasks(self.role, self.tasks)
            ):
                if any(self.has_unit_type(unit) for unit in template.units):
                    self.add_template(template)

    def set_templates(self, templates: list[GroundObjectTemplate]) -> None:
        self._templates = templates

    def has_unit_type(self, unit_type: UnitType[Any]) -> bool:
        return unit_type in self.ground_units or unit_type in self.ship_units

    @property
    def unit_types(self) -> Iterator[str]:
        for unit in self.ground_units:
            yield unit.dcs_id
        for ship in self.ship_units:
            yield ship.dcs_id
        for static in self.statics:
            yield static

    @classmethod
    def named(cls, name: str) -> UnitGroup:
        if not cls._loaded:
            cls._load_all()
        return cls._by_name[name]

    def generate(
        self,
        name: str,
        position: PointWithHeading,
        control_point: ControlPoint,
        game: Game,
    ) -> TheaterGroundObject:
        template = random.choice(self.templates)
        return template.generate(name, position, control_point, game)

    @classmethod
    def _load_all(cls) -> None:
        for file in Path("resources/units/unit_groups").glob("*.yaml"):
            if not file.is_file():
                continue

            with file.open(encoding="utf-8") as data_file:
                data = yaml.safe_load(data_file)

            group_role = GroupRole(data.get("role"))

            group_tasks = [GroupTask(n) for n in data.get("tasks", [])]

            ground_units = [
                GroundUnitType.named(n) for n in data.get("ground_units", [])
            ]
            ship_units = [ShipUnitType.named(n) for n in data.get("ship_units", [])]

            unit_group = UnitGroup(
                name=data.get("name"),
                ground_units=ground_units,
                ship_units=ship_units,
                statics=data.get("statics", []),
                role=group_role,
                tasks=group_tasks,
                template_names=data.get("templates", []),
            )

            cls._by_name[unit_group.name] = unit_group
            if group_role in cls._by_role:
                cls._by_role[group_role].append(unit_group)
            else:
                cls._by_role[group_role] = [unit_group]

        cls._loaded = True