pokemon/scripts/infinity-money-trainer.py

#!/usr/bin/env python3
"""Lightweight money editor for Pokémon Infinity (Essentials / RPG Maker XP).

Edits @money on PokeBattle_Trainer in Game_*.rxdata save slots under:
  %USERPROFILE%\\Saved Games\\Pokémon Infinity

Close the game before running. A timestamped .bak backup is created first.
"""

from __future__ import annotations

import argparse
import glob
import os
import shutil
import sys
from datetime import datetime
from pathlib import Path

try:
    from rubymarshal.reader import loads
    from rubymarshal.writer import writes
except ImportError:
    print("Missing dependency: pip install rubymarshal", file=sys.stderr)
    sys.exit(1)

DEFAULT_SAVE_DIR = Path.home() / "Saved Games" / "Pokémon Infinity"
ALT_SAVE_DIR = Path.home() / "Saved Games" / "Pokmon Infinity"  # typo seen in ludusavi manifest
DEFAULT_MONEY = 999_999
MAX_MONEY = 9_999_999


def resolve_save_dir(explicit: str | None) -> Path:
    if explicit:
        path = Path(explicit)
        if not path.is_dir():
            raise SystemExit(f"Save directory not found: {path}")
        return path
    for candidate in (DEFAULT_SAVE_DIR, ALT_SAVE_DIR):
        if candidate.is_dir():
            return candidate
    raise SystemExit(
        "Could not find Pokémon Infinity saves. Expected one of:\n"
        f"  {DEFAULT_SAVE_DIR}\n"
        f"  {ALT_SAVE_DIR}\n"
        "Use --save-dir if yours is elsewhere."
    )


def list_save_files(save_dir: Path) -> list[Path]:
    files = sorted(save_dir.glob("Game*.rxdata"))
    return [f for f in files if f.name != "Settings.rxdata"]


def read_money(path: Path) -> int:
    with path.open("rb") as fh:
        trainer = loads(fh.read())
    if getattr(trainer, "ruby_class_name", None) != "PokeBattle_Trainer":
        raise ValueError(f"{path.name} is not a trainer save (unexpected format)")
    if "@money" not in trainer.attributes:
        raise ValueError(f"{path.name} has no @money field")
    return int(trainer.attributes["@money"])


def write_money(path: Path, amount: int) -> None:
    if amount < 0 or amount > MAX_MONEY:
        raise ValueError(f"Money must be between 0 and {MAX_MONEY:,}")

    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = path.with_suffix(path.suffix + f".bak.{stamp}")
    shutil.copy2(path, backup)

    with path.open("rb") as fh:
        trainer = loads(fh.read())
    trainer.attributes["@money"] = amount
    with path.open("wb") as fh:
        fh.write(writes(trainer))

    print(f"  backup -> {backup.name}")


def cmd_show(save_dir: Path) -> None:
    files = list_save_files(save_dir)
    if not files:
        raise SystemExit(f"No Game*.rxdata files in {save_dir}")
    print(f"Save folder: {save_dir}\n")
    for path in files:
        try:
            money = read_money(path)
            print(f"  {path.name:16} ${money:,}")
        except Exception as exc:  # noqa: BLE001 - surface per-file issues
            print(f"  {path.name:16} (unreadable: {exc})")


def cmd_set(save_dir: Path, amount: int, slot: str | None) -> None:
    files = list_save_files(save_dir)
    if not files:
        raise SystemExit(f"No Game*.rxdata files in {save_dir}")

    if slot:
        matches = [f for f in files if f.stem.lower() == slot.lower() or f.name.lower() == slot.lower()]
        if not matches:
            matches = [f for f in files if slot in f.name]
        if not matches:
            raise SystemExit(f"No save matching slot {slot!r}. Available: {[f.name for f in files]}")
        targets = matches
    else:
        targets = files

    print(f"Save folder: {save_dir}")
    for path in targets:
        old = read_money(path)
        write_money(path, amount)
        print(f"  {path.name}: ${old:,} -> ${amount:,}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Set money in Pokémon Infinity save files (close the game first)."
    )
    parser.add_argument(
        "--save-dir",
        help="Override save directory (default: %%USERPROFILE%%\\Saved Games\\Pokémon Infinity)",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("show", help="Show money in each save slot")

    set_p = sub.add_parser("set", help="Set money in one or all slots")
    set_p.add_argument(
        "amount",
        nargs="?",
        type=int,
        default=DEFAULT_MONEY,
        help=f"Target money (default: {DEFAULT_MONEY:,})",
    )
    set_p.add_argument(
        "--slot",
        help="Only edit this slot (e.g. Game_1 or Game_1.rxdata). Default: all Game*.rxdata",
    )
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()
    save_dir = resolve_save_dir(args.save_dir)

    if args.command == "show":
        cmd_show(save_dir)
    elif args.command == "set":
        cmd_set(save_dir, args.amount, args.slot)


if __name__ == "__main__":
    main()