dcs-retribution/game/factions/faction.py

from __future__ import annotations

import copy
import itertools
import logging
import random
from dataclasses import dataclass, field
from typing import Optional, Dict, Type, List, Any, Iterator, TYPE_CHECKING

import dcs
from dcs.countries import country_dict
from dcs.unittype import ShipType

from game.data.building_data import (
    WW2_ALLIES_BUILDINGS,
    DEFAULT_AVAILABLE_BUILDINGS,
    WW2_GERMANY_BUILDINGS,
    WW2_FREE,
)
from game.data.doctrine import (
    Doctrine,
    MODERN_DOCTRINE,
    COLDWAR_DOCTRINE,
    WWII_DOCTRINE,
)
from game.data.units import UnitClass
from game.data.groups import GroupRole, GroupTask
from game import db
from game.dcs.aircrafttype import AircraftType
from game.dcs.groundunittype import GroundUnitType
from game.dcs.shipunittype import ShipUnitType
from game.dcs.unitgroup import UnitGroup
from game.dcs.unittype import UnitType
from gen.templates import (
    GroundObjectTemplates,
    GroundObjectTemplate,
    GroupTemplate,
)

if TYPE_CHECKING:
    from game.theater.start_generator import ModSettings


@dataclass
class Faction:
    #: List of locales to use when generating random names. If not set, Faker will
    #: choose the default locale.
    locales: Optional[List[str]]

    # Country used by this faction
    country: str = field(default="")

    # Nice name of the faction
    name: str = field(default="")

    # List of faction file authors
    authors: str = field(default="")

    # A description of the faction
    description: str = field(default="")

    # Available aircraft
    aircrafts: List[AircraftType] = field(default_factory=list)

    # Available awacs aircraft
    awacs: List[AircraftType] = field(default_factory=list)

    # Available tanker aircraft
    tankers: List[AircraftType] = field(default_factory=list)

    # Available frontline units
    frontline_units: List[GroundUnitType] = field(default_factory=list)

    # Available artillery units
    artillery_units: List[GroundUnitType] = field(default_factory=list)

    # Infantry units used
    infantry_units: List[GroundUnitType] = field(default_factory=list)

    # Logistics units used
    logistics_units: List[GroundUnitType] = field(default_factory=list)

    # Possible Air Defence units, Like EWRs
    air_defense_units: List[GroundUnitType] = field(default_factory=list)

    # A list of all supported sets of units
    preset_groups: list[UnitGroup] = field(default_factory=list)

    # Possible Missile site generators for this faction
    missiles: List[GroundUnitType] = field(default_factory=list)

    # Required mods or asset packs
    requirements: Dict[str, str] = field(default_factory=dict)

    # Possible carrier names
    carrier_names: List[str] = field(default_factory=list)

    # Possible helicopter carrier names
    helicopter_carrier_names: List[str] = field(default_factory=list)

    # Available Naval Units
    naval_units: List[ShipUnitType] = field(default_factory=list)

    # Whether this faction has JTAC access
    has_jtac: bool = field(default=False)

    # Unit to use as JTAC for this faction
    jtac_unit: Optional[AircraftType] = field(default=None)

    # doctrine
    doctrine: Doctrine = field(default=MODERN_DOCTRINE)

    # List of available building templates for this faction
    building_set: List[str] = field(default_factory=list)

    # List of default livery overrides
    liveries_overrides: Dict[AircraftType, List[str]] = field(default_factory=dict)

    #: Set to True if the faction should force the "Unrestricted satnav" option
    #: for the mission. This option enables GPS for capable aircraft regardless
    #: of the time period or operator. For example, the CJTF "countries" don't
    #: appear to have GPS capability, so they need this.
    #:
    #: Note that this option cannot be set per-side. If either faction needs it,
    #: both will use it.
    unrestricted_satnav: bool = False

    # All possible templates which can be generated by the faction
    templates: GroundObjectTemplates = field(default=GroundObjectTemplates())

    # All available unit_groups
    unit_groups: dict[GroupRole, list[UnitGroup]] = field(default_factory=dict)

    # Save all accessible units for performance increase
    _accessible_units: list[UnitType[Any]] = field(default_factory=list)

    def __getitem__(self, item: str) -> Any:
        return getattr(self, item)

    @property
    def accessible_units(self) -> Iterator[UnitType[Any]]:
        yield from self._accessible_units

    @property
    def air_defenses(self) -> list[str]:
        """Returns the Air Defense types"""
        air_defenses = [a.name for a in self.air_defense_units]
        air_defenses.extend(
            [pg.name for pg in self.preset_groups if pg.role == GroupRole.AntiAir]
        )
        return sorted(air_defenses)

    def has_access_to_unit_type(self, unit_type: str) -> bool:
        # GroundUnits
        if any(unit_type == u.dcs_id for u in self.accessible_units):
            return True

        # Statics
        if db.static_type_from_name(unit_type) is not None:
            # TODO Improve the statics checking
            return True
        return False

    def has_access_to_unit_class(self, unit_class: UnitClass) -> bool:
        return any(unit.unit_class is unit_class for unit in self.accessible_units)

    def _load_accessible_units(self, templates: GroundObjectTemplates) -> None:
        self._accessible_units = []
        all_units: Iterator[UnitType[Any]] = itertools.chain(
            self.ground_units,
            self.infantry_units,
            self.air_defense_units,
            self.naval_units,
            self.missiles,
            (
                ground_unit
                for preset_group in self.preset_groups
                for ground_unit in preset_group.ground_units
            ),
            (
                ship_unit
                for preset_group in self.preset_groups
                for ship_unit in preset_group.ship_units
            ),
        )
        for unit in all_units:
            if unit not in self._accessible_units:
                self._accessible_units.append(unit)

    def initialize(
        self, all_templates: GroundObjectTemplates, mod_settings: ModSettings
    ) -> None:
        # Apply the mod settings
        self._apply_mod_settings(mod_settings)
        # Load all accessible units and store them for performant later usage
        self._load_accessible_units(all_templates)
        # Load all faction compatible templates
        self._load_templates(all_templates)
        # Load Unit Groups
        self._load_unit_groups()

    def _add_unit_group(self, unit_group: UnitGroup, merge: bool = True) -> None:
        if not unit_group.templates:
            unit_group.load_templates(self)
        if not unit_group.templates:
            # Empty templates will throw an error on generation
            logging.error(
                f"Skipping Unit group {unit_group.name} as no templates are available to generate the group"
            )
            return
        if unit_group.role in self.unit_groups:
            for group in self.unit_groups[unit_group.role]:
                if merge and all(task in group.tasks for task in unit_group.tasks):
                    # Update existing group if same tasking
                    group.update_from_unit_group(unit_group)
                    return
            # Add new Unit_group
            self.unit_groups[unit_group.role].append(unit_group)
        else:
            self.unit_groups[unit_group.role] = [unit_group]

    def _load_unit_groups(self) -> None:
        # This function will create all the UnitGroups for the faction
        # It will create a unit group for each global Template, Building or
        # Legacy supported templates (not yet migrated from the generators).
        # For every preset_group there will be a separate UnitGroup so no mixed
        # UnitGroups will be generated for them. Special groups like complex SAM Systems
        self.unit_groups = {}

        # Generate UnitGroups for all global templates
        for role, template in self.templates.templates:
            if template.generic or role == GroupRole.Building:
                # Build groups for global templates and buildings
                self._add_group_for_template(role, template)

        # Add preset groups
        for preset_group in self.preset_groups:
            # Add as separate group, do not merge with generic groups!
            self._add_unit_group(preset_group, False)

    def _add_group_for_template(
        self, role: GroupRole, template: GroundObjectTemplate
    ) -> None:
        unit_group = UnitGroup(
            f"{role.value}: {', '.join([t.value for t in template.tasks])}",
            [u for u in template.units if isinstance(u, GroundUnitType)],
            [u for u in template.units if isinstance(u, ShipUnitType)],
            list(template.statics),
            role,
        )
        unit_group.tasks = template.tasks
        unit_group.set_templates([template])
        self._add_unit_group(unit_group)

    def initialize_group_template(
        self, group: GroupTemplate, faction_sensitive: bool = True
    ) -> bool:
        # Sensitive defines if the initialization should check if the unit is available
        # to this faction or not. It is disabled for migration only atm.
        unit_types = [
            t
            for t in group.unit_types
            if not faction_sensitive or self.has_access_to_unit_type(t)
        ]

        alternative_types = []
        for accessible_unit in self.accessible_units:
            if accessible_unit.unit_class in group.unit_classes:
                unit_types.append(accessible_unit.dcs_id)
            if accessible_unit.unit_class in group.alternative_classes:
                alternative_types.append(accessible_unit.dcs_id)

        if not unit_types and not alternative_types and not group.optional:
            raise StopIteration

        types = unit_types or alternative_types
        group.set_possible_types(types)
        return len(types) > 0

    def _load_templates(self, all_templates: GroundObjectTemplates) -> None:
        self.templates = GroundObjectTemplates()
        # This loads all templates which are usable by the faction
        for role, template in all_templates.templates:
            # Make a deep copy of a template and add it to the template_list.
            # This is required to have faction independent templates. Otherwise
            # the reference would be the same and changes would affect all.
            faction_template = copy.deepcopy(template)
            try:
                faction_template.groups[:] = [
                    group_template
                    for group_template in faction_template.groups
                    if self.initialize_group_template(group_template)
                ]
                if (
                    role == GroupRole.Building
                    and GroupTask.StrikeTarget in template.tasks
                    and faction_template.category not in self.building_set
                ):
                    # Special handling for strike targets. Skip if not supported by faction
                    continue
                if faction_template.groups:
                    self.templates.add_template(role, faction_template)
                    continue
            except StopIteration:
                pass

            logging.info(f"{self.name} can not use template {template.name}")

    @classmethod
    def from_json(cls: Type[Faction], json: Dict[str, Any]) -> Faction:
        faction = Faction(locales=json.get("locales"))

        faction.country = json.get("country", "/")
        if faction.country not in [c.name for c in country_dict.values()]:
            raise AssertionError(
                'Faction\'s country ("{}") is not a valid DCS country ID'.format(
                    faction.country
                )
            )

        faction.name = json.get("name", "")
        if not faction.name:
            raise AssertionError("Faction has no valid name")

        faction.authors = json.get("authors", "")
        faction.description = json.get("description", "")

        faction.aircrafts = [AircraftType.named(n) for n in json.get("aircrafts", [])]
        faction.awacs = [AircraftType.named(n) for n in json.get("awacs", [])]
        faction.tankers = [AircraftType.named(n) for n in json.get("tankers", [])]

        faction.aircrafts = list(
            set(faction.aircrafts + faction.awacs + faction.tankers)
        )

        faction.frontline_units = [
            GroundUnitType.named(n) for n in json.get("frontline_units", [])
        ]
        faction.artillery_units = [
            GroundUnitType.named(n) for n in json.get("artillery_units", [])
        ]
        faction.infantry_units = [
            GroundUnitType.named(n) for n in json.get("infantry_units", [])
        ]
        faction.logistics_units = [
            GroundUnitType.named(n) for n in json.get("logistics_units", [])
        ]
        faction.air_defense_units = [
            GroundUnitType.named(n) for n in json.get("air_defense_units", [])
        ]
        faction.missiles = [GroundUnitType.named(n) for n in json.get("missiles", [])]

        faction.preset_groups = [
            UnitGroup.named(n) for n in json.get("preset_groups", [])
        ]

        faction.requirements = json.get("requirements", {})

        faction.carrier_names = json.get("carrier_names", [])
        faction.helicopter_carrier_names = json.get("helicopter_carrier_names", [])

        faction.naval_units = [
            ShipUnitType.named(n) for n in json.get("naval_units", [])
        ]

        faction.has_jtac = json.get("has_jtac", False)
        jtac_name = json.get("jtac_unit", None)
        if jtac_name is not None:
            faction.jtac_unit = AircraftType.named(jtac_name)
        else:
            faction.jtac_unit = None

        # Load doctrine
        doctrine = json.get("doctrine", "modern")
        if doctrine == "modern":
            faction.doctrine = MODERN_DOCTRINE
        elif doctrine == "coldwar":
            faction.doctrine = COLDWAR_DOCTRINE
        elif doctrine == "ww2":
            faction.doctrine = WWII_DOCTRINE
        else:
            faction.doctrine = MODERN_DOCTRINE

        # Load the building set
        building_set = json.get("building_set", "default")
        if building_set == "default":
            faction.building_set = DEFAULT_AVAILABLE_BUILDINGS
        elif building_set == "ww2free":
            faction.building_set = WW2_FREE
        elif building_set == "ww2ally":
            faction.building_set = WW2_ALLIES_BUILDINGS
        elif building_set == "ww2germany":
            faction.building_set = WW2_GERMANY_BUILDINGS
        else:
            faction.building_set = DEFAULT_AVAILABLE_BUILDINGS

        # Load liveries override
        faction.liveries_overrides = {}
        liveries_overrides = json.get("liveries_overrides", {})
        for name, livery in liveries_overrides.items():
            aircraft = AircraftType.named(name)
            faction.liveries_overrides[aircraft] = [s.lower() for s in livery]

        faction.unrestricted_satnav = json.get("unrestricted_satnav", False)

        # Templates
        faction.templates = GroundObjectTemplates()

        return faction

    @property
    def ground_units(self) -> Iterator[GroundUnitType]:
        yield from self.artillery_units
        yield from self.frontline_units
        yield from self.logistics_units

    def infantry_with_class(self, unit_class: UnitClass) -> Iterator[GroundUnitType]:
        for unit in self.infantry_units:
            if unit.unit_class is unit_class:
                yield unit

    def groups_for_role_and_task(
        self, group_role: GroupRole, group_task: Optional[GroupTask] = None
    ) -> list[UnitGroup]:
        if group_role not in self.unit_groups:
            return []
        groups = []
        for unit_group in self.unit_groups[group_role]:
            if not group_task or group_task in unit_group.tasks:
                groups.append(unit_group)
        return groups

    def groups_for_role_and_tasks(
        self, group_role: GroupRole, tasks: list[GroupTask]
    ) -> list[UnitGroup]:
        groups = []
        for task in tasks:
            for group in self.groups_for_role_and_task(group_role, task):
                if group not in groups:
                    groups.append(group)
        return groups

    def random_group_for_role(self, group_role: GroupRole) -> Optional[UnitGroup]:
        unit_groups = self.groups_for_role_and_task(group_role)
        return random.choice(unit_groups) if unit_groups else None

    def random_group_for_role_and_task(
        self, group_role: GroupRole, group_task: GroupTask
    ) -> Optional[UnitGroup]:
        unit_groups = self.groups_for_role_and_task(group_role, group_task)
        return random.choice(unit_groups) if unit_groups else None

    def random_group_for_role_and_tasks(
        self, group_role: GroupRole, tasks: list[GroupTask]
    ) -> Optional[UnitGroup]:
        unit_groups = self.groups_for_role_and_tasks(group_role, tasks)
        return random.choice(unit_groups) if unit_groups else None

    def _apply_mod_settings(self, mod_settings: ModSettings) -> None:
        # aircraft
        if not mod_settings.a4_skyhawk:
            self.remove_aircraft("A-4E-C")
        if not mod_settings.hercules:
            self.remove_aircraft("Hercules")
        if not mod_settings.uh_60l:
            self.remove_aircraft("UH-60L")
            self.remove_aircraft("KC130J")
        if not mod_settings.f22_raptor:
            self.remove_aircraft("F-22A")
        if not mod_settings.f104_starfighter:
            self.remove_aircraft("VSN_F104G")
            self.remove_aircraft("VSN_F104S")
            self.remove_aircraft("VSN_F104S_AG")
        if not mod_settings.jas39_gripen:
            self.remove_aircraft("JAS39Gripen")
            self.remove_aircraft("JAS39Gripen_AG")
        if not mod_settings.su57_felon:
            self.remove_aircraft("Su-57")
        # frenchpack
        if not mod_settings.frenchpack:
            self.remove_vehicle("AMX10RCR")
            self.remove_vehicle("SEPAR")
            self.remove_vehicle("ERC")
            self.remove_vehicle("M120")
            self.remove_vehicle("AA20")
            self.remove_vehicle("TRM2000")
            self.remove_vehicle("TRM2000_Citerne")
            self.remove_vehicle("TRM2000_AA20")
            self.remove_vehicle("TRMMISTRAL")
            self.remove_vehicle("VABH")
            self.remove_vehicle("VAB_RADIO")
            self.remove_vehicle("VAB_50")
            self.remove_vehicle("VIB_VBR")
            self.remove_vehicle("VAB_HOT")
            self.remove_vehicle("VAB_MORTIER")
            self.remove_vehicle("VBL50")
            self.remove_vehicle("VBLANF1")
            self.remove_vehicle("VBL-radio")
            self.remove_vehicle("VBAE")
            self.remove_vehicle("VBAE_MMP")
            self.remove_vehicle("AMX-30B2")
            self.remove_vehicle("Tracma")
            self.remove_vehicle("JTACFP")
            self.remove_vehicle("SHERIDAN")
            self.remove_vehicle("Leclerc_XXI")
            self.remove_vehicle("Toyota_bleu")
            self.remove_vehicle("Toyota_vert")
            self.remove_vehicle("Toyota_desert")
            self.remove_vehicle("Kamikaze")
            self.remove_vehicle("AMX1375")
            self.remove_vehicle("AMX1390")
            self.remove_vehicle("VBCI")
            self.remove_vehicle("T62")
            self.remove_vehicle("T64BV")
            self.remove_vehicle("T72M")
            self.remove_vehicle("KORNET")
        # high digit sams
        if not mod_settings.high_digit_sams:
            self.remove_presets("SA-10B/S-300PS")
            self.remove_presets("SA-12/S-300V")
            self.remove_presets("SA-20/S-300PMU-1")
            self.remove_presets("SA-20B/S-300PMU-2")
            self.remove_presets("SA-23/S-300VM")
            self.remove_presets("SA-17")
            self.remove_presets("KS-19")

    def remove_aircraft(self, name: str) -> None:
        for i in self.aircrafts:
            if i.dcs_unit_type.id == name:
                self.aircrafts.remove(i)

    def remove_presets(self, name: str) -> None:
        for pg in self.preset_groups:
            if pg.name == name:
                self.preset_groups.remove(pg)

    def remove_vehicle(self, name: str) -> None:
        for i in self.frontline_units:
            if i.dcs_unit_type.id == name:
                self.frontline_units.remove(i)


def load_ship(name: str) -> Optional[Type[ShipType]]:
    if (ship := getattr(dcs.ships, name, None)) is not None:
        return ship
    logging.error(f"FACTION ERROR : Unable to find {name} in dcs.ships")
    return None


def load_all_ships(data: list[str]) -> List[Type[ShipType]]:
    items = []
    for name in data:
        item = load_ship(name)
        if item is not None:
            items.append(item)
    return items