#!/usr/bin/env python3
"""
Extract adjustment columns from input files and create trimmed versions.

Cross recoupment and transfers: Column AG (index 35) onward
Overrides: Column AD (index 29) onward
"""

import csv
from pathlib import Path

INPUT_DIR = Path(__file__).parent / "input"
OUTPUT_DIR = Path(__file__).parent / "output"
TRIMMED_DIR = Path(__file__).parent / "trimmed_input"

OUTPUT_DIR.mkdir(exist_ok=True)
TRIMMED_DIR.mkdir(exist_ok=True)

def extract_and_trim(input_file, adjustment_start_col, output_prefix):
    """
    Extract adjustment columns and create trimmed input file.

    Args:
        input_file: Path to input CSV file
        adjustment_start_col: Column index where adjustments start (0-based)
        output_prefix: Prefix for output files
    """
    input_path = INPUT_DIR / input_file
    original_adj_path = OUTPUT_DIR / f"original_{output_prefix}_adjustments.csv"
    trimmed_input_path = TRIMMED_DIR / input_file

    print(f"\nProcessing {input_file}...")
    print(f"  Adjustment columns start at index: {adjustment_start_col}")

    with open(input_path, 'r', encoding='utf-8') as infile:
        rows = list(csv.reader(infile))

    if not rows:
        print(f"  WARNING: Empty file!")
        return

    adjustment_rows = []
    trimmed_rows = []

    for row in rows:
        if len(row) > adjustment_start_col:
            adjustment_rows.append(row[adjustment_start_col:])
            trimmed_rows.append(row[:adjustment_start_col])
        else:
            adjustment_rows.append([])
            trimmed_rows.append(row)

    with open(original_adj_path, 'w', encoding='utf-8', newline='') as outfile:
        csv.writer(outfile).writerows(adjustment_rows)

    print(f"   Saved original adjustments to: {original_adj_path}")
    print(f"    Rows: {len(adjustment_rows)}, Columns: {len(adjustment_rows[0]) if adjustment_rows else 0}")

    with open(trimmed_input_path, 'w', encoding='utf-8', newline='') as outfile:
        csv.writer(outfile).writerows(trimmed_rows)

    print(f"   Saved trimmed input to: {trimmed_input_path}")
    print(f"    Rows: {len(trimmed_rows)}, Columns: {len(trimmed_rows[0]) if trimmed_rows else 0}")

def main():
    print("=" * 70)
    print("EXTRACTING ADJUSTMENTS AND CREATING TRIMMED INPUT FILES")
    print("=" * 70)

    # Cross recoupment & transfers: column AG (based on master sheet)
    extract_and_trim("input_cross_recoupment.csv", 35, "cross_recoupment")
    extract_and_trim("input_transfers.csv", 35, "transfers")

    # Overrides: column AD (based on master sheet)
    extract_and_trim("input_overrides.csv", 29, "overrides")

    print("\n" + "=" * 70)
    print("EXTRACTION COMPLETE!")
    print("=" * 70)
    print(f"\nOriginal adjustments saved in: {OUTPUT_DIR}")
    print(f"Trimmed input files saved in: {TRIMMED_DIR}")

if __name__ == "__main__":
    main()