A3Rebalancer/scripts/package_mod.py

"""
package_mod.py — Package output/ into a loadable Arma 3 mod.

Uses armake2 to pack each addon into a .pbo file.
Creates mod.cpp and installs to your Arma 3 addons directory.

Usage:
    python scripts/package_mod.py
    python scripts/package_mod.py --name "My Rebalance" --output "C:/Arma 3/!mods"
    python scripts/package_mod.py --dry-run

Settings are read from package_mod.ini in the repo root if it exists.
Command-line arguments override INI values.
"""

import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
OUTPUT_DIR = REPO_ROOT / "output"
INI_PATH = REPO_ROOT / "package_mod.ini"

DEFAULTS = {
    "mod_name": "A3Rebalancer",
    "mod_display_name": "A3 Rebalancer",
    "mod_description": "Normalizes non-RHS ammo and armor values to RHS baseline.",
    "mod_author": "REDCODE",
    "mod_version": "1.0",
    "arma_mods_dir": str(Path.home() / "Documents" / "Arma 3" / "!mods"),
}


def read_ini() -> dict:
    """Read package_mod.ini if it exists."""
    if not INI_PATH.exists():
        return {}
    config = {}
    current_section = None
    with open(INI_PATH, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or line.startswith(";"):
                continue
            if line.startswith("[") and line.endswith("]"):
                current_section = line[1:-1]
                continue
            if "=" in line:
                key, _, val = line.partition("=")
                key = key.strip()
                val = val.strip().strip('"').strip("'")
                if current_section:
                    key = f"{current_section}.{key}"
                config[key] = val
    return config


def generate_mod_cpp(display_name: str, description: str, author: str,
                     version: str) -> str:
    """Generate mod.cpp content."""
    desc_escaped = description.replace('"', '""')
    return f'''beta = 0;
action = "";
hideName = 0;
hidePicture = 0;
name = "{display_name}";
picture = "";
logo = "";
logoOver = "";
tooltip = "{desc_escaped}";
overview = "{desc_escaped}";
author = "{author}";
authorID = "";
version = "{version}";
versionArma[] = {{3, 0, 0}};
versionAI[] = {{3, 0, 0}};
versionExt[] = {{1, 0, 0}};
'''


def find_armake2() -> str:
    """Find armake2 binary."""
    # Check PATH first
    result = subprocess.run(
        ["where", "armake2"], capture_output=True, text=True)
    if result.returncode == 0:
        return result.stdout.strip().splitlines()[0]

    # Check tools/ directory
    tools_path = REPO_ROOT / "tools" / "armake2.exe"
    if tools_path.exists():
        return str(tools_path)

    print("ERROR: armake2 not found.")
    print("Install from: https://github.com/KoffeinFlummi/armake2")
    print("Or place armake2.exe in tools/")
    sys.exit(1)


def pack_pbo(armake2: str, source_dir: Path, target_pbo: Path,
             prefix: str) -> bool:
    """Pack a directory into a .pbo using armake2 build (with rapification)."""
    cmd = [
        armake2, "build", "-f",
        "-e", f"prefix={prefix}",
        str(source_dir),
        str(target_pbo),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"    ERROR packing {source_dir.name}: {result.stderr.strip()}")
        return False
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Package ArmADump output into an Arma 3 mod")
    parser.add_argument("--name", help="Mod folder name (default: A3Rebalancer)")
    parser.add_argument("--display-name", help="Mod display name")
    parser.add_argument("--description", help="Mod description")
    parser.add_argument("--author", help="Mod author")
    parser.add_argument("--version", help="Mod version")
    parser.add_argument("--output", help="Arma 3 mods directory")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show what would be done without building")
    args = parser.parse_args()

    # Read INI config
    ini = read_ini()

    # Resolve values: CLI > INI > defaults
    mod_name = args.name or ini.get("mod_name", DEFAULTS["mod_name"])
    display_name = (args.display_name or
                    ini.get("mod_display_name", DEFAULTS["mod_display_name"]))
    description = (args.description or
                   ini.get("mod_description", DEFAULTS["mod_description"]))
    author = args.author or ini.get("mod_author", DEFAULTS["mod_author"])
    version = args.version or ini.get("mod_version", DEFAULTS["mod_version"])
    mods_dir = args.output or ini.get("arma_mods_dir", DEFAULTS["arma_mods_dir"])

    print("=" * 60)
    print("ArmADump - Package Mod")
    print("=" * 60)
    print(f"  Mod name:      @{mod_name}")
    print(f"  Display name:  {display_name}")
    print(f"  Author:        {author}")
    print(f"  Version:       {version}")
    print(f"  Output:        {Path(mods_dir) / ('@' + mod_name)}")
    print()

    # Validate output dir exists
    if not OUTPUT_DIR.exists():
        print(f"ERROR: {OUTPUT_DIR} does not exist.")
        print("Run generate_patches.py first.")
        sys.exit(1)

    # Check for config.cpp files
    mod_dirs = [d for d in OUTPUT_DIR.iterdir()
                if d.is_dir() and (d / "config.cpp").exists()]
    if not mod_dirs:
        print(f"ERROR: No mod folders with config.cpp found in {OUTPUT_DIR}")
        sys.exit(1)

    print(f"  Found {len(mod_dirs)} addon folders")

    # Find armake2
    if not args.dry_run:
        armake2 = find_armake2()
        print(f"  armake2: {armake2}")
    print()

    # Build target path
    target_dir = Path(mods_dir) / f"@{mod_name}"
    addons_dir = target_dir / "addons"

    if args.dry_run:
        print(f"  [DRY RUN] Would create:")
        print(f"    {target_dir / 'mod.cpp'}")
        for d in sorted(mod_dirs):
            pbo_name = f"armadump_{d.name}.pbo"
            print(f"    {addons_dir / pbo_name}")
        print(f"\n  Total: {len(mod_dirs)} PBOs")
        return

    # Create directories
    addons_dir.mkdir(parents=True, exist_ok=True)

    # Generate mod.cpp
    mod_cpp = generate_mod_cpp(display_name, description, author, version)
    mod_cpp_path = target_dir / "mod.cpp"
    with open(mod_cpp_path, "w", encoding="utf-8") as f:
        f.write(mod_cpp)
    print(f"  Created {mod_cpp_path}")
    print()

    # Pack each addon into a PBO (all to temp first, then copy to final location)
    prefix_base = mod_name.lower()
    packed = 0
    failed = 0
    temp_dir = OUTPUT_DIR / "_pack_temp"
    addons_temp = temp_dir / "addons"
    addons_temp.mkdir(parents=True, exist_ok=True)

    for d in sorted(mod_dirs):
        pbo_name = f"armadump_{d.name}.pbo"
        temp_pbo = addons_temp / pbo_name
        prefix = f"{prefix_base}\\{d.name}"

        print(f"  Packing {d.name}...")
        if pack_pbo(armake2, d, temp_pbo, prefix):
            size_kb = temp_pbo.stat().st_size / 1024
            print(f"    -> {pbo_name} ({size_kb:.1f} KB)")
            packed += 1
        else:
            failed += 1

    if packed == 0:
        print("\nERROR: No PBOs packed.")
        sys.exit(1)

    # Copy mod.cpp to temp
    mod_cpp_temp = temp_dir / "mod.cpp"
    shutil.copy2(str(mod_cpp_path), str(mod_cpp_temp))

    # Swap: delete old dir, move temp into place
    print(f"\n  Installing to {target_dir}...")
    try:
        if target_dir.exists():
            shutil.rmtree(str(target_dir))
    except PermissionError:
        print("  WARNING: Could not remove old directory (Arma may have it locked).")
        print("  PBOs are ready in output/_pack_temp/ — close Arma and re-run, or copy manually.")

    try:
        shutil.move(str(temp_dir), str(target_dir))
    except Exception as e:
        print(f"  Could not move to {target_dir}: {e}")
        print(f"  PBOs are ready at: {temp_dir}")
        return

    print(f"\n{'=' * 60}")
    print(f"Done! {packed} PBOs packed, {failed} failed")
    print(f"  Mod: {target_dir}")
    print(f"  Load in Arma 3 launcher or via -mod=@{mod_name}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()