From a366aebbb233ba5f83e9dbff741245e193d4239b Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Fri, 21 Jul 2023 22:46:39 -0700
Subject: [PATCH] Add laser code config parsing and prop generation.

---
 game/dcs/aircrafttype.py          |  5 +++
 game/dcs/lasercodeconfig.py       | 48 +++++++++++++++++++++++
 tests/dcs/__init__.py             |  0
 tests/dcs/test_lasercodeconfig.py | 65 +++++++++++++++++++++++++++++++
 4 files changed, 118 insertions(+)
 create mode 100644 game/dcs/lasercodeconfig.py
 create mode 100644 tests/dcs/__init__.py
 create mode 100644 tests/dcs/test_lasercodeconfig.py

diff --git a/game/dcs/aircrafttype.py b/game/dcs/aircrafttype.py
index e31582f4..ce0d5168 100644
--- a/game/dcs/aircrafttype.py
+++ b/game/dcs/aircrafttype.py
@@ -16,6 +16,7 @@ from dcs.unittype import FlyingType
 from dcs.weapons_data import weapon_ids
 
 from game.data.units import UnitClass
+from game.dcs.lasercodeconfig import LaserCodeConfig
 from game.dcs.unittype import UnitType
 from game.persistency import user_custom_weapon_injections_dir
 from game.radio.channels import (
@@ -207,6 +208,7 @@ class AircraftType(UnitType[Type[FlyingType]]):
     has_built_in_target_pod: bool
 
     task_priorities: dict[FlightType, int]
+    laser_code_configs: list[LaserCodeConfig]
 
     _by_name: ClassVar[dict[str, AircraftType]] = {}
     _by_unit_type: ClassVar[dict[type[FlyingType], list[AircraftType]]] = defaultdict(
@@ -501,6 +503,9 @@ class AircraftType(UnitType[Type[FlyingType]]):
                 can_carry_crates=data.get("can_carry_crates", aircraft.helicopter),
                 has_built_in_target_pod=data.get("has_built_in_target_pod", False),
                 task_priorities=task_priorities,
+                laser_code_configs=[
+                    LaserCodeConfig.from_yaml(d) for d in data.get("laser_codes", [])
+                ],
             )
 
     @staticmethod
diff --git a/game/dcs/lasercodeconfig.py b/game/dcs/lasercodeconfig.py
new file mode 100644
index 00000000..9e8260db
--- /dev/null
+++ b/game/dcs/lasercodeconfig.py
@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class LaserCodeConfig(ABC):
+    def __init__(self, pylon: int) -> None:
+        self.pylon = pylon
+
+    @staticmethod
+    def from_yaml(data: dict[str, Any]) -> LaserCodeConfig:
+        pylon = data["pylon"]
+        if (property_def := data.get("property")) is not None:
+            return SinglePropertyLaserCodeConfig(
+                pylon, property_def["id"], int(property_def["digits"])
+            )
+        return MultiplePropertyLaserCodeConfig(
+            pylon, [(d["id"], d["digit"]) for d in data["properties"]]
+        )
+
+    @abstractmethod
+    def property_dict_for_code(self, code: int) -> dict[str, int]:
+        ...
+
+
+class SinglePropertyLaserCodeConfig(LaserCodeConfig):
+    def __init__(self, pylon: int, property_id: str, digits: int) -> None:
+        super().__init__(pylon)
+        self.property_id = property_id
+        self.digits = digits
+
+    def property_dict_for_code(self, code: int) -> dict[str, int]:
+        return {self.property_id: code % 10**self.digits}
+
+
+class MultiplePropertyLaserCodeConfig(LaserCodeConfig):
+    def __init__(
+        self, pylon: int, property_digit_mappings: list[tuple[str, int]]
+    ) -> None:
+        super().__init__(pylon)
+        self.property_digit_mappings = property_digit_mappings
+
+    def property_dict_for_code(self, code: int) -> dict[str, int]:
+        d = {}
+        for prop_id, idx in self.property_digit_mappings:
+            d[prop_id] = code // 10**idx % 10
+        return d
diff --git a/tests/dcs/__init__.py b/tests/dcs/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/dcs/test_lasercodeconfig.py b/tests/dcs/test_lasercodeconfig.py
new file mode 100644
index 00000000..f560f63a
--- /dev/null
+++ b/tests/dcs/test_lasercodeconfig.py
@@ -0,0 +1,65 @@
+from game.dcs.lasercodeconfig import (
+    SinglePropertyLaserCodeConfig,
+    MultiplePropertyLaserCodeConfig,
+    LaserCodeConfig,
+)
+
+
+def test_singlepropertylasercodeproperty() -> None:
+    config = SinglePropertyLaserCodeConfig(0, "code", 3)
+    assert config.property_dict_for_code(1688) == {"code": 688}
+    assert config.property_dict_for_code(1000) == {"code": 0}
+    assert config.property_dict_for_code(1234) == {"code": 234}
+    assert config.property_dict_for_code(1) == {"code": 1}
+
+
+def test_multiplepropertylasercodeproperty() -> None:
+    config = MultiplePropertyLaserCodeConfig(
+        0,
+        [
+            ("digit0", 0),
+            ("digit1", 1),
+            ("digit2", 2),
+        ],
+    )
+    assert config.property_dict_for_code(1688) == {
+        "digit0": 8,
+        "digit1": 8,
+        "digit2": 6,
+    }
+    assert config.property_dict_for_code(1000) == {
+        "digit0": 0,
+        "digit1": 0,
+        "digit2": 0,
+    }
+    assert config.property_dict_for_code(1234) == {
+        "digit0": 4,
+        "digit1": 3,
+        "digit2": 2,
+    }
+    assert config.property_dict_for_code(1) == {"digit0": 1, "digit1": 0, "digit2": 0}
+
+
+def test_lasercodeconfig_from_yaml() -> None:
+    config = LaserCodeConfig.from_yaml(
+        {"pylon": 0, "property": {"id": "code", "digits": 3}}
+    )
+    assert config.pylon == 0
+    assert config.property_dict_for_code(1688) == {"code": 688}
+
+    config = LaserCodeConfig.from_yaml(
+        {
+            "pylon": 1,
+            "properties": [
+                {"id": "digit0", "digit": 0},
+                {"id": "digit1", "digit": 1},
+                {"id": "digit2", "digit": 2},
+            ],
+        }
+    )
+    assert config.pylon == 1
+    assert config.property_dict_for_code(1688) == {
+        "digit0": 8,
+        "digit1": 8,
+        "digit2": 6,
+    }