From de325c12087ac96a664462e49dd6281f4f6a871e Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Mon, 28 Dec 2020 00:47:42 -0800
Subject: [PATCH] Refunding unfulfilled orders on capture.

Fixes https://github.com/Khopa/dcs_liberation/issues/682
---
 game/event/event.py          | 14 ++++++++++++++
 game/game.py                 |  6 ++++++
 game/theater/controlpoint.py |  8 ++++++--
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/game/event/event.py b/game/event/event.py
index 836a0984..61d7e69d 100644
--- a/game/event/event.py
+++ b/game/event/event.py
@@ -15,6 +15,7 @@ from game.operation.operation import Operation
 from game.theater import ControlPoint
 from gen import AirTaskingOrder
 from gen.ground_forces.combat_stance import CombatStance
+from ..db import PRICES
 from ..unitmap import UnitMap
 
 if TYPE_CHECKING:
@@ -362,6 +363,19 @@ class UnitsDeliveryEvent(Event):
         for k, v in units.items():
             self.units[k] = self.units.get(k, 0) + v
 
+    def refund_all(self) -> None:
+        while self.units:
+            unit_type, count = self.units.popitem()
+            try:
+                price = PRICES[unit_type]
+            except KeyError:
+                logging.error(f"Could not refund {unit_type.id}, price unknown")
+                continue
+
+            logging.info(
+                f"Refunding {count} {unit_type.id} at {self.to_cp.name}")
+            self.game.adjust_budget(price * count, player=self.to_cp.captured)
+
     def skip(self) -> None:
         for k, v in self.units.items():
             if self.to_cp.captured:
diff --git a/game/game.py b/game/game.py
index 069f22db..7482b335 100644
--- a/game/game.py
+++ b/game/game.py
@@ -189,6 +189,12 @@ class Game:
                                         front_line.control_point_a,
                                         front_line.control_point_b)
 
+    def adjust_budget(self, amount: float, player: bool) -> None:
+        if player:
+            self.budget += amount
+        else:
+            self.enemy_budget += amount
+
     def process_player_income(self):
         self.budget += Income(self, player=True).total
 
diff --git a/game/theater/controlpoint.py b/game/theater/controlpoint.py
index 8904e9b0..50efad63 100644
--- a/game/theater/controlpoint.py
+++ b/game/theater/controlpoint.py
@@ -21,16 +21,17 @@ from dcs.unittype import FlyingType
 
 from game import db
 from gen.ground_forces.ai_ground_planner_db import TYPE_SHORAD
-from gen.runways import RunwayAssigner, RunwayData
 from gen.ground_forces.combat_stance import CombatStance
+from gen.runways import RunwayAssigner, RunwayData
 from .base import Base
 from .missiontarget import MissionTarget
 from .theatergroundobject import (
     BaseDefenseGroundObject,
     EwrGroundObject,
+    GenericCarrierGroundObject,
     SamGroundObject,
     TheaterGroundObject,
-    VehicleGroupGroundObject, GenericCarrierGroundObject,
+    VehicleGroupGroundObject,
 )
 from ..weather import Conditions
 
@@ -366,6 +367,9 @@ class ControlPoint(MissionTarget, ABC):
 
     # TODO: Should be Airbase specific.
     def capture(self, game: Game, for_player: bool) -> None:
+        if self.pending_unit_deliveries is not None:
+            self.pending_unit_deliveries.refund_all()
+
         if for_player:
             self.captured = True
         else: