From 4aacc68f4ac97d93519695c6a10b05f91369204c Mon Sep 17 00:00:00 2001
From: Raffson <Raffson@users.noreply.github.com>
Date: Sun, 6 Oct 2024 22:29:29 +0200
Subject: [PATCH] Include layouts & groups from Saved Games folder

Partial #244
---
 game/armedforces/forcegroup.py | 16 ++++++++++---
 game/layout/layoutloader.py    | 41 +++++++++++++++++++++++++++-------
 game/persistency.py            |  8 +++++++
 3 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/game/armedforces/forcegroup.py b/game/armedforces/forcegroup.py
index 731dfbe1..51c132f1 100644
--- a/game/armedforces/forcegroup.py
+++ b/game/armedforces/forcegroup.py
@@ -9,6 +9,7 @@ from typing import Any, ClassVar, Iterator, Optional, TYPE_CHECKING, Type
 import yaml
 from dcs.unittype import ShipType, StaticType, UnitType as DcsUnitType, VehicleType
 
+from game import persistency
 from game.data.groups import GroupTask
 from game.dcs.groundunittype import GroundUnitType
 from game.dcs.helpers import static_type_from_name
@@ -317,7 +318,18 @@ class ForceGroup:
 
     @classmethod
     def _load_all(cls) -> None:
-        for file in Path("resources/groups").glob("*.yaml"):
+        locations = [
+            Path("resources/groups"),
+            persistency.groups_dir(),
+        ]
+        for path in locations:
+            cls._process_path(path)
+
+        cls._loaded = True
+
+    @classmethod
+    def _process_path(cls, path: Path) -> None:
+        for file in path.glob("*.yaml"):
             if not file.is_file():
                 raise RuntimeError(f"{file.name} is not a valid ForceGroup")
 
@@ -365,5 +377,3 @@ class ForceGroup:
             )
 
             cls._by_name[force_group.name] = force_group
-
-        cls._loaded = True
diff --git a/game/layout/layoutloader.py b/game/layout/layoutloader.py
index 035ef213..986b4290 100644
--- a/game/layout/layoutloader.py
+++ b/game/layout/layoutloader.py
@@ -5,6 +5,7 @@ import logging
 import pickle
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
+from copy import deepcopy
 from pathlib import Path
 from typing import Iterator
 
@@ -41,6 +42,8 @@ LAYOUT_TYPES = {
     GroupRole.DEFENSES: DefensesLayout,
 }
 
+LOCATIONS_TO_CHECK: list[Path] = []
+
 
 class LayoutLoader:
     # Map of all available layouts indexed by name
@@ -51,9 +54,19 @@ class LayoutLoader:
 
     def initialize(self) -> None:
         if not self._layouts:
+            self.initialize_locations_to_check()
             with logged_duration("Loading layouts"):
                 self.load_templates()
 
+    @staticmethod
+    def initialize_locations_to_check() -> None:
+        global LOCATIONS_TO_CHECK
+        if not LOCATIONS_TO_CHECK:
+            LOCATIONS_TO_CHECK = [
+                Path(LAYOUT_DIR),
+                persistency.layouts_dir(),
+            ]
+
     @property
     def layouts(self) -> Iterator[TgoLayout]:
         self.initialize()
@@ -83,14 +96,9 @@ class LayoutLoader:
         self._layouts = {}
         mappings: dict[str, list[LayoutMapping]] = defaultdict(list)
         with logged_duration("Parsing mapping yamls"):
-            for file in Path(LAYOUT_DIR).rglob("*.yaml"):
-                if not file.is_file():
-                    raise RuntimeError(f"{file.name} is not a file")
-                with file.open("r", encoding="utf-8") as f:
-                    mapping_dict = yaml.safe_load(f)
-
-                template_map = LayoutMapping.from_dict(mapping_dict, f.name)
-                mappings[template_map.layout_file].append(template_map)
+            for path_to_check in LOCATIONS_TO_CHECK:
+                for file in path_to_check.rglob("*.yaml"):
+                    self._process_yaml(file, mappings)
 
         with logged_duration(f"Parsing all layout miz multithreaded"):
             with ThreadPoolExecutor() as exe:
@@ -105,6 +113,15 @@ class LayoutLoader:
         logging.info(f"Imported {len(self._layouts)} layouts")
         self._dump_templates()
 
+    @staticmethod
+    def _process_yaml(file: Path, mappings: dict[str, list[LayoutMapping]]) -> None:
+        if not file.is_file():
+            raise RuntimeError(f"{file.name} is not a file")
+        with file.open("r", encoding="utf-8") as f:
+            mapping_dict = yaml.safe_load(f)
+        template_map = LayoutMapping.from_dict(mapping_dict, f.name)
+        mappings[template_map.layout_file].append(template_map)
+
     def _dump_templates(self) -> None:
         file = persistency.base_path() / LAYOUT_DUMP
         dump = (VERSION, self._layouts)
@@ -112,6 +129,14 @@ class LayoutLoader:
             pickle.dump(dump, fdata)
 
     def _load_from_miz(self, miz: str, mappings: list[LayoutMapping]) -> None:
+        path = Path(miz)
+        locations_to_check = deepcopy(LOCATIONS_TO_CHECK)
+        while not path.exists() and locations_to_check:
+            path = locations_to_check.pop() / miz
+            miz = path.absolute().as_posix()
+        if not path.exists():
+            logging.warning(f"Layout miz file not found: '{miz}'")
+            return
         template_position: dict[str, Point] = {}
         temp_mis = dcs.Mission()
         with logged_duration(f"Parsing {miz}"):
diff --git a/game/persistency.py b/game/persistency.py
index 5ade5dc1..053a80c0 100644
--- a/game/persistency.py
+++ b/game/persistency.py
@@ -125,6 +125,14 @@ def debug_dir() -> Path:
     return base_path() / "Retribution" / "Debug"
 
 
+def groups_dir() -> Path:
+    return base_path() / "Retribution" / "Groups"
+
+
+def layouts_dir() -> Path:
+    return base_path() / "Retribution" / "Layouts"
+
+
 def waypoint_debug_directory() -> Path:
     return debug_dir() / "Waypoints"