#1813: Add VPKs to the highest DLC folder, not dlc3 always.

src/packages/style_vpk.py

@@ -4,14 +4,17 @@
 """
 from __future__ import annotations
 
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, Iterator, Optional, TYPE_CHECKING
 from typing_extensions import Self
+
+from pathlib import Path
 import re
 import os
 import shutil
 
 from srctools import FileSystem, VPK
 import srctools.logger
+import trio
 
 import utils
 from packages import PakObject, ParseData, ExportData, NoVPKExport
@@ -21,6 +24,12 @@
 LOGGER = srctools.logger.get_logger(__name__, alias='packages.styleVPK')
 
 
+MARKER_FILENAME = 'bee2_vpk_autogen_marker.txt'
+MARKER_CONTENTS = b"""\
+This file marks the VPK as being autogenerated. 
+If you want to add files, see the vpk_override/ folder.
+"""
+
 VPK_OVERRIDE_README = """\
 Files in this folder will be written to the VPK during every BEE2 export.
 Use to override resources as you please.
@@ -29,29 +38,84 @@
 """
 
 
-# The folder we want to copy our VPKs to.
-VPK_FOLDER = {
-    # The last DLC released by Valve - this is the one that we
-    # overwrite with a VPK file.
-    utils.STEAM_IDS['PORTAL2']: 'portal2_dlc3',
-    utils.STEAM_IDS['DEST_AP']: 'portal2_dlc3',
+def iter_dlcs(root: Path) -> Iterator[Path]:
+    """Yield all mounted folders in order."""
+    yield root / 'portal2'
+    for dlc in range(1, 100):
+        yield root / f'portal2_dlc{dlc}'
 
-    # This doesn't have VPK files, and is higher priority.
-    utils.STEAM_IDS['APERTURE TAG']: 'portal2',
-}
 
-
-def clear_files(game: Game, folder: str) -> None:
+async def find_folder(game: Game) -> Path:
+    """Figure out which folder to use for this game."""
+    if game.unmarked_dlc3_vpk:
+        # Special case. This game had the old BEE behaviour, where it blindly wrote to DLC3
+        # and didn't mark the VPK. Just get rid of the folder, let our logic continue.
+        old_folder = 'portal2' if game.steamID == utils.STEAM_IDS['APERTURE TAG'] else 'portal2_dlc3'
+        # Keep the file, in case it actually was important.
+        for old_vpk in Path(game.root, old_folder).glob('*.vpk'):
+            if old_vpk.name.startswith('pak01'):
+                LOGGER.info('Removing old unmarked VPK: {}', old_vpk)
+                try:
+                    old_vpk.rename(old_vpk.with_name(f'bee_backup_{old_vpk.stem}.vpk'))
+                except FileNotFoundError:
+                    pass  # Never exported, very good.
+        game.unmarked_dlc3_vpk = False
+        game.save()
+
+    # We need to check if these are our VPK.
+    potentials: list[Path] = []
+    # If all are not ours, put it here.
+    fallback: Path
+
+    for game_folder in iter_dlcs(Path(game.root)):
+        vpk_filename = game_folder / 'pak01_dir.vpk'
+        if vpk_filename.exists():
+            potentials.append(vpk_filename)
+        else:
+            fallback = vpk_filename
+            break
+    else:
+        LOGGER.warning('Ran out of DLC folders??')
+        raise NoVPKExport()
+
+    # What we want to do is find the lowest-priority/first VPK that is a BEE one.
+    # But parsing is quite slow, and we expect the first three (p2, dlc1, dlc2) to all be fails.
+    # So we run the jobs, then wait for each in turn. That way later ones if they finish first
+    # will be ready.
+    async def worker(filename: Path, event: trio.Event) -> None:
+        """Parsing VPKs are expensive, do multiple concurrently."""
+        vpk = await trio.to_thread.run_sync(VPK, filename, cancellable=True)
+        results[filename] = MARKER_FILENAME in vpk
+        event.set()
+
+    results: dict[Path, bool | None] = dict.fromkeys(potentials, None)
+    events = [trio.Event() for _ in potentials]
+
+    event: trio.Event
+    filename: Path
+    async with trio.open_nursery() as nursery:
+        for filename, event in zip(potentials, events):
+            nursery.start_soon(worker, filename, event)
+        for filename, event in zip(potentials, events):
+            await event.wait()
+            if results[filename]:
+                LOGGER.info('Found BEE vpk: {}', filename)
+                return filename
+    LOGGER.info('No BEE vpk found, writing to: {}', fallback)
+    return fallback
+
+
+def clear_files(folder: Path) -> None:
     """Remove existing VPK files from the specified game folder.
 
      We want to leave other files - otherwise users will end up
      regenerating the sound cache every time they export.
     """
     os.makedirs(folder, exist_ok=True)
     try:
-        for file in os.listdir(folder):
-            if file[:6] == 'pak01_':
-                os.remove(os.path.join(folder, file))
+        for file in folder.iterdir():
+            if file.suffix == '.vpk' and file.stem.startswith('pak01_'):
+                file.unlink()
     except PermissionError:
         # The player might have Portal 2 open. Abort changing the VPK.
         LOGGER.warning("Couldn't replace VPK files. Is Portal 2 or Hammer open?")
@@ -99,36 +163,34 @@ async def export(exp_data: ExportData) -> None:
         else:
             sel_vpk = None
 
-        dest_folder = exp_data.game.abs_path(VPK_FOLDER.get(exp_data.game.steamID, 'portal2_dlc3'))
+        vpk_filename = await find_folder(exp_data.game)
+        LOGGER.info('VPK to write: {}', vpk_filename)
         try:
-            clear_files(exp_data.game, dest_folder)
+            clear_files(vpk_filename.parent)
         except PermissionError as exc:
             raise NoVPKExport() from exc  # We can't edit the VPK files - P2 is open..
 
-        if exp_data.game.steamID == utils.STEAM_IDS['PORTAL2']:
-            # In Portal 2, we make a dlc3 folder - this changes priorities,
-            # so the soundcache will be regenerated. Just copy the old one over.
-            sound_cache = os.path.join(
-                dest_folder, 'maps', 'soundcache', '_master.cache'
-            )
-            LOGGER.info('Sound cache: {}', sound_cache)
-            if not os.path.isfile(sound_cache):
-                LOGGER.info('Copying over soundcache file for DLC3..')
-                os.makedirs(os.path.dirname(sound_cache), exist_ok=True)
-                try:
-                    shutil.copy(
-                        exp_data.game.abs_path(
-                            'portal2_dlc2/maps/soundcache/_master.cache',
-                        ),
-                        sound_cache,
-                    )
-                except FileNotFoundError:
-                    # It's fine, this will be regenerated automatically
-                    pass
+        # When we make a DLC folder, this changes priorities,
+        # so the soundcache will be regenerated. Just copy the old one over.
+        sound_cache = Path(vpk_filename, '..', 'maps', 'soundcache', '_master.cache').resolve()
+        LOGGER.info('Sound cache: {}', sound_cache)
+        if not sound_cache.exists():
+            LOGGER.info('Copying over soundcache file for VPK folder..')
+            sound_cache.parent.mkdir(parents=True, exist_ok=True)
+            try:
+                shutil.copy(
+                    exp_data.game.abs_path('portal2_dlc2/maps/soundcache/_master.cache'),
+                    sound_cache,
+                )
+            except FileNotFoundError:
+                # It's fine, this will be regenerated automatically.
+                pass
 
         # Generate the VPK.
-        vpk_file = VPK(os.path.join(dest_folder, 'pak01_dir.vpk'), mode='w')
+        vpk_file = VPK(vpk_filename, mode='w')
         with vpk_file:
+            # Write the marker, so we can identify this later. Always put it in the _dir.vpk.
+            vpk_file.add_file(MARKER_FILENAME, MARKER_CONTENTS, arch_index=None)
             if sel_vpk is not None:
                 for file in sel_vpk.fsys.walk_folder(sel_vpk.dir):
                     with file.open_bin() as open_file:
@@ -144,8 +206,8 @@ async def export(exp_data: ExportData) -> None:
             override_folder = exp_data.game.abs_path('vpk_override')
             os.makedirs(override_folder, exist_ok=True)
 
-            # Also write a file to explain what it's for..
-            with open(os.path.join(override_folder, 'BEE2_README.txt'), 'w') as f:
+            # Also write a file to explain what it's for...
+            with open(os.path.join(override_folder, 'BEE2_README.txt'), 'w', encoding='utf8') as f:
                 f.write(VPK_OVERRIDE_README)
 
             # Matches pak01_038.vpk, etc. These shouldn't be opened.