From a2c10f1c7a74b97181c4eaf901d1afcc3d270b01 Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Thu, 27 May 2021 18:51:53 -0700
Subject: [PATCH] Check for compatible squadrons when buying planes.

https://github.com/dcs-liberation/dcs_liberation/issues/276
---
 game/procurement.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/game/procurement.py b/game/procurement.py
index a65879a8..29c1ba15 100644
--- a/game/procurement.py
+++ b/game/procurement.py
@@ -50,6 +50,7 @@ class ProcurementAi:
 
         self.game = game
         self.is_player = for_player
+        self.air_wing = game.air_wing_for(for_player)
         self.faction = faction
         self.manage_runways = manage_runways
         self.manage_front_line = manage_front_line
@@ -163,23 +164,31 @@ class ProcurementAi:
 
         return budget
 
-    def _affordable_aircraft_of_types(
+    def _affordable_aircraft_for_task(
         self,
-        types: List[Type[FlyingType]],
+        task: FlightType,
         airbase: ControlPoint,
         number: int,
         max_price: float,
     ) -> Optional[Type[FlyingType]]:
         best_choice: Optional[Type[FlyingType]] = None
-        for unit in [u for u in types if u in self.faction.aircrafts]:
+        for unit in aircraft_for_task(task):
+            if unit not in self.faction.aircrafts:
+                continue
             if db.PRICES[unit] * number > max_price:
                 continue
             if not airbase.can_operate(unit):
                 continue
 
-            # Affordable and compatible. To keep some variety, skip with a 50/50
-            # chance. Might be a good idea to have the chance to skip based on
-            # the price compared to the rest of the choices.
+            for squadron in self.air_wing.squadrons_for(unit):
+                if task in squadron.mission_types:
+                    break
+            else:
+                continue
+
+            # Affordable, compatible, and we have a squadron capable of the task. To
+            # keep some variety, skip with a 50/50 chance. Might be a good idea to have
+            # the chance to skip based on the price compared to the rest of the choices.
             best_choice = unit
             if random.choice([True, False]):
                 break
@@ -188,8 +197,8 @@ class ProcurementAi:
     def affordable_aircraft_for(
         self, request: AircraftProcurementRequest, airbase: ControlPoint, budget: float
     ) -> Optional[Type[FlyingType]]:
-        return self._affordable_aircraft_of_types(
-            aircraft_for_task(request.task_capability), airbase, request.number, budget
+        return self._affordable_aircraft_for_task(
+            request.task_capability, airbase, request.number, budget
         )
 
     def fulfill_aircraft_request(