dev/src/clew/clew.py

from __future__ import annotations

import argparse
import logging
import os
import shutil
import subprocess
import sys
from collections import deque
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from rpycdec.rpa import extract_rpa
from src.clew import extract_renpy_ast as extract_renpy_ast_script

EXTRACTOR = Path(extract_renpy_ast_script.__file__).resolve()
LOG = logging.getLogger(__name__)

def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="clew",
        description="Extract Ren'Py game files.",
    )
    parser.add_argument(
        "--game-root",
        required=True,
        type=Path,
        help="Game installation root to scan.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path("output"),
        help="Output root. Defaults to repository output/.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Replace Clew-managed output directories.",
    )
    parser.add_argument(
        "--no-extract-renpy-ast",
        action="store_true",
        help="Skip Ren'Py AST extraction.",
    )
    return parser.parse_args(argv)


def resolve_output_root(output_root: Path) -> Path:
    if output_root.is_absolute():
        return output_root.resolve()

    workspace = os.environ.get("BUILD_WORKSPACE_DIRECTORY")
    if workspace:
        return (Path(workspace) / output_root).resolve()

    return output_root.resolve()


def get_output_directory(output_root: Path, folder: str) -> Path:
    directory = output_root / folder
    directory.mkdir(parents=True, exist_ok=True)
    return directory


def extract_archive(archive: Path, destination: Path) -> None:
    destination.mkdir(parents=True, exist_ok=True)
    with archive.open("rb") as handle:
        extract_rpa(handle, output_dir=str(destination))


def extract_game_files(game_root: Path, output_path: Path) -> None:
    queue = deque(game_root.iterdir())

    while queue:
        item = queue.popleft()
        destination = output_path / item.relative_to(game_root)

        if item.is_dir():
            destination.mkdir(parents=True, exist_ok=True)
            queue.extend(item.iterdir())
        elif item.is_file() and item.suffix.lower() == ".rpa":
            extract_archive(item, destination.with_suffix(""))
        elif item.is_file():
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(item, destination)


def find(root: Path, predicate: Callable[[Path], bool]) -> list[Path]:
    matches = []
    for current, _directories, files in os.walk(root):
        current_path = Path(current)
        if predicate(current_path):
            matches.append(current_path)
        for file_name in files:
            file_path = current_path / file_name
            if predicate(file_path):
                matches.append(file_path)
    return sorted(matches)

def is_game_root(path: Path) -> bool:
    return (
        (path / "game").is_dir()
        and (path / "renpy").is_dir()
        and (path / "lib" / "linux-x86_64" / "python").is_file())


def is_renpy_script(path: Path) -> bool:
    return path.is_file() and path.suffix.lower() in {".rpyc", ".rpymc"}


def run_extract_renpy_ast(
    bundled_python: Path,
    game_root: Path,
    compiled_script: Path,
    output: Path,
) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    command = [
        "steam-run",
        str(bundled_python),
        str(EXTRACTOR),
        "--game-root",
        str(game_root),
        "--input",
        str(compiled_script),
        "--output",
        str(output),
    ]
    subprocess.run(command, check=True)


def extract_game_ast(extracted_root: Path, ast_root: Path, overwrite: bool) -> None:
    game_roots = find(extracted_root, is_game_root)
    if not game_roots:
        LOG.warning(f"no Ren'Py game roots found under {extracted_root}")
        return

    tasks = [
        (game_root, script_path)
        for game_root in game_roots
        for script_path in find(game_root / "game", is_renpy_script)
    ]

    with ThreadPoolExecutor() as executor:
        futures = {}
        for game_root, script_path in tasks:
            output_path = (
                ast_root
                / game_root.relative_to(extracted_root)
                / script_path.relative_to(game_root).with_suffix(script_path.suffix + ".json")
            )
            if output_path.exists() and not overwrite:
                LOG.info(f"skipping existing AST extraction: {output_path}")
                continue

            future = executor.submit(
                run_extract_renpy_ast,
                game_root / "lib" / "linux-x86_64" / "python",
                game_root,
                script_path,
                output_path,
            )
            futures[future] = script_path

        for future in as_completed(futures):
            future.result()
            LOG.info(f"finished AST extraction: {futures[future]}")


def run(args: argparse.Namespace) -> None:
    game_root = args.game_root.resolve()
    output_root = resolve_output_root(args.output_root)

    if not game_root.is_dir():
        LOG.critical(f"game root is not a directory: {game_root}")
        raise SystemExit(2)

    extracted = get_output_directory(output_root, "extracted")
    extract_game_files(game_root, extracted)

    if not args.no_extract_renpy_ast:
        ast = get_output_directory(output_root, "ast")
        extract_game_ast(extracted, ast, args.overwrite)


if __name__ == "__main__":
    args = parse_args(sys.argv[1:])

    logging.basicConfig(
        level=logging.INFO,
        format="[%(asctime)s] %(levelname)s %(filename)s:%(lineno)d | %(message)s",
        datefmt="%Y-%m-%d][%H:%M:%S",
    )

    run(args)