From f112f227ccc7832f3f43b7198730b63f36c47cb9 Mon Sep 17 00:00:00 2001
From: Raffson <Raffson@users.noreply.github.com>
Date: Sun, 5 Jan 2025 14:46:20 +0100
Subject: [PATCH] Update armed forces when preset-group is added

---
 qt_ui/windows/AirWingDialog.py                | 35 +++++++++++++------
 .../newgame/WizardPages/QFactionSelection.py  |  6 ++--
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/qt_ui/windows/AirWingDialog.py b/qt_ui/windows/AirWingDialog.py
index 9aa850c2..4b0eef31 100644
--- a/qt_ui/windows/AirWingDialog.py
+++ b/qt_ui/windows/AirWingDialog.py
@@ -18,7 +18,9 @@ from PySide6.QtWidgets import (
     QPushButton,
 )
 
+from game.armedforces.armedforces import ArmedForces
 from game.ato.flight import Flight
+from game.factions import Faction
 from game.server import EventStream
 from game.sim import GameUpdateEvents
 from game.squadrons import Squadron
@@ -270,20 +272,24 @@ class AirWingTabs(QTabWidget):
             w = QWidget(layout=layout)
             self.addTab(w, "Cheats")
 
+        qfu_ownfor = QFactionUnits(
+            game_model.game.coalition_for(True).faction,
+            self,
+            show_jtac=True,
+        )
+        qfu_ownfor.preset_groups_changed.connect(self.preset_group_updated_ownfor)
         self.addTab(
-            QFactionUnits(
-                game_model.game.coalition_for(True).faction,
-                self,
-                show_jtac=True,
-            ),
+            qfu_ownfor,
             "Faction OWNFOR",
         )
+        qfu_opfor = QFactionUnits(
+            game_model.game.coalition_for(False).faction,
+            self,
+            show_jtac=True,
+        )
+        qfu_opfor.preset_groups_changed.connect(self.preset_group_updated_opfor)
         self.addTab(
-            QFactionUnits(
-                game_model.game.coalition_for(False).faction,
-                self,
-                show_jtac=False,
-            ),
+            qfu_opfor,
             "Faction OPFOR",
         )
 
@@ -293,6 +299,15 @@ class AirWingTabs(QTabWidget):
         EventStream.put_nowait(events)
         self.game_model.ato_model.on_sim_update(events)
 
+    def preset_group_updated_ownfor(self, f: Faction) -> None:
+        self.preset_group_updated(f, player=True)
+
+    def preset_group_updated_opfor(self, f: Faction) -> None:
+        self.preset_group_updated(f, player=False)
+
+    def preset_group_updated(self, f: Faction, player: bool) -> None:
+        self.game_model.game.coalition_for(player).armed_forces = ArmedForces(f)
+
 
 class AirWingDialog(QDialog):
     """Dialog window showing the player's air wing."""
diff --git a/qt_ui/windows/newgame/WizardPages/QFactionSelection.py b/qt_ui/windows/newgame/WizardPages/QFactionSelection.py
index e74a10cd..f0abe4ba 100644
--- a/qt_ui/windows/newgame/WizardPages/QFactionSelection.py
+++ b/qt_ui/windows/newgame/WizardPages/QFactionSelection.py
@@ -1,12 +1,11 @@
 from __future__ import unicode_literals
 
 import json
-
 from copy import deepcopy
 from typing import Union, Callable, Set, Optional, List
 
 from PySide6 import QtWidgets, QtGui
-from PySide6.QtCore import Qt
+from PySide6.QtCore import Qt, Signal
 from PySide6.QtWidgets import (
     QScrollArea,
     QWidget,
@@ -37,6 +36,8 @@ from qt_ui.windows.newgame.jinja_env import jinja_env
 
 
 class QFactionUnits(QScrollArea):
+    preset_groups_changed = Signal(Faction)
+
     def __init__(self, faction: Faction, parent=None, show_jtac: bool = False):
         super().__init__()
         self.setWidgetResizable(True)
@@ -302,6 +303,7 @@ class QFactionUnits(QScrollArea):
             # invalidate the cached property
             del self.faction.__dict__["accessible_units"]
         self.updateFaction(self.faction)
+        self.preset_groups_changed.emit(self.faction)
 
     def updateFaction(self, faction: Faction):
         self.faction = faction