From 186aa55f4a13bd76ff6c56779a65dae10af966aa Mon Sep 17 00:00:00 2001
From: Faraphel <rc60650@hotmail.com>
Date: Mon, 18 Jul 2022 19:11:11 +0200
Subject: [PATCH] custom tracks are now converted into szs in the cache

---
 source/mkw/ExtractedGame.py |  4 +++-
 source/mkw/Game.py          |  9 +++++++--
 source/mkw/ModConfig.py     | 18 ++++++++++++++++++
 source/wt/szs.py            | 21 +++++++++++++++++++++
 4 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/source/mkw/ExtractedGame.py b/source/mkw/ExtractedGame.py
index 75bd45e..5bfa035 100644
--- a/source/mkw/ExtractedGame.py
+++ b/source/mkw/ExtractedGame.py
@@ -36,7 +36,9 @@ class ExtractedGame:
         destination_path.mkdir(parents=True, exist_ok=True)
         yield {"description": "Extracting original tracks...", "determinate": False}
         for track_file in (self.path / "files/Race/Course/").glob("*.szs"):
-            track_file.rename(destination_path / track_file.name)
+            yield {"description": f"Extracting original tracks ({track_file.name})...", "determinate": False}
+            if not (destination_path / track_file.name).exists(): track_file.rename(destination_path / track_file.name)
+            else: track_file.unlink()
 
     def install_mystuff(self) -> Generator[dict, None, None]:
         """
diff --git a/source/mkw/Game.py b/source/mkw/Game.py
index 1c2417d..726ed74 100644
--- a/source/mkw/Game.py
+++ b/source/mkw/Game.py
@@ -98,6 +98,10 @@ class Game:
         cache_directory: Path = Path("./.cache")
         cache_directory.mkdir(parents=True, exist_ok=True)
 
+        cache_autoadd_directory = cache_directory / "autoadd/"
+        cache_ogtracks_directory = cache_directory / "original-tracks/"
+        cache_cttracks_directory = cache_directory / f"custom-tracks/"
+
         # get the directory where the game will be extracted
         extracted_game = ExtractedGame(self.get_output_directory(dest, mod_config), self)
 
@@ -108,8 +112,9 @@ class Game:
         yield from self.extract(extracted_game.path)
 
         # prepare the cache
-        yield from extracted_game.extract_autoadd(cache_directory / "autoadd/")
-        yield from extracted_game.extract_original_tracks(cache_directory / "original-tracks/")
+        yield from extracted_game.extract_autoadd(cache_autoadd_directory)
+        yield from extracted_game.extract_original_tracks(cache_ogtracks_directory)
+        yield from mod_config.normalize_all_tracks(cache_autoadd_directory, cache_cttracks_directory)
 
         # patch the game
         yield from extracted_game.install_mystuff()
diff --git a/source/mkw/ModConfig.py b/source/mkw/ModConfig.py
index 593814e..9be0b9f 100644
--- a/source/mkw/ModConfig.py
+++ b/source/mkw/ModConfig.py
@@ -8,6 +8,7 @@ from source.mkw.Cup import Cup
 from source.mkw.Track import Track
 import json
 
+from source.wt.szs import SZSPath
 
 CT_ICON_SIZE: int = 128
 
@@ -242,3 +243,20 @@ class ModConfig:
         for i, cticon in enumerate(cticons): full_cticon.paste(cticon, (0, i * CT_ICON_SIZE))
 
         return full_cticon
+
+    def normalize_all_tracks(self, autoadd_path: "Path | str", destination_path: "Path | str") -> Generator[dict, None, None]:
+        """
+        Convert all tracks of the mod to szs into the destination_path
+        :param autoadd_path: autoadd directory
+        :param destination_path: destination where the files are converted
+        """
+        yield {"description": "Normalizing track..."}
+        destination_path = Path(destination_path)
+        destination_path.mkdir(parents=True, exist_ok=True)
+        for track_file in filter(lambda file: file.is_file(), (self.path.parent / "_TRACKS").rglob("*")):
+            yield {"description": f"Normalizing track \"{track_file.name}\"..."}
+            SZSPath(track_file).normalize(
+                autoadd_path,
+                destination_path / track_file.with_suffix(".szs").name,
+                format="szs"
+            )
diff --git a/source/wt/szs.py b/source/wt/szs.py
index cd7fd50..d884f9b 100644
--- a/source/wt/szs.py
+++ b/source/wt/szs.py
@@ -53,6 +53,27 @@ class SZSPath:
     def __eq__(self, other: "SZSPath") -> bool:
         return self.path == other.path
 
+    def normalize(self, autoadd_path: "Path | str", destination_path: "Path | str", format: str = "szs") -> "SZSPath":
+        """
+        Convert the file into a another format
+        :param format: format to convert the file
+        :param autoadd_path: Autoadd directory
+        :param destination_path: destination of the converted file
+        :return: new path of the file
+        """
+        if not destination_path.exists() or \
+                (destination_path.exists() and destination_path.stat().st_mtime < self.path.stat().st_mtime):
+            # if the destination_path exists and is less recent than this source file, update it.
+            _tools_run(
+                "NORMALIZE",
+                self.path,
+                "--autoadd-path", autoadd_path,
+                "--DEST", destination_path,
+                f"--{format}",
+                "--overwrite"
+            )
+        return SZSPath(destination_path)
+
     def cat(self, subfile: str) -> bytes:
         """
         Run the cat command (read a subfile) and return the output