A3Rebalancer/scripts/extract_csvs.py

"""
extract_csvs.py — Extract CSV data from Arma 3 .rpt log file.

Parses the .rpt log and extracts CSV blocks delimited by
--- ... START --- / --- ... END --- markers.

Usage:
    python scripts/extract_csvs.py path/to/arma3.rpt
    python scripts/extract_csvs.py path/to/arma3.rpt --output-dir ./extracted

Output files (in repo root by default):
    ammo_raw.csv
    armor_raw.csv
    weapons_raw.csv
    magazines_raw.csv
"""

import argparse
import re
import sys
from pathlib import Path

# Arma 3 .rpt lines start with "HH:MM:SS " timestamp prefix
RPT_TIMESTAMP_RE = re.compile(r'^\d{2}:\d{2}:\d{2}\s+')

# Marker patterns -> output filename
MARKERS = {
    "AMMO EXPORT CSV":    "ammo_raw.csv",
    "ARMOR EXPORT CSV":   "armor_raw.csv",
    "WEAPONS EXPORT CSV": "weapons_raw.csv",
    "MAGAZINES EXPORT CSV": "magazines_raw.csv",
}

# Also detect the old sheet-style markers for backwards compat
LEGACY_MARKERS = {
    "CONFIG DUMP CSV SHEET 3 (AMMO)": "ammo_raw.csv",
    "CONFIG DUMP SHEET 1":            "weapons_raw.csv",
    "CONFIG DUMP SHEET 2":            "armor_raw.csv",
}

# Header line prefix to skip (the diag_log text("name;...") header)
HEADER_PREFIXES = ["name;", "name ", "source_mod;"]


def extract_from_rpt(rpt_path: Path, output_dir: Path) -> dict[str, int]:
    """Parse .rpt file and extract CSV blocks.

    Returns dict of filename -> number of data rows extracted.
    """
    if not rpt_path.exists():
        print(f"Error: {rpt_path} not found")
        sys.exit(1)

    print(f"Reading: {rpt_path}")
    print(f"File size: {rpt_path.stat().st_size / 1024 / 1024:.1f} MB")

    # Read the full file
    with open(rpt_path, "r", encoding="utf-8", errors="replace") as f:
        lines = f.readlines()

    print(f"Total lines: {len(lines)}")

    # Find all START/END marker positions
    # Format: --- AMMO EXPORT CSV START ---
    #         --- AMMO EXPORT CSV END ---
    start_pattern = re.compile(r'^\s*---\s+(.+?)\s+START\s+---\s*$')
    end_pattern = re.compile(r'^\s*---\s+(.+?)\s+END\s+---\s*$')

    # Build list of (start_line, end_line, marker_name)
    blocks = []
    active_block = None

    for i, line in enumerate(lines):
        # Strip RPT timestamp prefix (e.g. "17:55:18 ")
        line_clean = RPT_TIMESTAMP_RE.sub('', line).strip()
        # Strip surrounding quotes (.rpt wraps diag_log strings in "...")
        if line_clean.startswith('"') and line_clean.endswith('"'):
            line_clean = line_clean[1:-1].strip()

        start_match = start_pattern.match(line_clean)
        if start_match:
            marker = start_match.group(1).strip()
            active_block = {"marker": marker, "start": i + 1, "end": None}
            continue

        end_match = end_pattern.match(line_clean)
        if end_match:
            marker = end_match.group(1).strip()
            if active_block and active_block["marker"] == marker:
                active_block["end"] = i  # end line (exclusive)
                blocks.append(active_block)
                active_block = None

    print(f"Found {len(blocks)} CSV blocks in .rpt")

    # Map marker names to output filenames
    marker_to_file = {}
    for marker_name, filename in {**MARKERS, **LEGACY_MARKERS}.items():
        marker_to_file[marker_name] = filename

    # Extract each block
    results = {}
    output_dir.mkdir(parents=True, exist_ok=True)

    for block in blocks:
        marker = block["marker"]
        filename = marker_to_file.get(marker)

        if not filename:
            # Try partial match
            for key, fname in marker_to_file.items():
                if key in marker:
                    filename = fname
                    break

        if not filename:
            print(f"  Skipping unrecognized marker: {marker}")
            continue

        # Extract lines between START and END (exclusive of marker lines)
        data_lines = []
        header_found = False
        for line in lines[block["start"]:block["end"]]:
            # Strip RPT timestamp prefix
            stripped = RPT_TIMESTAMP_RE.sub('', line).strip()
            # Strip surrounding quotes
            if stripped.startswith('"') and stripped.endswith('"'):
                stripped = stripped[1:-1].strip()
            # Skip empty lines
            if not stripped:
                continue
            # Remove diag_log wrapper if present: text("...") -> ...
            cleaned = stripped
            if cleaned.startswith('text("') and cleaned.endswith('")'):
                cleaned = cleaned[6:-2]
            elif cleaned.startswith("text('") and cleaned.endswith("')"):
                cleaned = cleaned[6:-2]
            # Keep only the first header row (diag_log text("name;..."))
            if any(cleaned.startswith(p) for p in HEADER_PREFIXES):
                if not header_found:
                    data_lines.append(cleaned)
                    header_found = True
                continue
            data_lines.append(cleaned)

        out_path = output_dir / filename
        with open(out_path, "w", encoding="utf-8", newline="") as f:
            for line in data_lines:
                f.write(line + "\n")

        results[filename] = len(data_lines)
        print(f"  {filename:25s} -> {len(data_lines):6d} rows  ({out_path})")

    return results


def main():
    parser = argparse.ArgumentParser(
        description="Extract CSV data from Arma 3 .rpt log file")
    parser.add_argument(
        "rpt_file",
        help="Path to Arma 3 .rpt log file")
    parser.add_argument(
        "--output-dir", "-o",
        default=".",
        help="Output directory for CSV files (default: repo root)")
    args = parser.parse_args()

    rpt_path = Path(args.rpt_file).resolve()
    output_dir = Path(args.output_dir).resolve()

    print("=" * 60)
    print("ArmADump - Extract CSVs from .rpt log")
    print("=" * 60)

    results = extract_from_rpt(rpt_path, output_dir)

    print(f"\n{'=' * 60}")
    if results:
        print(f"Extracted {len(results)} CSV files to {output_dir}")
        for fname, count in sorted(results.items()):
            print(f"  {fname}: {count} rows")
        print("\nNext step: python scripts/split_by_mod.py")
    else:
        print("No CSV blocks found in .rpt file!")
        print("Make sure the extraction scripts were run in Arma first.")
    print("=" * 60)


if __name__ == "__main__":
    main()