"""Benchmark runner for executing IR evaluation."""
import time
import json
from pathlib import Path
from typing import List, Dict, Any, Callable
from dataclasses import dataclass, asdict

from candlekeep.eval.metrics import (
    calculate_reciprocal_rank,
    calculate_ndcg,
    calculate_ndcg_graded,
    calculate_hit_rate,
    calculate_precision_at_k
)


@dataclass
class EvalQuery:
    query: str
    expected_sources: List[str]
    category: str
    difficulty: str
    graded_relevance: Dict[str, int] | None = None  # chunk_id -> grade (0-3)


@dataclass
class EvalResult:
    query: str
    category: str
    difficulty: str
    retrieved_sources: List[str]
    rr: float
    ndcg_5: float
    hit_rate_1: float
    hit_rate_5: float
    precision_5: float
    latency_ms: float
    tokens: int
    ndcg_5_graded: float = -1.0  # -1.0 = no graded data available
    retrieved_chunk_ids: List[str] | None = None  # source:chunk_index keys


class BenchmarkRunner:
    """Runs IR evaluation against a search function."""
    
    def __init__(self, search_fn: Callable):
        self.search_fn = search_fn
        
    def run_suite(self, queries: List[EvalQuery], k: int = 5) -> List[EvalResult]:
        results = []
        for q in queries:
            start_time = time.time()
            # The search_fn should return a list of objects with a .metadata['source'] attribute
            search_results = self.search_fn(q.query, k)
            latency = (time.time() - start_time) * 1000
            
            retrieved_sources = []
            retrieved_chunk_ids = []
            for r in search_results:
                source = r.metadata.get('source', '')
                # Normalize source path to match expected_sources format
                if 'tests/fixtures' in source:
                    source = source[source.index('tests/fixtures'):]
                retrieved_sources.append(source)
                # Build chunk-level ID for graded relevance
                chunk_idx = r.metadata.get('chunk_index', 0)
                retrieved_chunk_ids.append(f"{source}:{chunk_idx}")
            
            ground_truth = set(q.expected_sources)
            
            # Calculate tokens (rough estimate: chars / 4)
            total_chars = sum(len(r.text) for r in search_results)
            tokens = total_chars // 4

            # Graded nDCG (when annotations are available)
            ndcg_g = -1.0
            if q.graded_relevance:
                ndcg_g = calculate_ndcg_graded(retrieved_chunk_ids, q.graded_relevance, k)
            
            results.append(EvalResult(
                query=q.query,
                category=q.category,
                difficulty=q.difficulty,
                retrieved_sources=retrieved_sources,
                rr=calculate_reciprocal_rank(retrieved_sources, ground_truth),
                ndcg_5=calculate_ndcg(retrieved_sources, ground_truth, 5),
                hit_rate_1=calculate_hit_rate(retrieved_sources, ground_truth, 1),
                hit_rate_5=calculate_hit_rate(retrieved_sources, ground_truth, 5),
                precision_5=calculate_precision_at_k(retrieved_sources, ground_truth, 5),
                latency_ms=latency,
                tokens=tokens,
                ndcg_5_graded=ndcg_g,
                retrieved_chunk_ids=retrieved_chunk_ids,
            ))
        return results

    def summarize(self, results: List[EvalResult]) -> Dict[str, Any]:
        if not results:
            return {}

        from candlekeep.eval.statistics import bootstrap_ci

        mrr_scores = [r.rr for r in results]
        ndcg_scores = [r.ndcg_5 for r in results]
        hr5_scores = [r.hit_rate_5 for r in results]

        mrr_mean, mrr_lo, mrr_hi = bootstrap_ci(mrr_scores)
        ndcg_mean, ndcg_lo, ndcg_hi = bootstrap_ci(ndcg_scores)
        hr5_mean, hr5_lo, hr5_hi = bootstrap_ci(hr5_scores)

        summary = {
            "mrr": mrr_mean,
            "mrr_ci": [mrr_lo, mrr_hi],
            "avg_ndcg_5": ndcg_mean,
            "avg_ndcg_5_ci": [ndcg_lo, ndcg_hi],
            "avg_hit_rate_1": sum(r.hit_rate_1 for r in results) / len(results),
            "avg_hit_rate_5": hr5_mean,
            "avg_hit_rate_5_ci": [hr5_lo, hr5_hi],
            "avg_precision_5": sum(r.precision_5 for r in results) / len(results),
            "avg_latency_ms": sum(r.latency_ms for r in results) / len(results),
            "avg_tokens": sum(r.tokens for r in results) / len(results),
            "total_tokens": sum(r.tokens for r in results),
            "by_category": {},
            "total_queries": len(results)
        }

        # Graded nDCG (only when annotations are present)
        graded_scores = [r.ndcg_5_graded for r in results if r.ndcg_5_graded >= 0]
        if graded_scores:
            g_mean, g_lo, g_hi = bootstrap_ci(graded_scores)
            summary["avg_ndcg_5_graded"] = g_mean
            summary["avg_ndcg_5_graded_ci"] = [g_lo, g_hi]
        
        categories = set(r.category for r in results)
        for cat in categories:
            cat_results = [r for r in results if r.category == cat]
            cat_summary = {
                "mrr": sum(r.rr for r in cat_results) / len(cat_results),
                "avg_ndcg_5": sum(r.ndcg_5 for r in cat_results) / len(cat_results),
                "avg_hit_rate_5": sum(r.hit_rate_5 for r in cat_results) / len(cat_results),
                "count": len(cat_results)
            }
            cat_graded = [r.ndcg_5_graded for r in cat_results if r.ndcg_5_graded >= 0]
            if cat_graded:
                cat_summary["avg_ndcg_5_graded"] = sum(cat_graded) / len(cat_graded)
            summary["by_category"][cat] = cat_summary
            
        return summary

    def save_report(self, results: List[EvalResult], summary: Dict[str, Any], output_path: str):
        report = {
            "summary": summary,
            "details": [asdict(r) for r in results]
        }
        Path(output_path).write_text(json.dumps(report, indent=2))
        print(f"Report saved to {output_path}")
