#!/usr/bin/env python3
"""Enable or disable Pokémon Infinity debug mode by patching Data/Scripts.rxdata."""

from __future__ import annotations

import argparse
import shutil
import sys
import zlib
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)

RGSS2_SCRIPT_INDEX = 4
DEBUG_FALSE = b"$DEBUG=false"
DEBUG_TRUE = b"$DEBUG=true"


def patch_scripts(path: Path, enable: bool) -> bool:
    with path.open("rb") as fh:
        scripts = loads(fh.read())

    entry = scripts[RGSS2_SCRIPT_INDEX]
    name = entry[1].decode() if isinstance(entry[1], bytes) else str(entry[1])
    if name != "RGSS2Compatibility":
        raise ValueError(f"Unexpected script at index {RGSS2_SCRIPT_INDEX}: {name!r}")

    source = zlib.decompress(entry[2])
    has_true = DEBUG_TRUE in source
    has_false = DEBUG_FALSE in source

    if enable:
        if has_true and not has_false:
            print("Debug mode already enabled.")
            return False
        if not has_false:
            raise ValueError("Could not find $DEBUG=false in RGSS2Compatibility")
        patched = source.replace(DEBUG_FALSE, DEBUG_TRUE, 1)
    else:
        if has_false and not has_true:
            print("Debug mode already disabled.")
            return False
        if not has_true:
            raise ValueError("Could not find $DEBUG=true in RGSS2Compatibility")
        patched = source.replace(DEBUG_TRUE, DEBUG_FALSE, 1)
    scripts[RGSS2_SCRIPT_INDEX] = [entry[0], entry[1], zlib.compress(patched, 9)]

    with path.open("wb") as fh:
        fh.write(writes(scripts))

    return True


def main() -> None:
    parser = argparse.ArgumentParser(description="Toggle Pokémon Infinity debug mode.")
    parser.add_argument(
        "game_dir",
        nargs="?",
        default=r"C:\Users\MattC\roms\windows\Pokemon Infinity",
        help="Path to the extracted game folder (contains Game.exe)",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--enable", action="store_true", default=True, help="Enable debug (default)")
    group.add_argument("--disable", action="store_true", help="Restore normal mode")
    args = parser.parse_args()

    game_dir = Path(args.game_dir)
    scripts_path = game_dir / "Data" / "Scripts.rxdata"
    if not (game_dir / "Game.exe").is_file():
        raise SystemExit(f"Game.exe not found under {game_dir}")
    if not scripts_path.is_file():
        raise SystemExit(f"Scripts.rxdata not found at {scripts_path}")

    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = scripts_path.with_suffix(f".rxdata.bak.{stamp}")
    shutil.copy2(scripts_path, backup)
    print(f"Backup: {backup}")

    changed = patch_scripts(scripts_path, enable=not args.disable)
    if changed:
        state = "enabled" if not args.disable else "disabled"
        print(f"Debug mode {state}.")
        if not args.disable:
            print()
            print("In-game: open the pause menu -> Debug -> Set Money / Add Item")
            print("Load screen also gets a Debug option. Back up saves before experimenting.")


if __name__ == "__main__":
    main()