#!/usr/bin/env python3
"""
Fixed migration script for AsyncAuditQueue API changes.
"""

import re
import sys
from pathlib import Path

def fix_file_content(content: str) -> str:
    """Apply all fixes to file content."""

    # Fix 1: AsyncAuditQueue::new() - add parameters and .await
    # Pattern: AsyncAuditQueue::new(wal_path_expr)
    # Replace: AsyncAuditQueue::new(wal_path_expr, Arc::clone(&pool), 100, 100).await.expect("...")
    content = re.sub(
        r'AsyncAuditQueue::new\(([^)]+)\)(?!\s*\.await)',
        r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")',
        content
    )

    # Fix 2: Remove Arc::new() wrapper from AsyncAuditQueue (it returns Arc now)
    content = re.sub(
        r'Arc::new\(AsyncAuditQueue::new\(([^)]+), Arc::clone\(&pool\), 100, 100\)\.await\.expect\("Failed to create AsyncAuditQueue"\)\)',
        r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")',
        content
    )

    # Fix 3: AuditTrailEngine::new() - add pool and wal_path parameters
    # Only fix if it doesn't already have .await
    content = re.sub(
        r'AuditTrailEngine::new\(([^,)]+)\)(?!\s*\.await)',
        r'AuditTrailEngine::new(\1, Arc::clone(&pool), wal_path.clone()).await.expect("Failed to create AuditTrailEngine")',
        content
    )

    # Fix 4: PostgresPool::new() with 2 args -> PostgresPool::connect() with 1 arg
    content = re.sub(
        r'PostgresPool::new\(([^,)]+),\s*\d+\)',
        r'PostgresPool::connect(\1).await.expect("Failed to connect to PostgreSQL")',
        content
    )

    # Fix 5: Add missing pool setup at start of test functions if needed
    # This is done more carefully per-function
    lines = content.split('\n')
    result = []
    i = 0

    while i < len(lines):
        line = lines[i]
        result.append(line)

        # Check if this is an async test function
        if 'async fn test_' in line or 'async fn bench_' in line:
            # Look for opening brace
            j = i + 1
            while j < len(lines) and '{' not in lines[j]:
                result.append(lines[j])
                j += 1

            if j < len(lines):
                result.append(lines[j])  # Add the line with {

                # Check if next few lines already have pool setup
                has_pool = False
                for k in range(j + 1, min(j + 5, len(lines))):
                    if 'let pool =' in lines[k] or 'setup_test_postgres_pool' in lines[k]:
                        has_pool = True
                        break

                # Check if function uses AsyncAuditQueue or AuditTrailEngine
                needs_pool = False
                for k in range(j + 1, min(j + 30, len(lines))):
                    if 'AsyncAuditQueue::new' in lines[k] or 'AuditTrailEngine::new' in lines[k]:
                        needs_pool = True
                        break

                # Add pool setup if needed
                if needs_pool and not has_pool:
                    result.append('    // Setup PostgreSQL pool for audit persistence')
                    result.append('    let pool = setup_test_postgres_pool().await;')
                    result.append('    let wal_path = std::env::temp_dir().join(format!("foxhunt_audit_test_{}.wal", uuid::Uuid::new_v4()));')
                    result.append('')

                i = j

        i += 1

    content = '\n'.join(result)

    # Fix 6: Remove old AuditTrailConfig fields
    old_fields_patterns = [
        r'\s+enabled:\s*\w+,?\s*\n',
        r'\s+compression_algorithm:\s*CompressionAlgorithm::\w+,?\s*\n',
        r'\s+encryption_algorithm:\s*EncryptionAlgorithm::\w+,?\s*\n',
        r'\s+encryption_key:\s*vec!\[.*?\],?\s*\n',
        r'\s+postgres_pool:\s*[^,\n]+,?\s*\n',
        r'\s+file_path:\s*[^,\n]+,?\s*\n',
        r'\s+enable_checksums:\s*\w+,?\s*\n',
        r'\s+enable_tamper_detection:\s*\w+,?\s*\n',
        r'\s+enable_best_execution_tracking:\s*\w+,?\s*\n',
        r'\s+enable_mifid_reporting:\s*\w+,?\s*\n',
    ]

    for pattern in old_fields_patterns:
        content = re.sub(pattern, '', content)

    # Fix 7: Fix enum variants
    content = content.replace('EncryptionAlgorithm::Aes256Gcm', 'EncryptionAlgorithm::AES256GCM')

    # Fix 8: Remove duplicate PostgresPool::connect().await.expect().await patterns
    content = re.sub(
        r'PostgresPool::connect\(([^)]+)\)\.await\.expect\("([^"]+)"\)\.await',
        r'PostgresPool::connect(\1).await.expect("\2")',
        content
    )

    return content

def process_file(file_path: Path) -> bool:
    """Process a single test file."""
    print(f"Processing {file_path.name}...")

    try:
        original = file_path.read_text()
        fixed = fix_file_content(original)

        if fixed != original:
            file_path.write_text(fixed)
            print(f"  ✅ Fixed {file_path.name}")
            return True
        else:
            print(f"  ⏭️  No changes for {file_path.name}")
            return False

    except Exception as e:
        print(f"  ❌ Error: {e}")
        return False

def main():
    test_dir = Path(__file__).parent.parent / 'trading_engine' / 'tests'

    if not test_dir.exists():
        print(f"Error: {test_dir} not found")
        return 1

    test_files = sorted(test_dir.glob('*.rs'))

    if not test_files:
        print(f"Error: No test files in {test_dir}")
        return 1

    print(f"Found {len(test_files)} test files\n")

    updated = sum(1 for f in test_files if process_file(f))

    print(f"\n{'='*60}")
    print(f"Updated {updated}/{len(test_files)} files")
    print(f"{'='*60}")

    return 0

if __name__ == '__main__':
    sys.exit(main())