"""
Composite benchmark evaluator.

This module provides a composite evaluator that can run multiple benchmarks
with weighted scores to provide a comprehensive evaluation.
"""

from typing import Any, Dict, Optional

from loguru import logger

# Import specific evaluator implementations
from .browsecomp import BrowseCompEvaluator
from .simpleqa import SimpleQAEvaluator


class CompositeBenchmarkEvaluator:
    """
    Evaluator that combines multiple benchmarks with weighted scores.

    This evaluator runs multiple benchmark types and combines their scores
    according to specified weights, enabling comprehensive evaluation across
    different metrics and tasks.
    """

    def __init__(self, benchmark_weights: Optional[Dict[str, float]] = None):
        """
        Initialize with benchmark weights.

        Args:
            benchmark_weights: Dictionary mapping benchmark names to weights
                Default: {"simpleqa": 1.0}
        """
        # Default to SimpleQA only if no weights provided
        self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0}

        # Create evaluators for available benchmarks
        self.evaluators = {
            "simpleqa": SimpleQAEvaluator(),
            "browsecomp": BrowseCompEvaluator(),
        }

        # Normalize weights to sum to 1.0
        total_weight = sum(self.benchmark_weights.values())
        if total_weight <= 0:
            logger.warning(
                "Total benchmark weight is zero or negative. Using default weights."
            )
            self.normalized_weights = {"simpleqa": 1.0}
        else:
            self.normalized_weights = {
                k: w / total_weight for k, w in self.benchmark_weights.items()
            }

        # Log the weights being used
        logger.info(
            f"Using normalized benchmark weights: {self.normalized_weights}"
        )

    def evaluate(
        self,
        system_config: Dict[str, Any],
        num_examples: int,
        output_dir: str,
    ) -> Dict[str, Any]:
        """
        Run all requested benchmarks and compute weighted score.

        Args:
            system_config: Configuration parameters for the system under test
            num_examples: Number of benchmark examples to evaluate
            output_dir: Directory to save evaluation results

        Returns:
            Dictionary with combined metrics and individual benchmark results
        """
        all_results = {}
        combined_score = 0.0

        # Run each benchmark with weight > 0
        for benchmark_name, weight in self.normalized_weights.items():
            if weight > 0 and benchmark_name in self.evaluators:
                evaluator = self.evaluators[benchmark_name]

                try:
                    # Run benchmark evaluation
                    result = evaluator.evaluate(
                        system_config=system_config,
                        num_examples=num_examples,
                        output_dir=output_dir,
                    )

                    # Store individual results
                    all_results[benchmark_name] = result

                    # Calculate weighted contribution to combined score
                    quality_score = result.get("quality_score", 0.0)
                    weighted_contribution = quality_score * weight

                    logger.info(
                        f"Benchmark {benchmark_name}: score={quality_score:.4f}, "
                        f"weight={weight:.2f}, contribution={weighted_contribution:.4f}"
                    )

                    # Add to combined score
                    combined_score += weighted_contribution

                except Exception as e:
                    logger.exception(
                        f"Error running {benchmark_name} benchmark"
                    )
                    all_results[benchmark_name] = {
                        "benchmark_type": benchmark_name,
                        "error": str(e),
                        "quality_score": 0.0,
                    }

        # Return combined results
        return {
            "quality_score": combined_score,
            "benchmark_results": all_results,
            "benchmark_weights": self.normalized_weights,
            "combined_score": combined_score,
        }
