"""Load test scenarios for Story 6-8.

This module implements 4 load testing scenarios:
1. Constant Load: 1000 RPS sustained for 60 seconds
2. Ramp-Up: Gradual increase from 100 to 1000 RPS over 30s, sustain 30s
3. Burst Load: 1000 RPS baseline with 2000 RPS burst
4. Spike Recovery: Multiple burst cycles to test recovery
"""

import logging
import statistics
import time
from typing import Dict, List, Tuple

import psutil

from mca_sdk import MCAClient
from .performance_helpers import get_collector_endpoint

logger = logging.getLogger(__name__)


def run_load_scenario(
    scenario_name: str,
    rps_schedule: List[Tuple[float, int]],
    collector_endpoint: str = None,
) -> Dict:
    """Run a load test scenario with variable RPS schedule.

    Args:
        scenario_name: Name of the scenario for logging
        rps_schedule: List of (duration_seconds, target_rps) tuples
        collector_endpoint: OTel collector endpoint (uses env var or localhost if None)

    Returns:
        Dict containing test results and statistics
    """
    if collector_endpoint is None:
        collector_endpoint = get_collector_endpoint()

    # Initialize client
    client = MCAClient(
        service_name=f"load_test_{scenario_name}",
        model_id="load_test_model",
        team_name="performance_team",
        model_category="internal",
        model_type="regression",
        collector_endpoint=collector_endpoint,
        buffering_enabled=True,
        max_queue_size=20000,  # Large queue for burst scenarios
    )

    # Tracking
    latencies: List[float] = []
    request_timestamps: List[float] = []  # Absolute timestamps for each request
    cpu_samples: List[Tuple[float, float]] = []  # (timestamp, cpu_percent)
    request_count = 0
    process = psutil.Process()

    total_duration = sum(duration for duration, _ in rps_schedule)
    logger.info(f"Starting load scenario: {scenario_name}")
    logger.info(f"  Total duration: {total_duration}s")
    logger.info(f"  Collector: {collector_endpoint}")
    logger.info(f"  Schedule: {rps_schedule}")

    start_time = time.time()
    schedule_start = start_time

    try:
        for phase_idx, (phase_duration, target_rps) in enumerate(rps_schedule):
            phase_start = time.time()
            phase_requests = 0
            interval = 1.0 / target_rps if target_rps > 0 else 0.001

            logger.info(f"  Phase {phase_idx + 1}: {target_rps} RPS for {phase_duration}s")

            # Use absolute time scheduling to avoid drift
            next_request_time = time.perf_counter()

            while time.time() - phase_start < phase_duration:
                # Capture timestamp BEFORE request
                request_timestamp = time.time()

                # Record prediction (measure SDK overhead)
                req_start = time.perf_counter()
                client.record_prediction(latency=0.001)
                req_end = time.perf_counter()

                sdk_overhead_ms = (req_end - req_start) * 1000
                latencies.append(sdk_overhead_ms)
                request_timestamps.append(request_timestamp)
                request_count += 1
                phase_requests += 1

                # Sample CPU periodically (non-blocking)
                if request_count % 100 == 0:
                    cpu = process.cpu_percent(interval=None)  # Non-blocking instantaneous read
                    cpu_samples.append((time.time(), cpu))

                # Schedule next request using absolute time (drift compensation)
                next_request_time += interval
                now = time.perf_counter()
                sleep_time = next_request_time - now

                if sleep_time > 0:
                    time.sleep(sleep_time)
                elif sleep_time < -0.1:
                    # If we're more than 100ms behind, reset to avoid perpetual catchup
                    next_request_time = now

            phase_elapsed = time.time() - phase_start
            phase_rps = phase_requests / phase_elapsed if phase_elapsed > 0 else 0
            logger.info(
                f"    Phase {phase_idx + 1} complete: {phase_requests} requests, {phase_rps:.1f} actual RPS"
            )

    except KeyboardInterrupt:
        logger.warning("Test interrupted by user")

    finally:
        elapsed_time = time.time() - start_time
        actual_rps = request_count / elapsed_time if elapsed_time > 0 else 0

        # Shutdown client first to flush queues
        client.shutdown()

        # Get final buffer stats AFTER shutdown to include any drops during flush
        buffer_stats = client.buffer_stats() if hasattr(client, "buffer_stats") else {}

        # Calculate total dropped items
        dropped_count = 0
        if buffer_stats and "per_signal_stats" in buffer_stats:
            for signal_name, signal_stats in buffer_stats["per_signal_stats"].items():
                dropped_count += signal_stats.get("items_dropped", 0)

    # Calculate statistics
    latencies.sort()
    p50 = statistics.median(latencies) if latencies else 0
    p95 = (
        statistics.quantiles(latencies, n=20)[18]
        if len(latencies) > 20
        else (latencies[-1] if latencies else 0)
    )
    p99 = (
        statistics.quantiles(latencies, n=100)[98]
        if len(latencies) > 100
        else (latencies[-1] if latencies else 0)
    )
    avg = statistics.mean(latencies) if latencies else 0

    # CPU analysis
    cpu_values = [cpu for _, cpu in cpu_samples]
    avg_cpu = statistics.mean(cpu_values) if cpu_values else 0
    max_cpu = max(cpu_values) if cpu_values else 0

    # CPU stability check (with absolute threshold to handle low baseline)
    cpu_stable = True
    cpu_increase_percent = 0
    cpu_absolute_increase = 0
    if len(cpu_samples) >= 10:
        first_half = cpu_values[: len(cpu_values) // 2]
        second_half = cpu_values[len(cpu_values) // 2 :]
        avg_first_half = statistics.mean(first_half)
        avg_second_half = statistics.mean(second_half)
        cpu_absolute_increase = avg_second_half - avg_first_half
        cpu_increase_percent = (
            ((avg_second_half - avg_first_half) / avg_first_half * 100) if avg_first_half > 0 else 0
        )

        # Stable if absolute increase <5% OR relative increase <10% (whichever is more lenient)
        cpu_stable = (cpu_absolute_increase < 5.0) or (cpu_increase_percent < 10)

    # Results
    results = {
        "scenario": scenario_name,
        "duration_seconds": elapsed_time,
        "total_requests": request_count,
        "actual_rps": actual_rps,
        "dropped_count": dropped_count,
        "request_timestamps": request_timestamps,  # For real throughput timeline
        "latency": {
            "p50_ms": p50,
            "p95_ms": p95,
            "p99_ms": p99,
            "avg_ms": avg,
            "min_ms": min(latencies) if latencies else 0,
            "max_ms": max(latencies) if latencies else 0,
            "all_latencies": latencies,  # For graphing
        },
        "cpu": {
            "average_percent": avg_cpu,
            "max_percent": max_cpu,
            "samples": len(cpu_samples),
            "stable": cpu_stable,
            "increase_percent": cpu_increase_percent,
            "timeline": cpu_samples,  # For graphing
        },
        "acceptance_criteria": {
            "p95_under_100ms": p95 < 100,
            "no_drops": dropped_count == 0,
            "cpu_stable": cpu_stable,
        },
    }

    # Print summary
    logger.info(f"\n{'=' * 60}")
    logger.info(f"SCENARIO: {scenario_name}")
    logger.info(f"{'=' * 60}")
    logger.info(f"Duration: {elapsed_time:.1f}s")
    logger.info(f"Requests: {request_count:,}")
    logger.info(f"Actual RPS: {actual_rps:.1f}")
    logger.info(f"Dropped: {dropped_count}")
    logger.info(f"Latency p95: {p95:.3f}ms")
    logger.info(f"CPU avg: {avg_cpu:.2f}%")
    logger.info(f"CPU stable: {'YES' if cpu_stable else 'NO'}")
    logger.info(f"{'=' * 60}\n")

    return results


def scenario_constant_load(duration_seconds: int = 60, target_rps: int = 1000) -> Dict:
    """Scenario 1: Constant load at target RPS for specified duration.

    Args:
        duration_seconds: Test duration in seconds (default 60)
        target_rps: Target requests per second (default 1000)

    Returns:
        Dict containing test results
    """
    return run_load_scenario(
        scenario_name="constant_load",
        rps_schedule=[(duration_seconds, target_rps)],
    )


def scenario_ramp_up(ramp_duration: int = 30, sustain_duration: int = 30) -> Dict:
    """Scenario 2: Ramp up from 100 to 1000 RPS, then sustain.

    Args:
        ramp_duration: Duration of ramp-up phase in seconds
        sustain_duration: Duration to sustain peak load in seconds

    Returns:
        Dict containing test results
    """
    # Break ramp into 5-second intervals with increasing RPS
    ramp_steps = ramp_duration // 5
    rps_increments = [100 + (900 // ramp_steps) * i for i in range(ramp_steps + 1)]

    schedule = []
    for rps in rps_increments[:-1]:
        schedule.append((5, rps))

    # Add sustain phase at 1000 RPS
    schedule.append((sustain_duration, 1000))

    return run_load_scenario(
        scenario_name="ramp_up",
        rps_schedule=schedule,
    )


def scenario_burst_load(
    baseline_duration: int = 20, burst_duration: int = 10, recovery_duration: int = 20
) -> Dict:
    """Scenario 3: Baseline load with burst spike.

    Args:
        baseline_duration: Initial baseline duration in seconds
        burst_duration: Burst spike duration in seconds
        recovery_duration: Recovery duration in seconds

    Returns:
        Dict containing test results
    """
    schedule = [
        (baseline_duration, 1000),  # Baseline
        (burst_duration, 2000),  # Burst to 2x
        (recovery_duration, 1000),  # Recovery
    ]

    return run_load_scenario(
        scenario_name="burst_load",
        rps_schedule=schedule,
    )


def scenario_spike_recovery(num_cycles: int = 3, cycle_duration: int = 20) -> Dict:
    """Scenario 4: Multiple burst cycles to test recovery behavior.

    Args:
        num_cycles: Number of burst cycles
        cycle_duration: Duration of each cycle phase in seconds

    Returns:
        Dict containing test results
    """
    schedule = []

    for i in range(num_cycles):
        schedule.append((cycle_duration // 2, 1000))  # Baseline
        schedule.append((cycle_duration // 2, 2000))  # Burst

    return run_load_scenario(
        scenario_name="spike_recovery",
        rps_schedule=schedule,
    )


if __name__ == "__main__":
    import argparse

    # Configure logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    parser = argparse.ArgumentParser(description="Run load test scenarios")
    parser.add_argument(
        "--scenario",
        choices=["constant", "ramp", "burst", "spike", "all"],
        default="all",
        help="Scenario to run",
    )
    parser.add_argument(
        "--duration",
        type=int,
        default=60,
        help="Duration for constant load scenario (seconds)",
    )
    args = parser.parse_args()

    scenarios = {
        "constant": lambda: scenario_constant_load(duration_seconds=args.duration),
        "ramp": scenario_ramp_up,
        "burst": scenario_burst_load,
        "spike": scenario_spike_recovery,
    }

    if args.scenario == "all":
        results = []
        for name, func in scenarios.items():
            logger.info(f"\n\nRunning scenario: {name}")
            result = func()
            results.append(result)

            # Brief pause between scenarios
            time.sleep(5)
    else:
        results = [scenarios[args.scenario]()]

    # Summary
    logger.info("\n\n" + "=" * 60)
    logger.info("ALL SCENARIOS COMPLETE")
    logger.info("=" * 60)
    for result in results:
        status = "PASS" if all(result["acceptance_criteria"].values()) else "FAIL"
        logger.info(f"{result['scenario']:20s} - {status}")
    logger.info("=" * 60)
