A3Rebalancer/scripts/split_by_mod.py

"""
split_by_mod.py — Split raw Arma 3 CSV exports into per-mod folders.

Reads raw CSV files (extracted from .rpt log) and organizes them by source mod.
Each raw CSV has a 'source_mod' column (first column) that identifies which
CfgPatches class defined that entry.

Usage:
    python scripts/split_by_mod.py

Input files (in repo root):
    ammo_raw.csv      — from ammo.sqf
    armor_raw.csv     — from armor.sqf
    weapons_raw.csv   — from weapons.sqf
    magazines_raw.csv — from magazines.sqf

Output structure:
    data/
    ├── _rhs/
    │   ├── ammo.csv
    │   ├── armor.csv
    │   ├── weapons.csv
    │   ├── magazines.csv
    ├── _vanilla/
    │   ├── ammo.csv
    │   ├── ...
    ├── [mod_name]/
    │   ├── ammo.csv
    │   ├── ...
"""

import csv
import os
import re
import sys
from collections import defaultdict
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = REPO_ROOT / "data"

# Map (input_filename, output_filename)
CSV_FILES = [
    ("ammo_raw.csv", "ammo.csv"),
    ("armor_raw.csv", "armor.csv"),
    ("weapons_raw.csv", "weapons.csv"),
    ("magazines_raw.csv", "magazines.csv"),
]

# Known CfgPatches -> folder name mappings
# Prefix with _ for baseline mods (rhs, vanilla) so they sort first
# Order matters: more specific prefixes first
MOD_FOLDER_PREFIXES = [
    # RHS family
    ("rhs_", "_rhs"),
    ("rhsusf_", "_rhsusf"),
    ("rhsgref_", "_rhsgref"),
    ("rhssaf_", "_rhssaf"),
    # Vanilla
    ("A3_", "_vanilla"),
    ("A3a3_", "_vanilla"),
    # ACE
    ("ace_", "_ace"),
    # CBA
    ("cba_", "_cba"),
]

# Exact matches for edge cases
MOD_FOLDER_EXACT = {
    "ace_main": "_ace",
    "ace_ballistics": "_ace",
    "cba_common": "_cba",
    "cba_xeh": "_cba",
}


def sanitize_folder_name(mod_name: str) -> str:
    """Convert a CfgPatches class name to a safe folder name.

    Uses prefix matching for known mod families, then falls back to
    extracting the first meaningful segment.
    """
    # Exact match first
    if mod_name in MOD_FOLDER_EXACT:
        return MOD_FOLDER_EXACT[mod_name]

    # Prefix matching for known families
    for prefix, folder in MOD_FOLDER_PREFIXES:
        if mod_name.startswith(prefix):
            return folder

    # Heuristic: take the first segment before underscore,
    # stripping common suffixes
    parts = mod_name.lower().split("_")
    suffixes = {"main", "weapons", "data", "config", "core", "compat",
                "f", "e", "c", "add", "addons"}
    clean = [p for p in parts if p not in suffixes and p]
    if clean:
        return clean[0]

    return mod_name.lower()


def parse_rpt_csv(filepath: Path) -> list[list[str]]:
    """Parse a semicolon-delimited CSV file extracted from .rpt log.

    Handles lines that may have been wrapped or have trailing semicolons.
    """
    rows = []
    if not filepath.exists():
        print(f"  Warning: {filepath} not found, skipping")
        return rows

    with open(filepath, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # Split on semicolons, strip whitespace
            parts = [p.strip() for p in line.split(";")]
            # Remove trailing empty fields (from trailing semicolons)
            while parts and parts[-1] == "":
                parts.pop()
            if parts:
                rows.append(parts)

    return rows


def split_csv(input_name: str, output_name: str):
    """Split a single raw CSV file into per-mod files."""
    input_path = REPO_ROOT / input_name
    print(f"\nProcessing {input_name}...")

    rows = parse_rpt_csv(input_path)
    if not rows:
        return

    # First row is header
    header = rows[0]
    data_rows = rows[1:]

    # source_mod is always the first column
    mod_groups = defaultdict(list)
    for row in data_rows:
        if len(row) < 2:
            continue
        mod_name = row[0].strip().strip('"')
        folder = sanitize_folder_name(mod_name)
        mod_groups[folder].append(row)

    print(f"  Found {len(mod_groups)} unique mods, {len(data_rows)} total entries")

    # Write per-mod files
    for folder, group_rows in sorted(mod_groups.items()):
        mod_dir = DATA_DIR / folder
        mod_dir.mkdir(parents=True, exist_ok=True)
        out_path = mod_dir / output_name

        with open(out_path, "w", encoding="utf-8", newline="") as f:
            writer = csv.writer(f, delimiter=";")
            writer.writerow(header)
            writer.writerows(group_rows)

        print(f"  {folder:30s} -> {len(group_rows):5d} entries  ({out_path.relative_to(REPO_ROOT)})")


def main():
    print("=" * 60)
    print("ArmADump - Split CSVs by Mod")
    print("=" * 60)
    print(f"Repository root: {REPO_ROOT}")
    print(f"Output directory: {DATA_DIR}")

    DATA_DIR.mkdir(parents=True, exist_ok=True)

    for input_name, output_name in CSV_FILES:
        split_csv(input_name, output_name)

    print("\n" + "=" * 60)
    print("Done! Per-mod CSVs written to data/")
    print("=" * 60)


if __name__ == "__main__":
    main()