#!/usr/bin/env python3
"""
Comprehensive fix for all audit trail test files.
Fixes the AsyncAuditQueue and AuditTrailEngine API changes from Wave 107.
"""

import re
import sys
from pathlib import Path

# Files to fix
TEST_FILES = [
    "trading_engine/tests/async_audit_queue_tests.rs",
    "trading_engine/tests/audit_trail_persistence_test.rs",
    "trading_engine/tests/audit_persistence_tests.rs",
    "trading_engine/tests/audit_persistence_comprehensive.rs",
    "trading_engine/tests/audit_retention_tests.rs",
]

def fix_audit_trail_persistence_test(content):
    """Fix audit_trail_persistence_test.rs specific issues"""

    # Pattern 1: Fix AuditTrailEngine::new() - old API without await
    # OLD: let audit_engine = AuditTrailEngine::new(audit_config);
    # NEW: let wal_path = ...; let audit_engine = AuditTrailEngine::new(audit_config, postgres_pool.clone(), wal_path).await.unwrap();

    lines = content.split('\n')
    fixed_lines = []

    for i, line in enumerate(lines):
        # Fix AuditTrailEngine::new() calls without .await
        if 'let audit_engine = AuditTrailEngine::new(audit_config);' in line:
            indent = len(line) - len(line.lstrip())
            spaces = ' ' * indent
            fixed_lines.append(f'{spaces}let wal_path = std::env::temp_dir().join(format!("audit_wal_{{}}.log", uuid::Uuid::new_v4()));')
            fixed_lines.append(f'{spaces}let audit_engine = AuditTrailEngine::new(audit_config, postgres_pool.clone(), wal_path).await.unwrap();')
            continue

        # Remove .set_postgres_pool() calls (no longer needed)
        if '.set_postgres_pool(' in line and 'audit_engine' in line:
            continue

        # Fix AuditTrailConfig - remove old fields
        if 'real_time_persistence:' in line:
            continue

        fixed_lines.append(line)

    return '\n'.join(fixed_lines)

def fix_async_audit_queue_tests(content):
    """Fix async_audit_queue_tests.rs specific issues"""

    lines = content.split('\n')
    fixed_lines = []

    for i, line in enumerate(lines):
        # Fix AsyncAuditQueue::new() - now async and returns Result
        # OLD: let queue = Arc::new(AsyncAuditQueue::new(wal_path.clone()));
        # NEW: let queue = AsyncAuditQueue::new(wal_path.clone(), pool.clone(), 100, 100).await.unwrap();
        if 'Arc::new(AsyncAuditQueue::new(' in line and 'wal_path' in line:
            indent = len(line) - len(line.lstrip())
            spaces = ' ' * indent
            # Extract variable name
            var_match = re.search(r'let (\w+)', line)
            if var_match:
                var_name = var_match.group(1)
                fixed_lines.append(f'{spaces}let {var_name} = AsyncAuditQueue::new(wal_path.clone(), pool.clone(), 100, 100).await.unwrap();')
                continue

        # Fix submit calls that need to be await
        if '.submit(' in line and 'queue' in line and not '.await' in line and not 'expect' in line:
            line = line.replace('.submit(', '.submit(').replace(')?;', ').await?;').replace(');', ').await.unwrap();')

        fixed_lines.append(line)

    return '\n'.join(fixed_lines)

def fix_audit_config_fields(content):
    """Remove deprecated AuditTrailConfig fields"""

    deprecated_fields = [
        'enabled:',
        'compression_algorithm:',
        'encryption_algorithm:',
        'encryption_key:',
        'postgres_pool:',
        'file_path:',
        'enable_checksums:',
        'enable_tamper_detection:',
        'enable_best_execution_tracking:',
        'enable_mifid_reporting:',
        'compression_enabled:',
        'encryption_enabled:',
        'retention_days:',
        'real_time_persistence:',
    ]

    lines = content.split('\n')
    fixed_lines = []
    skip_next = False

    for i, line in enumerate(lines):
        if skip_next:
            skip_next = False
            continue

        # Skip lines with deprecated fields
        if any(field in line for field in deprecated_fields):
            continue

        fixed_lines.append(line)

    return '\n'.join(fixed_lines)

def fix_file(file_path):
    """Fix a single test file"""

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

    original_content = content

    # Apply fixes based on file
    if 'audit_trail_persistence_test' in file_path:
        content = fix_audit_trail_persistence_test(content)
    elif 'async_audit_queue_tests' in file_path:
        content = fix_async_audit_queue_tests(content)

    # Apply common fixes to all files
    content = fix_audit_config_fields(content)

    # Fix PostgresPool::new() calls - should use PostgresPool::new(config).await
    content = re.sub(
        r'PostgresPool::new\(postgres_config, (\d+)\)',
        r'PostgresPool::new(postgres_config)',
        content
    )

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

def main():
    base_path = Path('/home/jgrusewski/Work/foxhunt')

    fixed_count = 0
    for test_file in TEST_FILES:
        file_path = base_path / test_file
        if file_path.exists():
            if fix_file(str(file_path)):
                fixed_count += 1
        else:
            print(f"❌ File not found: {file_path}")

    print(f"\n📊 Summary: Fixed {fixed_count}/{len(TEST_FILES)} files")

    if fixed_count > 0:
        sys.exit(0)
    else:
        sys.exit(1)

if __name__ == '__main__':
    main()