"""
Benchmark Runner

Runs benchmark evaluations against the SAGE RAG system
and generates evaluation reports.
"""

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

from sage.evaluation.datasets import (
    BenchmarkDataset,
    BenchmarkQuestion,
    Difficulty,
    QuestionCategory,
)
from sage.evaluation.metrics import (
    EvaluationMetrics,
    MetricsCalculator,
)
from sage.rag import RAGResponse, SAGERAGPipeline
from sage.rag.validation import CitationExtractor

logger = logging.getLogger(__name__)


@dataclass
class QuestionResult:
    """Result of evaluating a single question."""

    question: BenchmarkQuestion
    response: RAGResponse | None
    metrics: EvaluationMetrics
    error: str | None = None

    @property
    def passed(self) -> bool:
        """Check if question passed evaluation."""
        return self.error is None and self.metrics.overall_score >= 0.7


@dataclass
class BenchmarkResult:
    """Complete benchmark evaluation result."""

    dataset_name: str
    timestamp: str = ""
    total_questions: int = 0
    questions_passed: int = 0
    questions_failed: int = 0
    questions_errored: int = 0

    # Aggregate metrics
    accuracy: float = 0.0
    avg_citation_f1: float = 0.0
    avg_retrieval_precision: float = 0.0
    avg_latency_ms: float = 0.0
    avg_tokens: float = 0.0

    # Detailed results
    results: list[QuestionResult] = field(default_factory=list)

    # Breakdown by category
    by_category: dict[str, dict] = field(default_factory=dict)
    by_difficulty: dict[str, dict] = field(default_factory=dict)

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()

    def calculate_aggregates(self) -> None:
        """Calculate aggregate metrics from results."""
        if not self.results:
            return

        self.total_questions = len(self.results)
        self.questions_passed = sum(1 for r in self.results if r.passed)
        self.questions_errored = sum(1 for r in self.results if r.error)
        self.questions_failed = self.total_questions - self.questions_passed - self.questions_errored

        # Calculate averages
        valid_results = [r for r in self.results if r.error is None]
        if valid_results:
            self.accuracy = sum(r.metrics.answer.correctness_score for r in valid_results) / len(valid_results)
            self.avg_citation_f1 = sum(r.metrics.citation.f1 for r in valid_results) / len(valid_results)
            self.avg_retrieval_precision = sum(r.metrics.retrieval.precision_at_k for r in valid_results) / len(valid_results)
            self.avg_latency_ms = sum(r.metrics.latency_ms for r in valid_results) / len(valid_results)
            self.avg_tokens = sum(r.metrics.tokens_used for r in valid_results) / len(valid_results)

        # Calculate breakdown by category
        self._calculate_breakdown()

    def _calculate_breakdown(self) -> None:
        """Calculate metrics breakdown by category and difficulty."""
        # By category
        for category in QuestionCategory:
            cat_results = [r for r in self.results if r.question.category == category and r.error is None]
            if cat_results:
                self.by_category[category.value] = {
                    "count": len(cat_results),
                    "accuracy": sum(r.metrics.answer.correctness_score for r in cat_results) / len(cat_results),
                    "passed": sum(1 for r in cat_results if r.passed),
                }

        # By difficulty
        for difficulty in Difficulty:
            diff_results = [r for r in self.results if r.question.difficulty == difficulty and r.error is None]
            if diff_results:
                self.by_difficulty[difficulty.value] = {
                    "count": len(diff_results),
                    "accuracy": sum(r.metrics.answer.correctness_score for r in diff_results) / len(diff_results),
                    "passed": sum(1 for r in diff_results if r.passed),
                }

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "dataset_name": self.dataset_name,
            "timestamp": self.timestamp,
            "summary": {
                "total_questions": self.total_questions,
                "questions_passed": self.questions_passed,
                "questions_failed": self.questions_failed,
                "questions_errored": self.questions_errored,
                "pass_rate": self.questions_passed / self.total_questions if self.total_questions > 0 else 0,
            },
            "metrics": {
                "accuracy": self.accuracy,
                "avg_citation_f1": self.avg_citation_f1,
                "avg_retrieval_precision": self.avg_retrieval_precision,
                "avg_latency_ms": self.avg_latency_ms,
                "avg_tokens": self.avg_tokens,
            },
            "by_category": self.by_category,
            "by_difficulty": self.by_difficulty,
            "detailed_results": [
                {
                    "question_id": r.question.id,
                    "question": r.question.question,
                    "passed": r.passed,
                    "error": r.error,
                    "metrics": r.metrics.to_dict() if r.metrics else None,
                }
                for r in self.results
            ],
        }


class BenchmarkRunner:
    """
    Runs benchmark evaluations against SAGE RAG.

    Evaluates a dataset of questions and computes
    comprehensive metrics.

    Example:
        runner = BenchmarkRunner(pipeline)
        result = await runner.run(dataset)
        print(f"Accuracy: {result.accuracy:.1%}")
    """

    def __init__(
        self,
        pipeline: SAGERAGPipeline,
        calculator: MetricsCalculator | None = None,
    ):
        """
        Initialize benchmark runner.

        Args:
            pipeline: RAG pipeline to evaluate
            calculator: Optional custom metrics calculator
        """
        self.pipeline = pipeline
        self.calculator = calculator or MetricsCalculator()
        self.citation_extractor = CitationExtractor()

    async def run(
        self,
        dataset: BenchmarkDataset,
        max_questions: int | None = None,
        parallel: int = 1,
    ) -> BenchmarkResult:
        """
        Run benchmark evaluation.

        Args:
            dataset: Benchmark dataset to evaluate
            max_questions: Maximum questions to evaluate
            parallel: Number of parallel evaluations (1 = sequential)

        Returns:
            BenchmarkResult with all metrics
        """
        logger.info(f"Starting benchmark: {dataset.name} ({len(dataset)} questions)")

        questions = list(dataset.questions)
        if max_questions:
            questions = questions[:max_questions]

        results = []

        if parallel > 1:
            # Run in parallel batches
            for i in range(0, len(questions), parallel):
                batch = questions[i:i + parallel]
                batch_results = await asyncio.gather(
                    *[self._evaluate_question(q) for q in batch],
                    return_exceptions=True
                )
                for q, r in zip(batch, batch_results, strict=True):
                    if isinstance(r, Exception):
                        results.append(QuestionResult(
                            question=q,
                            response=None,
                            metrics=EvaluationMetrics(),
                            error=str(r),
                        ))
                    else:
                        results.append(r)
        else:
            # Sequential evaluation
            for question in questions:
                try:
                    result = await self._evaluate_question(question)
                    results.append(result)
                except Exception as e:
                    logger.error(f"Error evaluating {question.id}: {e}")
                    results.append(QuestionResult(
                        question=question,
                        response=None,
                        metrics=EvaluationMetrics(),
                        error=str(e),
                    ))

        # Build result
        benchmark_result = BenchmarkResult(
            dataset_name=dataset.name,
            results=results,
        )
        benchmark_result.calculate_aggregates()

        logger.info(
            f"Benchmark complete: {benchmark_result.accuracy:.1%} accuracy, "
            f"{benchmark_result.questions_passed}/{benchmark_result.total_questions} passed"
        )

        return benchmark_result

    async def _evaluate_question(self, question: BenchmarkQuestion) -> QuestionResult:
        """Evaluate a single question."""
        logger.debug(f"Evaluating question: {question.id}")

        start_time = time.time()

        # Get RAG response
        response = await self.pipeline.query(
            question=question.question,
            validate=True,
        )

        latency_ms = (time.time() - start_time) * 1000

        # Calculate metrics
        metrics = EvaluationMetrics(
            query_id=question.id,
            query=question.question,
            latency_ms=latency_ms,
            tokens_used=response.total_tokens,
        )

        # Answer metrics
        metrics.answer = self.calculator.calculate_answer_metrics(
            generated=response.answer,
            expected=question.expected_answer,
            key_points=question.key_points,
        )

        # Citation metrics
        found_citations = self.citation_extractor.extract(response.answer)
        metrics.citation = self.calculator.calculate_citation_metrics(
            found_citations=found_citations,
            expected_citations=question.expected_citations,
        )

        # Retrieval metrics
        retrieved_docs = response.sources_used
        metrics.retrieval = self.calculator.calculate_retrieval_metrics(
            retrieved_doc_ids=retrieved_docs,
            expected_doc_ids=question.source_documents,
        )

        # Safety metrics
        metrics.safety = self.calculator.calculate_safety_metrics(
            response=response.answer,
            context_available=len(response.sources_used) > 0,
        )

        return QuestionResult(
            question=question,
            response=response,
            metrics=metrics,
        )

    async def run_single(self, question: BenchmarkQuestion) -> QuestionResult:
        """
        Evaluate a single question.

        Args:
            question: Question to evaluate

        Returns:
            QuestionResult
        """
        return await self._evaluate_question(question)


def save_benchmark_result(result: BenchmarkResult, path: Path) -> None:
    """
    Save benchmark result to JSON file.

    Args:
        result: Benchmark result to save
        path: Output path
    """
    with open(path, "w") as f:
        json.dump(result.to_dict(), f, indent=2)

    logger.info(f"Saved benchmark result to {path}")


def load_benchmark_result(path: Path) -> dict:
    """
    Load benchmark result from JSON file.

    Args:
        path: Path to result file

    Returns:
        Result dictionary
    """
    with open(path) as f:
        return json.load(f)
