#!/usr/bin/env python3
"""
Parallel fuzzing runner for long-duration testing.
Spawns multiple processes to saturate CPU cores.
"""

import sys
import multiprocessing as mp
import time
import json
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from test_generator import TestCaseGenerator
from differential_tester import DifferentialTester


def worker_process(worker_id: int, duration_seconds: int, batch_size: int, output_dir: str):
    """Worker process that runs fuzzing tests."""
    print(f"Worker {worker_id} started (PID: {mp.current_process().pid})")

    gen = TestCaseGenerator(seed=None)  # Random seed
    tester = DifferentialTester(log_file=f"{output_dir}/worker_{worker_id}.log")

    start_time = time.time()
    end_time = start_time + duration_seconds
    batch_num = 0

    while time.time() < end_time:
        batch_num += 1

        # Generate and run test batch
        test_cases = gen.generate_fuzzing_cases(count=batch_size)

        for test_case in test_cases:
            try:
                tester.run_test_case(test_case)
            except Exception as e:
                tester.results['errors'] += 1
                tester.log(f"ERROR: {e}")

        # Save checkpoint every 10 batches
        if batch_num % 10 == 0:
            checkpoint_file = f"{output_dir}/worker_{worker_id}_checkpoint_{batch_num}.json"
            tester.save_results(checkpoint_file)

            elapsed = time.time() - start_time
            remaining = end_time - time.time()
            print(f"Worker {worker_id}: Batch {batch_num}, "
                  f"Elapsed: {elapsed/3600:.2f}h, "
                  f"Remaining: {remaining/3600:.2f}h, "
                  f"Passed: {tester.results['passed']}, "
                  f"Failed: {tester.results['failed']}")

    # Save final results
    final_file = f"{output_dir}/worker_{worker_id}_final.json"
    tester.save_results(final_file)

    print(f"Worker {worker_id} completed: "
          f"Passed={tester.results['passed']}, "
          f"Failed={tester.results['failed']}, "
          f"Errors={tester.results['errors']}")

    return {
        'worker_id': worker_id,
        'passed': tester.results['passed'],
        'failed': tester.results['failed'],
        'errors': tester.results['errors']
    }


def run_parallel_fuzzing(duration_hours: int = 24, num_workers: int = None, batch_size: int = 50):
    """Run parallel fuzzing tests."""
    if num_workers is None:
        num_workers = max(1, mp.cpu_count() - 1)  # Leave one core free

    output_dir = "/home/ubuntu-exp/cronproject/croniter-rs/tests/differential/fuzzing_results"
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    print(f"{'='*80}")
    print("Parallel Fuzzing Test Runner")
    print(f"{'='*80}")
    print(f"Duration: {duration_hours} hours")
    print(f"Workers: {num_workers}")
    print(f"Batch size: {batch_size}")
    print(f"Output directory: {output_dir}")
    print(f"Start time: {datetime.now()}")
    print(f"{'='*80}\n")

    duration_seconds = duration_hours * 3600

    # Start worker processes
    with mp.Pool(processes=num_workers) as pool:
        results = []
        for worker_id in range(num_workers):
            result = pool.apply_async(worker_process,
                                     (worker_id, duration_seconds, batch_size, output_dir))
            results.append(result)

        # Wait for all workers to complete
        pool.close()
        pool.join()

        # Collect results
        worker_results = [r.get() for r in results]

    # Aggregate results
    total_passed = sum(r['passed'] for r in worker_results)
    total_failed = sum(r['failed'] for r in worker_results)
    total_errors = sum(r['errors'] for r in worker_results)
    total_tests = total_passed + total_failed + total_errors

    print(f"\n{'='*80}")
    print("Parallel Fuzzing Complete")
    print(f"{'='*80}")
    print(f"End time: {datetime.now()}")
    print(f"Total tests: {total_tests}")
    print(f"Passed: {total_passed} ({100*total_passed/total_tests if total_tests > 0 else 0:.1f}%)")
    print(f"Failed: {total_failed} ({100*total_failed/total_tests if total_tests > 0 else 0:.1f}%)")
    print(f"Errors: {total_errors} ({100*total_errors/total_tests if total_tests > 0 else 0:.1f}%)")
    print(f"\nResults saved to: {output_dir}/")
    print(f"{'='*80}")

    # Save aggregate results
    aggregate = {
        'duration_hours': duration_hours,
        'num_workers': num_workers,
        'batch_size': batch_size,
        'start_time': datetime.now().isoformat(),
        'total_tests': total_tests,
        'total_passed': total_passed,
        'total_failed': total_failed,
        'total_errors': total_errors,
        'worker_results': worker_results
    }

    with open(f"{output_dir}/aggregate_results.json", 'w') as f:
        json.dump(aggregate, f, indent=2)

    return total_failed == 0 and total_errors == 0


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Parallel fuzzing test runner')
    parser.add_argument('--duration', type=int, default=24,
                       help='Duration in hours (default: 24)')
    parser.add_argument('--workers', type=int, default=None,
                       help='Number of worker processes (default: CPU count - 1)')
    parser.add_argument('--batch-size', type=int, default=50,
                       help='Test cases per batch (default: 50)')

    args = parser.parse_args()

    success = run_parallel_fuzzing(
        duration_hours=args.duration,
        num_workers=args.workers,
        batch_size=args.batch_size
    )

    sys.exit(0 if success else 1)
