From bc26eb3f5e9676efb53785dc83a36e8c851d9ea0 Mon Sep 17 00:00:00 2001
From: Raffson <Raffson@users.noreply.github.com>
Date: Sun, 21 Jan 2024 13:36:38 +0100
Subject: [PATCH] Fix naval aircraft not always returning to carrier

---
 changelog.md                                             | 1 +
 game/missiongenerator/aircraft/waypoints/landingpoint.py | 7 ++++++-
 game/missiongenerator/tgogenerator.py                    | 1 +
 game/theater/controlpoint.py                             | 6 ++++++
 4 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/changelog.md b/changelog.md
index 869426f4..21449798 100644
--- a/changelog.md
+++ b/changelog.md
@@ -46,6 +46,7 @@
 * **[Mission Generation]** Fix infinite loop when using "Fast-Forward to first contact"
 * **[Capture Logic]** Release all parking slots when an airbase is captured
 * **[Modding]** Swedish Military Assets Pack air defence presets are now correctly removed from the faction when the mod is disabled.
+* **[Mission Generation]** Naval aircraft not always returning to carrier
 
 # Retribution v1.2.1 (hotfix)
 
diff --git a/game/missiongenerator/aircraft/waypoints/landingpoint.py b/game/missiongenerator/aircraft/waypoints/landingpoint.py
index 69c1e71d..f7b353d2 100644
--- a/game/missiongenerator/aircraft/waypoints/landingpoint.py
+++ b/game/missiongenerator/aircraft/waypoints/landingpoint.py
@@ -1,5 +1,6 @@
 from dcs.point import MovingPoint, PointAction
 
+from game.theater import NavalControlPoint
 from .pydcswaypointbuilder import PydcsWaypointBuilder
 
 
@@ -9,5 +10,9 @@ class LandingPointBuilder(PydcsWaypointBuilder):
         waypoint.type = "Land"
         waypoint.action = PointAction.Landing
         if (control_point := self.waypoint.control_point) is not None:
-            waypoint.airdrome_id = control_point.airdrome_id_for_landing
+            if isinstance(control_point, NavalControlPoint):
+                waypoint.helipad_id = control_point.airdrome_id_for_landing
+                waypoint.link_unit = control_point.airdrome_id_for_landing
+            else:
+                waypoint.airdrome_id = control_point.airdrome_id_for_landing
         return waypoint
diff --git a/game/missiongenerator/tgogenerator.py b/game/missiongenerator/tgogenerator.py
index 8a1f61b8..a01f33e1 100644
--- a/game/missiongenerator/tgogenerator.py
+++ b/game/missiongenerator/tgogenerator.py
@@ -590,6 +590,7 @@ class GenericCarrierGenerator(GroundObjectGenerator):
                         f"Error generating carrier group for {self.control_point.name}"
                     )
                 ship_group.units[0].type = carrier_type.id
+                self.control_point.carrier_id = ship_group.units[0].id
                 if self.control_point.tacan is None:
                     tacan = self.tacan_registry.alloc_for_band(
                         TacanBand.X, TacanUsage.TransmitReceive
diff --git a/game/theater/controlpoint.py b/game/theater/controlpoint.py
index 65852d99..69908e44 100644
--- a/game/theater/controlpoint.py
+++ b/game/theater/controlpoint.py
@@ -1336,6 +1336,8 @@ class Airfield(ControlPoint, CTLD):
 class NavalControlPoint(
     ControlPoint, ABC, Link4Container, TacanContainer, ICLSContainer
 ):
+    carrier_id: Optional[int] = None
+
     @property
     def is_fleet(self) -> bool:
         return True
@@ -1427,6 +1429,10 @@ class NavalControlPoint(
             return ControlPointStatus.Damaged
         return ControlPointStatus.Functional
 
+    @property
+    def airdrome_id_for_landing(self) -> Optional[int]:
+        return self.carrier_id
+
 
 class Carrier(NavalControlPoint):
     def __init__(