Atlas-Install/source/mkw/ExtractedGame.py

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

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


Game: any


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[dict, 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 {"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[dict, 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 {"description": "Extracting original tracks...", "determinate": False}
        for track_file in (self.path / "files/Race/Course/").glob("*.szs"):
            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, mystuff_path: "Path | str") -> Generator[dict, 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 {"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[dict, None, None]:
        """
        Install multiple mystuff patch
        :param mystuff_paths: paths to all the mystuff patch
        """
        yield {"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[dict, None, None]:
        """
        Prepare special files for the patch
        :return: the special files dict
        """
        yield {"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[dict, None, None]:
        """
        Prepare main.dol and StaticR.rel files (clean them and add lecode)
        """
        yield {"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[dict, None, None]:
        """
        Repack all the .d directory into .szs files.
        """
        yield {"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 {"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[dict, 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 {"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) -> Generator[dict, None, None]:
        """
        Install all patchs of the mod_config into the game
        :param mod_config: the mod to install
        """
        yield {"description": "Installing all Patch...", "determinate": False}

        # prepare special files data
        yield from self.prepare_special_file(mod_config)

        # for all directory that are in the root of the mod, and don't start with an underscore,
        # for all the subdirectory named "_PATCH", apply the patch
        for part_directory in mod_config.get_mod_directory().glob("[!_]*"):
            for patch_directory in part_directory.glob("_PATCH/"):
                yield from Patch(patch_directory, mod_config, self._special_file).install(self)