Atlas-Install/source/mkw/ExtractedGame.py

import shutil
from io import BytesIO
from pathlib import Path
from typing import Generator, IO, TYPE_CHECKING

from source.mkw.ModConfig import ModConfig
from source.mkw.Patch.Patch import Patch
from source.progress import Progress
from source.wt import szs, lec, wit
from source.wt.wstrt import StrPath

if TYPE_CHECKING:
    from source.mkw.Game import Game


class PathOutsideMod(Exception):
    def __init__(self, forbidden_path: Path, allowed_range: Path):
        super().__init__(f"Error : path {forbidden_path} outside of allowed range {allowed_range}")


class ExtractedGame:
    """
    Class that represents an extracted game
    """

    def __init__(self, path: "Path | str", original_game: "Game" = None):
        self.path = Path(path)
        self.original_game = original_game
        self._special_file: dict[str, IO] = {}

    def extract_autoadd(self, destination_path: "Path | str") -> Generator[Progress, None, None]:
        """
        Extract all the autoadd files from the game to destination_path
        :param destination_path: directory where the autoadd files will be extracted
        """
        yield Progress(description="Extracting autoadd files...", determinate=False)
        szs.autoadd(self.path / "files/Race/Course/", destination_path)

    def extract_original_tracks(self, destination_path: "Path | str") -> Generator[Progress, None, None]:
        """
        Move all the original tracks to the destination path
        :param destination_path: destination of the track
        """
        destination_path = Path(destination_path)
        destination_path.mkdir(parents=True, exist_ok=True)
        yield Progress(description="Extracting original tracks...", determinate=False)
        for track_file in (self.path / "files/Race/Course/").glob("*.szs"):
            yield Progress(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, mystuff_path: "Path | str") -> Generator[Progress, None, None]:
        """
        Install mystuff directory. If any files of the game have the same name as a file at the root of the MyStuff
        Patch, then it is copied.
        :mystuff_path: path to the MyStuff directory
        :return:
        """
        yield Progress(description=f"Installing MyStuff '{mystuff_path}'...", determinate=False)
        mystuff_path = Path(mystuff_path)

        mystuff_rootfiles: dict[str, Path] = {}
        for mystuff_subpath in mystuff_path.glob("*"):
            if mystuff_subpath.is_file(): mystuff_rootfiles[mystuff_subpath.name] = mystuff_subpath
            else: shutil.copytree(mystuff_subpath, self.path / f"files/{mystuff_subpath.name}", dirs_exist_ok=True)

        for game_file in filter(lambda file: file.is_file(), (self.path / "files/").rglob("*")):
            if (mystuff_file := mystuff_rootfiles.get(game_file.name)) is None: continue
            shutil.copy(mystuff_file, game_file)

    def install_multiple_mystuff(self, mystuff_paths: list["Path | str"]) -> Generator[Progress, None, None]:
        """
        Install multiple mystuff patch
        :param mystuff_paths: paths to all the mystuff patch
        """
        yield Progress(description="Installing all the mystuff patchs")

        for mystuff_path in mystuff_paths:
            yield from self.install_mystuff(mystuff_path)

    def prepare_special_file(self, mod_config: ModConfig) -> Generator[Progress, None, None]:
        """
        Prepare special files for the patch
        :return: the special files dict
        """
        yield Progress(description="Preparing ct_icon special file...", determinate=False)
        ct_icons = BytesIO()
        mod_config.get_full_cticon().save(ct_icons, format="PNG")
        ct_icons.seek(0)
        self._special_file["ct_icons"] = ct_icons

    def prepare_dol(self) -> Generator[Progress, None, None]:
        """
        Prepare main.dol and StaticR.rel files (clean them and add lecode)
        """
        yield Progress(description="Preparing main.dol...", determinate=False)
        StrPath(self.path / "sys/main.dol").patch(clean_dol=True, add_lecode=True)

    def recreate_all_szs(self) -> Generator[Progress, None, None]:
        """
        Repack all the .d directory into .szs files.
        """
        yield Progress(description=f"Repacking all szs", determinate=False)

        for extracted_szs in filter(lambda path: path.is_dir(), self.path.rglob("*.d")):
            # for every directory that end with a .d in the extracted game, recreate the szs
            yield Progress(description=f"Repacking {extracted_szs} to szs", determinate=False)

            szs.create(extracted_szs, extracted_szs.with_suffix(".szs"), overwrite=True)
            shutil.rmtree(str(extracted_szs.resolve()))

    def patch_lecode(self, mod_config: ModConfig, cache_directory: Path | str,
                     cttracks_directory: Path | str, ogtracks_directory: Path | str) -> Generator[Progress, None, None]:
        """
        install lecode on the mod
        :param cttracks_directory: directory to the customs tracks
        :param ogtracks_directory: directory to the originals tracks
        :param cache_directory: Path to the cache
        :param mod_config: mod configuration
        """
        yield Progress(description="Patching LECODE.bin")
        cache_directory = Path(cache_directory)
        cttracks_directory = Path(cttracks_directory)
        ogtracks_directory = Path(ogtracks_directory)

        # write ctfile data to a file in the cache
        ct_file = (cache_directory / "ctfile.lpar")
        ct_file.write_text(mod_config.get_ctfile(template="-"))

        lpar_dir: Path = mod_config.path.parent / "_LPAR/"
        lpar: Path = lpar_dir / mod_config.multiple_safe_eval(mod_config.lpar_template)()
        if not lpar.is_relative_to(lpar_dir): raise PathOutsideMod(lpar, lpar_dir)

        for lecode_file in (self.path / "files/rel/").glob("lecode-*.bin"):
            lec.patch(
                lecode_file=lecode_file,
                ct_file=ct_file,
                lpar=lpar,
                game_tracks_directory=self.path / "files/Race/Course/",
                copy_tracks_directories=[ogtracks_directory, cttracks_directory]
            )

    def _install_all_patch(self, mod_config: ModConfig, patch_directory_name: str) -> Generator[Progress, None, None]:
        """
        for all directory that are in the root of the mod, and don't start with an underscore,
        for all the subdirectory named by the patch_directory_name, apply the patch
        :param mod_config: the mod to install
        """
        # yield an empty dict so that if nothing is yielded by the Patch, still is considered a generator
        yield Progress()

        for part_directory in mod_config.get_mod_directory().glob("[!_]*"):
            for patch_directory in part_directory.glob(patch_directory_name):
                yield from Patch(patch_directory, mod_config, self._special_file).install(self)

    def install_all_prepatch(self, mod_config: ModConfig) -> Generator[Progress, None, None]:
        """
        Install all patchs of the mod_config into the game.
        Used before the lecode patch is applied
        :param mod_config: the mod to install
        """
        yield Progress(description="Installing all Pre-Patch...", determinate=False)
        yield from self._install_all_patch(mod_config, "_PREPATCH/")

    def install_all_patch(self, mod_config: ModConfig) -> Generator[Progress, None, None]:
        """
        Install all patchs of the mod_config into the game.
        Used after the lecode patch is applied
        :param mod_config: the mod to install
        """
        yield Progress(description="Installing all Patch...", determinate=False)
        yield from self._install_all_patch(mod_config, "_PATCH/")

    def convert_to(self, output_type: wit.Extension) -> Generator[Progress, None, wit.WITPath | None]:
        """
        Convert the extracted game to another format
        :param output_type: path to the destination of the game
        :output_type: format of the destination game
        """
        yield Progress(description=f"Converting game to {output_type}", determinate=False)
        if output_type == wit.Extension.FST: return

        destination_file = self.path.with_suffix(self.path.suffix + output_type.value)
        dest_stem: str = destination_file.stem

        i: int = 0
        while destination_file.exists():
            i += 1
            destination_file = destination_file.with_stem(dest_stem + f" ({i})")

        converted_game: wit.WITPath = wit.copy(
            source_directory=self.path,
            destination_file=destination_file,
        )

        yield Progress(description="Deleting the extracted game...", determinate=False)
        shutil.rmtree(self.path)

        return converted_game