#!/usr/bin/env python3
"""
Fix audit_compliance.rs by updating all AuditTrailEngine::new() calls.
This script performs surgical fixes to update the API from:
    AuditTrailEngine::new(config)
to:
    AuditTrailEngine::new(config, pool, wal_path).await.unwrap()
"""

import re
import sys

def fix_audit_compliance(file_path):
    with open(file_path, 'r') as f:
        content = f.read()

    original_content = content

    # Pattern 1: Find test functions and add wal_path after pg_pool setup
    # We'll insert WAL path right before AuditTrailEngine::new() calls

    # Find all AuditTrailEngine::new(config) patterns
    pattern = r'let audit_engine = AuditTrailEngine::new\(config\)\.await\.unwrap\(\);'

    def replacement(match):
        # Extract indentation
        return '''let wal_path = std::env::temp_dir().join(format!("audit_wal_{}.log", uuid::Uuid::new_v4()));
    let audit_engine = AuditTrailEngine::new(config, pg_pool.clone().unwrap(), wal_path).await.unwrap();'''

    # First, let's find all occurrences and their line positions
    lines = content.split('\n')
    fixed_lines = []

    for i, line in enumerate(lines):
        if 'let audit_engine = AuditTrailEngine::new(config).await.unwrap();' in line:
            # Get indentation
            indent = len(line) - len(line.lstrip())
            spaces = ' ' * indent

            # Add wal_path line before audit_engine
            fixed_lines.append(f'{spaces}let wal_path = std::env::temp_dir().join(format!("audit_wal_{{}}.log", uuid::Uuid::new_v4()));')
            # Fix the audit_engine line
            fixed_lines.append(f'{spaces}let audit_engine = AuditTrailEngine::new(config, pg_pool.clone().unwrap(), wal_path).await.unwrap();')
        else:
            fixed_lines.append(line)

    content = '\n'.join(fixed_lines)

    # Count changes
    changes = content.count('pg_pool.clone().unwrap(), wal_path')

    # Write back
    if content != original_content:
        with open(file_path, 'w') as f:
            f.write(content)
        print(f"✅ Fixed {changes} AuditTrailEngine::new() calls in {file_path}")
        return True
    else:
        print(f"⚠️  No changes needed in {file_path}")
        return False

if __name__ == '__main__':
    file_path = '/home/jgrusewski/Work/foxhunt/trading_engine/tests/audit_compliance.rs'
    if fix_audit_compliance(file_path):
        sys.exit(0)
    else:
        sys.exit(1)