"""
Aggregation strategies for multiple experiment runs.

When multiple experiments match the same filter conditions,
users can choose how to aggregate the results.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional

from .models import WorkloadRunData


class AggregationStrategy(str, Enum):
    """Strategy for aggregating multiple experiment runs."""

    ALL = "all"  # Show all runs
    AVERAGE = "average"  # Average metrics across runs
    BEST = "best"  # Best performing run (highest TPS/tpmC, lowest latency)
    LATEST = "latest"  # Most recent run
    MEDIAN = "median"  # Median metrics across runs


@dataclass
class AggregatedMetrics:
    """Aggregated metrics from multiple runs.

    Contains both the aggregated values and metadata about the aggregation.
    """

    # Aggregated metric values
    metrics: Dict  # e.g., {"tpmC": 7500, "p95": 100.5}

    # Metadata
    strategy: AggregationStrategy
    run_count: int  # Number of runs aggregated
    task_ids: List[int]  # IDs of tasks included in aggregation
    workflow_execution_ids: List[str]  # Execution IDs included


def aggregate_runs(
    runs: List[WorkloadRunData],
    strategy: AggregationStrategy,
    workload_type: str,
) -> Optional[WorkloadRunData]:
    """Aggregate multiple runs based on strategy.

    Args:
        runs: List of runs to aggregate (same workload, same thread count)
        strategy: Aggregation strategy to use
        workload_type: Type of workload (e.g., "bench-tpcc", "bench-sysbench")

    Returns:
        Single WorkloadRunData representing the aggregated result, or None if no valid runs
    """
    if not runs:
        return None

    # Filter out runs without metrics (but keep both SUCCESS and FAILURE runs)
    # Failed runs may still have useful performance data
    runs_with_metrics = [r for r in runs if r.metrics]
    if not runs_with_metrics:
        return None

    if strategy == AggregationStrategy.ALL:
        # For ALL strategy, don't aggregate - caller should handle separately
        return None
    elif strategy == AggregationStrategy.LATEST:
        return _aggregate_latest(runs_with_metrics)
    elif strategy == AggregationStrategy.BEST:
        return _aggregate_best(runs_with_metrics, workload_type)
    elif strategy == AggregationStrategy.AVERAGE:
        return _aggregate_average(runs_with_metrics, workload_type)
    elif strategy == AggregationStrategy.MEDIAN:
        return _aggregate_median(runs_with_metrics, workload_type)

    return None


def _aggregate_latest(runs: List[WorkloadRunData]) -> WorkloadRunData:
    """Get the most recent run.

    Args:
        runs: List of runs with metrics (includes both SUCCESS and FAILURE)

    Returns:
        Most recent run
    """
    # Sort by completed_at or workflow_execution_id (descending)
    sorted_runs = sorted(
        runs,
        key=lambda r: (r.completed_at or "", r.workflow_execution_id or ""),
        reverse=True,
    )
    return sorted_runs[0]


def _aggregate_best(runs: List[WorkloadRunData], workload_type: str) -> WorkloadRunData:
    """Get the best performing run.

    Best is defined as:
    - For TPCC: Highest tpmC
    - For Sysbench: Highest TPS
    - If primary metric is tied, prefer lowest P95 latency

    Args:
        runs: List of runs with metrics (includes both SUCCESS and FAILURE)
        workload_type: Type of workload

    Returns:
        Best performing run
    """
    if "tpcc" in workload_type.lower():
        # Sort by tpmC (descending), then by P95 latency (ascending)
        sorted_runs = sorted(
            runs,
            key=lambda r: (
                -r.metrics.get("tpmC", 0),
                r.metrics.get("latency", {}).get("p95", float("inf")),
            ),
        )
    else:
        # For sysbench and others, sort by TPS (descending), then by P95 latency (ascending)
        sorted_runs = sorted(
            runs,
            key=lambda r: (
                -r.metrics.get("tps", 0),
                r.metrics.get("latency", {}).get("p95", float("inf")),
            ),
        )

    return sorted_runs[0]


def _aggregate_average(
    runs: List[WorkloadRunData], workload_type: str
) -> WorkloadRunData:
    """Calculate average metrics across runs.

    Args:
        runs: List of runs with metrics (includes both SUCCESS and FAILURE)
        workload_type: Type of workload

    Returns:
        WorkloadRunData with averaged metrics
    """
    if not runs:
        return None

    # Use first run as template
    template = runs[0]

    # Calculate average metrics
    averaged_metrics = _calculate_average_metrics(runs, workload_type)

    # Create aggregated run
    return WorkloadRunData(
        node_name=f"{template.node_name} (avg of {len(runs)} runs)",
        workload_type=template.workload_type,
        params=template.params,
        metrics=averaged_metrics,
        status="SUCCEEDED",
        task_id=None,  # Not a single task
        workflow_execution_id=f"aggregated_{len(runs)}_runs",
        started_at=None,
        completed_at=None,
        client_metrics=_calculate_average_client_metrics(runs),
    )


def _aggregate_median(
    runs: List[WorkloadRunData], workload_type: str
) -> WorkloadRunData:
    """Calculate median metrics across runs.

    Args:
        runs: List of runs with metrics (includes both SUCCESS and FAILURE)
        workload_type: Type of workload

    Returns:
        WorkloadRunData with median metrics
    """
    if not runs:
        return None

    # Use first run as template
    template = runs[0]

    # Calculate median metrics
    median_metrics = _calculate_median_metrics(runs, workload_type)

    # Create aggregated run
    return WorkloadRunData(
        node_name=f"{template.node_name} (median of {len(runs)} runs)",
        workload_type=template.workload_type,
        params=template.params,
        metrics=median_metrics,
        status="SUCCEEDED",
        task_id=None,
        workflow_execution_id=f"aggregated_{len(runs)}_runs",
        started_at=None,
        completed_at=None,
        client_metrics=_calculate_median_client_metrics(runs),
    )


def _calculate_average_metrics(runs: List[WorkloadRunData], workload_type: str) -> Dict:
    """Calculate average of all numeric metrics.

    Args:
        runs: List of runs
        workload_type: Type of workload

    Returns:
        Dictionary with averaged metrics
    """
    if not runs:
        return {}

    averaged = {}

    # Handle TPCC metrics
    if "tpcc" in workload_type.lower():
        tpmc_values = [r.metrics.get("tpmC", 0) for r in runs if r.metrics.get("tpmC")]
        if tpmc_values:
            averaged["tpmC"] = sum(tpmc_values) / len(tpmc_values)

        tps_values = [r.metrics.get("tps", 0) for r in runs if r.metrics.get("tps")]
        if tps_values:
            averaged["tps"] = sum(tps_values) / len(tps_values)

    # Handle Sysbench metrics
    else:
        tps_values = [r.metrics.get("tps", 0) for r in runs if r.metrics.get("tps")]
        if tps_values:
            averaged["tps"] = sum(tps_values) / len(tps_values)

        qps_values = [r.metrics.get("qps", 0) for r in runs if r.metrics.get("qps")]
        if qps_values:
            averaged["qps"] = sum(qps_values) / len(qps_values)

    # Handle latency metrics (nested dict)
    latency_metrics = ["p95", "p99", "avg", "min", "max"]
    latency_dict = {}
    for metric in latency_metrics:
        values = [
            r.metrics.get("latency", {}).get(metric, 0)
            for r in runs
            if r.metrics.get("latency", {}).get(metric)
        ]
        if values:
            latency_dict[metric] = sum(values) / len(values)

    if latency_dict:
        averaged["latency"] = latency_dict

    return averaged


def _calculate_median_metrics(runs: List[WorkloadRunData], workload_type: str) -> Dict:
    """Calculate median of all numeric metrics.

    Args:
        runs: List of runs
        workload_type: Type of workload

    Returns:
        Dictionary with median metrics
    """
    if not runs:
        return {}

    median = {}

    def _median(values: List[float]) -> float:
        """Calculate median of a list of values."""
        sorted_values = sorted(values)
        n = len(sorted_values)
        if n % 2 == 0:
            return (sorted_values[n // 2 - 1] + sorted_values[n // 2]) / 2
        else:
            return sorted_values[n // 2]

    # Handle TPCC metrics
    if "tpcc" in workload_type.lower():
        tpmc_values = [r.metrics.get("tpmC", 0) for r in runs if r.metrics.get("tpmC")]
        if tpmc_values:
            median["tpmC"] = _median(tpmc_values)

        tps_values = [r.metrics.get("tps", 0) for r in runs if r.metrics.get("tps")]
        if tps_values:
            median["tps"] = _median(tps_values)

    # Handle Sysbench metrics
    else:
        tps_values = [r.metrics.get("tps", 0) for r in runs if r.metrics.get("tps")]
        if tps_values:
            median["tps"] = _median(tps_values)

        qps_values = [r.metrics.get("qps", 0) for r in runs if r.metrics.get("qps")]
        if qps_values:
            median["qps"] = _median(qps_values)

    # Handle latency metrics
    latency_metrics = ["p95", "p99", "avg", "min", "max"]
    latency_dict = {}
    for metric in latency_metrics:
        values = [
            r.metrics.get("latency", {}).get(metric, 0)
            for r in runs
            if r.metrics.get("latency", {}).get(metric)
        ]
        if values:
            latency_dict[metric] = _median(values)

    if latency_dict:
        median["latency"] = latency_dict

    return median


def _calculate_average_client_metrics(runs: List[WorkloadRunData]) -> Optional[Dict]:
    """Calculate average client metrics across runs.

    Args:
        runs: List of runs

    Returns:
        Dictionary with averaged client metrics, or None if no client metrics
    """
    runs_with_metrics = [r for r in runs if r.client_metrics]
    if not runs_with_metrics:
        return None

    averaged = {}

    # Get all metric keys from first run
    metric_keys = runs_with_metrics[0].client_metrics.keys()

    for key in metric_keys:
        values = [
            r.client_metrics.get(key, 0)
            for r in runs_with_metrics
            if r.client_metrics.get(key) is not None
        ]
        if values:
            averaged[key] = sum(values) / len(values)

    return averaged if averaged else None


def _calculate_median_client_metrics(runs: List[WorkloadRunData]) -> Optional[Dict]:
    """Calculate median client metrics across runs.

    Args:
        runs: List of runs

    Returns:
        Dictionary with median client metrics, or None if no client metrics
    """
    runs_with_metrics = [r for r in runs if r.client_metrics]
    if not runs_with_metrics:
        return None

    def _median(values: List[float]) -> float:
        sorted_values = sorted(values)
        n = len(sorted_values)
        if n % 2 == 0:
            return (sorted_values[n // 2 - 1] + sorted_values[n // 2]) / 2
        else:
            return sorted_values[n // 2]

    median_metrics = {}

    # Get all metric keys from first run
    metric_keys = runs_with_metrics[0].client_metrics.keys()

    for key in metric_keys:
        values = [
            r.client_metrics.get(key, 0)
            for r in runs_with_metrics
            if r.client_metrics.get(key) is not None
        ]
        if values:
            median_metrics[key] = _median(values)

    return median_metrics if median_metrics else None


def group_runs_by_thread_count(
    runs: List[WorkloadRunData],
) -> Dict[int, List[WorkloadRunData]]:
    """Group runs by thread count parameter.

    Args:
        runs: List of runs to group

    Returns:
        Dictionary mapping thread count -> list of runs
    """
    grouped = {}

    for run in runs:
        # Extract thread count from params
        thread_count = (
            run.params.get("thread")
            or run.params.get("threads")
            or run.params.get("concurrency")
        )

        if thread_count is None:
            continue

        if thread_count not in grouped:
            grouped[thread_count] = []

        grouped[thread_count].append(run)

    return grouped


if __name__ == "__main__":
    # Test aggregation logic
    print("Testing Aggregation Logic:\n")

    # Create sample runs
    sample_runs = [
        WorkloadRunData(
            node_name="tpcc_1k_100th",
            workload_type="bench-tpcc",
            params={"thread": 100, "warehouses": 1000},
            metrics={"tpmC": 7500, "tps": 125, "latency": {"p95": 100.5, "p99": 150.2}},
            status="SUCCEEDED",
            task_id=1,
            workflow_execution_id="exec1",
        ),
        WorkloadRunData(
            node_name="tpcc_1k_100th",
            workload_type="bench-tpcc",
            params={"thread": 100, "warehouses": 1000},
            metrics={"tpmC": 7800, "tps": 130, "latency": {"p95": 95.3, "p99": 145.8}},
            status="SUCCEEDED",
            task_id=2,
            workflow_execution_id="exec2",
        ),
        WorkloadRunData(
            node_name="tpcc_1k_100th",
            workload_type="bench-tpcc",
            params={"thread": 100, "warehouses": 1000},
            metrics={"tpmC": 7600, "tps": 127, "latency": {"p95": 98.7, "p99": 148.5}},
            status="SUCCEEDED",
            task_id=3,
            workflow_execution_id="exec3",
        ),
    ]

    # Test BEST aggregation
    best_run = aggregate_runs(sample_runs, AggregationStrategy.BEST, "bench-tpcc")
    print("Test 1: BEST aggregation")
    print(f"  Best tpmC: {best_run.metrics.get('tpmC')}")
    print(f"  Task ID: {best_run.task_id}")
    print()

    # Test AVERAGE aggregation
    avg_run = aggregate_runs(sample_runs, AggregationStrategy.AVERAGE, "bench-tpcc")
    print("Test 2: AVERAGE aggregation")
    print(f"  Average tpmC: {avg_run.metrics.get('tpmC'):.2f}")
    print(f"  Average P95: {avg_run.metrics.get('latency', {}).get('p95'):.2f}ms")
    print(f"  Node name: {avg_run.node_name}")
    print()

    # Test MEDIAN aggregation
    median_run = aggregate_runs(sample_runs, AggregationStrategy.MEDIAN, "bench-tpcc")
    print("Test 3: MEDIAN aggregation")
    print(f"  Median tpmC: {median_run.metrics.get('tpmC'):.2f}")
    print(f"  Median P95: {median_run.metrics.get('latency', {}).get('p95'):.2f}ms")
    print()

    # Test LATEST aggregation
    latest_run = aggregate_runs(sample_runs, AggregationStrategy.LATEST, "bench-tpcc")
    print("Test 4: LATEST aggregation")
    print(f"  Latest tpmC: {latest_run.metrics.get('tpmC')}")
    print(f"  Workflow ID: {latest_run.workflow_execution_id}")
    print()

    # Test grouping by thread count
    mixed_runs = sample_runs + [
        WorkloadRunData(
            node_name="tpcc_1k_200th",
            workload_type="bench-tpcc",
            params={"thread": 200, "warehouses": 1000},
            metrics={"tpmC": 8500, "tps": 142},
            status="SUCCEEDED",
            task_id=4,
            workflow_execution_id="exec4",
        ),
    ]

    grouped = group_runs_by_thread_count(mixed_runs)
    print("Test 5: Group by thread count")
    for thread_count, runs in sorted(grouped.items()):
        print(f"  Thread {thread_count}: {len(runs)} runs")

    print("\n✓ Aggregation tests completed!")