Mist/backend/src/mist/core/chain_replay.py

"""Cold reconstruction of arbitrary historical versions via delta-chain replay.

Ported from: JoshSteam CDN/Utils/prepare_game_version.py
Bug fixes / changes during port:
    - All hardcoded paths (GAMES_DIR/CACHE_DIR/TEMP_DIR/hdiff_path) replaced with
      `mist.core.paths` helpers + `settings`
    - Type hints
    - Behavior preserved: find closest cached version, copy it, walk forward
      applying hdiff patches at each step, verify final state, cache result

The key insight: we never need to store every version's full files. The base
version (.7z) + the delta chain gives us any historical version on demand,
and the result gets cached for next time.
"""

from __future__ import annotations

import concurrent.futures
import json
import os
import shutil
import subprocess
from pathlib import Path
from typing import Any

from mist.config import settings
from mist.core import paths
from mist.core.compression import extract_7z
from mist.core.manifest import verify_files


# ---- Version-history helpers (legacy JSON path; callers should prefer the DB) ----


def get_base_version(game_title: str) -> str:
    path = paths.game_version_history_path(game_title)
    if not path.exists():
        raise FileNotFoundError(f"Version file for {game_title} not found at {path}")
    with open(path) as f:
        versions: list[str] = json.load(f)
    if not versions:
        raise ValueError(f"No versions found for {game_title}")
    return versions[0]


def get_versions_up_to(game_title: str, target_version: str) -> list[str]:
    path = paths.game_version_history_path(game_title)
    if not path.exists():
        raise FileNotFoundError(f"Version file for {game_title} not found at {path}")
    with open(path) as f:
        versions: list[str] = json.load(f)
    try:
        idx = versions.index(target_version)
    except ValueError:
        return []
    return versions[0 : idx + 1]


def get_versions_in_range(game_title: str, from_version: str, to_version: str) -> list[str]:
    path = paths.game_version_history_path(game_title)
    if not path.exists():
        raise FileNotFoundError(f"Version file for {game_title} not found at {path}")
    with open(path) as f:
        versions: list[str] = json.load(f)
    try:
        from_idx = versions.index(from_version)
        to_idx = versions.index(to_version)
    except ValueError:
        return []
    return versions[from_idx + 1 : to_idx + 1]


# ---- Cache lookup ----


def find_game_folder_in_cache(game_title: str, version: str) -> bool:
    """Is the reconstructed-version folder for (title, version) present in the cache?"""
    return paths.cached_game_version_dir(game_title, version).exists()


def find_closest_cached_version(game_title: str, version_list: list[str]) -> str | None:
    """Of the candidate versions in `version_list`, return the latest one we've already
    reconstructed and have sitting in cache (if any)."""
    cache_dir = settings.cache_dir
    cached_versions: list[str] = []
    if not cache_dir.exists():
        return None
    for folder in os.listdir(cache_dir):
        if not folder.startswith(game_title):
            continue
        folder_version = folder[len(game_title) :].strip()
        if folder_version in version_list:
            cached_versions.append(folder_version)
    if not cached_versions:
        return None
    cached_versions.sort(key=lambda v: version_list.index(v))
    return cached_versions[-1]


# ---- Patch application ----


def _copy_folder(src: Path, dest: Path) -> None:
    if not src.exists():
        raise FileNotFoundError(f"Source folder not found: {src}")
    shutil.copytree(src, dest)


def _apply_hdiff_patch(
    patch: str, data: dict[str, str], temp_delta_folder: Path, temp_game_dir: Path
) -> bool:
    patch_path = temp_delta_folder / patch
    relative_path = data.get("relative_path")
    if not relative_path:
        return False
    original_file_path = temp_game_dir / relative_path
    if not original_file_path.exists():
        return False
    temp_file_path = temp_delta_folder / f"{patch}.tmp"
    try:
        subprocess.run(
            [str(settings.hpatchz_path), "-f", str(original_file_path), str(patch_path), str(temp_file_path)],
            check=True,
        )
        os.replace(temp_file_path, original_file_path)
        return True
    except (subprocess.CalledProcessError, OSError):
        return False


def apply_update(game_title: str, version: str, temp_game_dir: Path) -> None:
    """Apply the delta_manifest for `version` to the working folder `temp_game_dir`.

    Copies the delta folder out of the NAS into temp_delta_folder, applies in parallel,
    handles new/deleted files, then deletes the temp_delta_folder.
    """
    delta_folder = paths.deltas_dir(game_title, version)
    temp_delta_folder = settings.temp_dir / f"{game_title} Update" / version
    if not delta_folder.exists():
        return

    _copy_folder(delta_folder, temp_delta_folder)
    with open(temp_delta_folder / "delta_manifest.json") as f:
        delta_manifest: dict[str, Any] = json.load(f)

    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
        futures = [
            executor.submit(_apply_hdiff_patch, patch, data, temp_delta_folder, temp_game_dir)
            for patch, data in delta_manifest["modified_files"].items()
        ]
        for f in concurrent.futures.as_completed(futures):
            f.result()

    new_files_folder = temp_delta_folder / "new_files"
    for new_file in delta_manifest["new_files"]:
        src = new_files_folder / new_file["file"]
        dst = temp_game_dir / new_file["relative_path"]
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(src, dst)

    for deleted in delta_manifest["deleted_files"]:
        target = temp_game_dir / deleted["relative_path"]
        if target.exists():
            target.unlink()

    shutil.rmtree(temp_delta_folder)


# ---- Orchestrator ----


def prepare_game_version(game_title: str, version: str) -> bool:
    """Reconstruct (game_title, version) into the cache folder and verify it.

    Strategy:
        1. If version == base, copy/extract the base archive.
        2. Otherwise: find the closest cached version, copy it, walk forward.
        3. If nothing is cached, extract the base and walk forward from there.
        4. Verify against the per-version manifest. On success, move to cache.
    """
    base_version = get_base_version(game_title)
    update_path: list[str] | None = None
    temp_game_dir: Path

    if version == base_version:
        if find_game_folder_in_cache(game_title, base_version):
            _copy_folder(
                paths.cached_game_version_dir(game_title, base_version),
                settings.temp_dir / f"{game_title} {version}",
            )
            temp_game_dir = settings.temp_dir / f"{game_title} {version}"
        else:
            base_archive = paths.base_archive_path(game_title)
            temp_game_dir = settings.temp_dir / f"{game_title} {base_version}"
            if not base_archive.exists():
                return False
            temp_game_dir.mkdir(parents=True, exist_ok=True)
            extract_7z(base_archive, temp_game_dir)
    else:
        update_path_from_base = get_versions_up_to(game_title, version)
        closest = find_closest_cached_version(game_title, update_path_from_base)
        if closest:
            _copy_folder(
                paths.cached_game_version_dir(game_title, closest),
                settings.temp_dir / f"{game_title} {closest}",
            )
            temp_game_dir = settings.temp_dir / f"{game_title} {closest}"
            update_path = get_versions_in_range(game_title, closest, version)
        else:
            base_archive = paths.base_archive_path(game_title)
            temp_game_dir = settings.temp_dir / f"{game_title} {base_version}"
            if not base_archive.exists():
                return False
            temp_game_dir.mkdir(parents=True, exist_ok=True)
            extract_7z(base_archive, temp_game_dir)
            update_path = get_versions_up_to(game_title, version)
            if update_path:
                update_path.pop(0)  # we already have the base extracted

    if update_path:
        for next_version in update_path:
            apply_update(game_title, next_version, temp_game_dir)
            new_folder = settings.temp_dir / f"{game_title} {next_version}"
            os.rename(temp_game_dir, new_folder)
            temp_game_dir = new_folder
        update_workspace = settings.temp_dir / f"{game_title} Update"
        if update_workspace.exists():
            shutil.rmtree(update_workspace)

    requested_manifest = paths.manifest_path(game_title, version)
    results = verify_files(temp_game_dir, requested_manifest)
    if results["all_files_verified"]:
        output_dir = paths.cached_game_version_dir(game_title, version)
        shutil.move(str(temp_game_dir), str(output_dir))
        return True
    shutil.rmtree(temp_game_dir)
    return False