# src/openevalkit/core.py

"""Main evaluation function."""
from typing import List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError
import time

from openevalkit.config import EvalConfig
from openevalkit.dataset import Dataset
from openevalkit.scorers.base import Scorer
from openevalkit.judges.base import Judge
from openevalkit.result import EvaluationResult
from openevalkit.errors import ConfigError

# Optional: tqdm for progress bar
try:
    from tqdm import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False


def evaluate(
    dataset: Dataset,
    scorers: Optional[List[Scorer]] = None,
    judges: Optional[List[Judge]] = None,
    config: Optional[EvalConfig] = None
) -> EvaluationResult:
    """
    Evaluate a dataset with scorers and/or judges.

    Args:
        dataset: Dataset to evaluate
        scorers: Optional list of scorers to apply
        judges: Optional list of judges to apply
        config: Optional evaluation configuration
    
    Returns:
        EvaluationResult with scores, judgments, and aggregates
    
    Raises:
        ConfigError: If configuration is invalid
    
    Examples:
        >>> from openevalkit import Dataset, evaluate, EvalConfig
        >>> from openevalkit.scorers import ExactMatch
        >>> from openevalkit.judges import LLMJudge, LLMConfig, Rubric
        >>> 
        >>> dataset = Dataset.from_jsonl("data.jsonl")
        >>> 
        >>> # With scorers only
        >>> results = evaluate(dataset, scorers=[ExactMatch()])
        >>> 
        >>> # With judges only
        >>> rubric = Rubric(criteria=["helpfulness"], scale="0-1")
        >>> judge = LLMJudge(LLMConfig(model="gpt-4o"), rubric)
        >>> results = evaluate(dataset, judges=[judge])
        >>> 
        >>> # With both and custom config
        >>> results = evaluate(
        ...     dataset,
        ...     scorers=[ExactMatch()],
        ...     judges=[judge],
        ...     config=EvalConfig(
        ...         concurrency=10,
        ...         cache_enabled=True,
        ...         verbose=True,
        ...         seed=42
        ...     )
        ... )
    """
    # Use default config if not provided
    if config is None:
        config = EvalConfig()
    
    # Set seed for reproducibility
    if config.seed is not None:
        import random
        import numpy as np
        random.seed(config.seed)
        np.random.seed(config.seed)
        if config.verbose:
            print(f"  Seed: {config.seed}")
    
    # Initialize lists
    scorers = scorers or []
    judges = judges or []
    
    # Validate: at least one scorer or judge
    if not scorers and not judges:
        raise ConfigError("Must provide at least one scorer or judge")
    
    if len(dataset) == 0:
        raise ConfigError("Dataset is empty")
    
    # Warn if progress_bar requested but tqdm not available
    if config.progress_bar and not TQDM_AVAILABLE:
        import warnings
        warnings.warn(
            "progress_bar=True but tqdm not installed. "
            "Install with: pip install tqdm"
        )
        config.progress_bar = False
    
    # Initialize cache
    cache = None
    if config.cache_enabled:
        from openevalkit.cache import Cache
        
        cache = Cache(
            cache_dir=config.cache_dir,
            max_size_mb=config.cache_max_size_mb
        )
        
        # Auto-cleanup old entries
        if config.cache_max_age_days > 0:
            deleted = cache.clear_old(days=config.cache_max_age_days)
            if config.verbose and deleted > 0:
                print(f"  Cache: Cleared {deleted} old entries")
        
        if config.verbose:
            stats = cache.stats()
            print(f"  Cache: {stats['entries']} entries ({stats['size_mb']:.2f} MB)")
            print(f"  Cache location: {stats['db_path']}")
    
    if config.verbose:
        print(f"Starting evaluation:")
        print(f"  Dataset: {len(dataset)} runs")
        print(f"  Scorers: {[s.name for s in scorers]}")
        print(f"  Judges: {[j.name for j in judges]}")
        print(f"  Concurrency: {config.concurrency}")
        if config.timeout:
            print(f"  Timeout: {config.timeout}s per run")
    
    # Validate scorers that require reference
    for scorer in scorers:
        if scorer.requires_reference:
            missing = [run.id for run in dataset if run.reference is None]
            if missing:
                raise ConfigError(
                    f"{scorer.name} requires 'reference' field. "
                    f"Missing in {len(missing)} runs: {missing[:5]}...\n\n"
                    f"Solutions:\n"
                    f"  1. Add references to your dataset\n"
                    f"  2. Use scorers that don't need references"
                )
    
    start_time = time.time()
    failed_runs = []
    
    # Process runs (parallel or sequential based on config)
    if config.concurrency == 1:
        # Sequential execution
        all_scores, all_judgments, failed = _evaluate_sequential(
            dataset, scorers, judges, config, cache
        )
    else:
        # Parallel execution
        all_scores, all_judgments, failed = _evaluate_parallel(
            dataset, scorers, judges, config, cache
        )
    
    failed_runs.extend(failed)
    
    # Aggregate scores
    aggregates = {}
    for scorer in scorers:
        if scorer.name in all_scores:
            agg = scorer.aggregate(all_scores[scorer.name])
            aggregates.update(agg)
    
    # Aggregate judgments
    for judge in judges:
        if judge.name in all_judgments:
            agg = judge.aggregate(all_judgments[judge.name])
            # Prefix with judge name to avoid collisions
            aggregates.update({f"{judge.name}_{k}": v for k, v in agg.items()})
    
    elapsed = time.time() - start_time
    
    if config.verbose:
        print(f"\nEvaluation complete in {elapsed:.2f}s")
        print(f"  Successful: {len(dataset) - len(failed_runs)}/{len(dataset)}")
        if failed_runs:
            print(f"  Failed: {len(failed_runs)}")
    
    return EvaluationResult(
        scores=all_scores,
        judgments=all_judgments,
        aggregates=aggregates,
        failed_runs=failed_runs,
        metadata={
            "duration": elapsed,
            "num_runs": len(dataset),
            "concurrency": config.concurrency,
            "cache_enabled": config.cache_enabled,
            "seed": config.seed
        }
    )


def _evaluate_sequential(dataset, scorers, judges, config, cache=None):
    """Sequential evaluation (one run at a time)."""
    from openevalkit.cache import make_cache_key
    
    all_scores = {scorer.name: [] for scorer in scorers}
    all_judgments = {judge.name: [] for judge in judges}
    failed_runs = []
    
    # Progress bar
    iterator = tqdm(dataset, desc="Evaluating") if config.progress_bar else dataset
    
    for run in iterator:
        # Score with each scorer
        for scorer in scorers:
            try:
                # Check cache if enabled and scorer is cacheable
                if config.cache_enabled and scorer.cacheable and cache:
                    cache_key = make_cache_key(run, scorer)
                    cached_result = cache.get(cache_key)
                    
                    if cached_result is not None:
                        score_result = cached_result
                    else:
                        # Cache miss - compute and store
                        score_result = _run_with_timeout(
                            scorer.score,
                            config.timeout,
                            run=run
                        )
                        cache.set(cache_key, score_result)
                else:
                    # No caching
                    score_result = _run_with_timeout(
                        scorer.score,
                        config.timeout,
                        run=run
                    )
                
                all_scores[scorer.name].append(score_result)
                
            except TimeoutError:
                failed_runs.append((run.id, scorer.name, TimeoutError(f"Timeout after {config.timeout}s")))
                if config.fail_fast:
                    raise
            except Exception as e:
                failed_runs.append((run.id, scorer.name, e))
                if config.fail_fast:
                    raise
        
        # Judge with each judge
        for judge in judges:
            try:
                # Check cache if enabled and judge is cacheable
                if config.cache_enabled and judge.cacheable and cache:
                    cache_key = make_cache_key(run, judge)
                    cached_result = cache.get(cache_key)
                    
                    if cached_result is not None:
                        judgment = cached_result
                    else:
                        # Cache miss - compute and store
                        judgment = _run_with_timeout(
                            judge.judge,
                            config.timeout,
                            run,
                            judge.rubric
                        )
                        cache.set(cache_key, judgment)
                else:
                    # No caching
                    judgment = _run_with_timeout(
                        judge.judge,
                        config.timeout,
                        run,
                        judge.rubric
                    )
                
                all_judgments[judge.name].append(judgment)
                
            except TimeoutError:
                failed_runs.append((run.id, judge.name, TimeoutError(f"Timeout after {config.timeout}s")))
                if config.fail_fast:
                    raise
            except Exception as e:
                failed_runs.append((run.id, judge.name, e))
                if config.fail_fast:
                    raise
    
    return all_scores, all_judgments, failed_runs


def _evaluate_parallel(dataset, scorers, judges, config, cache=None):
    """Parallel evaluation (multiple runs at once)."""
    from openevalkit.cache import make_cache_key
    
    all_scores = {scorer.name: [None] * len(dataset) for scorer in scorers}
    all_judgments = {judge.name: [None] * len(dataset) for judge in judges}
    failed_runs = []
    
    def process_run(idx, run):
        """Process a single run with all scorers and judges."""
        run_scores = {}
        run_judgments = {}
        run_failures = []
        
        # Score with each scorer
        for scorer in scorers:
            try:
                # Check cache if enabled and scorer is cacheable
                if config.cache_enabled and scorer.cacheable and cache:
                    cache_key = make_cache_key(run, scorer)
                    cached_result = cache.get(cache_key)
                    
                    if cached_result is not None:
                        score_result = cached_result
                    else:
                        # Cache miss - compute and store
                        score_result = _run_with_timeout(
                            scorer.score,
                            config.timeout,
                            run=run
                        )
                        cache.set(cache_key, score_result)
                else:
                    # No caching
                    score_result = _run_with_timeout(
                        scorer.score,
                        config.timeout,
                        run=run
                    )
                
                run_scores[scorer.name] = score_result
                
            except TimeoutError:
                run_failures.append((run.id, scorer.name, TimeoutError(f"Timeout after {config.timeout}s")))
                if config.fail_fast:
                    raise
            except Exception as e:
                run_failures.append((run.id, scorer.name, e))
                if config.fail_fast:
                    raise
        
        # Judge with each judge
        for judge in judges:
            try:
                # Check cache if enabled and judge is cacheable
                if config.cache_enabled and judge.cacheable and cache:
                    cache_key = make_cache_key(run, judge)
                    cached_result = cache.get(cache_key)
                    
                    if cached_result is not None:
                        judgment = cached_result
                    else:
                        # Cache miss - compute and store
                        judgment = _run_with_timeout(
                            judge.judge,
                            config.timeout,
                            run,
                            judge.rubric
                        )
                        cache.set(cache_key, judgment)
                else:
                    # No caching
                    judgment = _run_with_timeout(
                        judge.judge,
                        config.timeout,
                        run,
                        judge.rubric
                    )
                
                run_judgments[judge.name] = judgment
                
            except TimeoutError:
                run_failures.append((run.id, judge.name, TimeoutError(f"Timeout after {config.timeout}s")))
                if config.fail_fast:
                    raise
            except Exception as e:
                run_failures.append((run.id, judge.name, e))
                if config.fail_fast:
                    raise
        
        return idx, run_scores, run_judgments, run_failures
    
    # Execute in parallel
    with ThreadPoolExecutor(max_workers=config.concurrency) as executor:
        futures = {
            executor.submit(process_run, idx, run): idx
            for idx, run in enumerate(dataset)
        }
        
        # Progress bar for completed futures
        if config.progress_bar:
            futures_iter = tqdm(
                as_completed(futures),
                total=len(dataset),
                desc="Evaluating"
            )
        else:
            futures_iter = as_completed(futures)
        
        for future in futures_iter:
            try:
                idx, run_scores, run_judgments, run_failures = future.result()
                
                # Store results in order
                for scorer_name, score in run_scores.items():
                    all_scores[scorer_name][idx] = score
                
                for judge_name, judgment in run_judgments.items():
                    all_judgments[judge_name][idx] = judgment
                
                failed_runs.extend(run_failures)
                
            except Exception as e:
                if config.fail_fast:
                    raise
                # If fail_fast is False, error already captured in run_failures
    
    return all_scores, all_judgments, failed_runs


def _run_with_timeout(func, timeout, *args, **kwargs):
    """Run function with optional timeout."""
    if timeout is None:
        return func(*args, **kwargs)
    
    # Use ThreadPoolExecutor for timeout
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(func, *args, **kwargs)
        try:
            return future.result(timeout=timeout)
        except TimeoutError:
            raise TimeoutError(f"Function {func.__name__} exceeded timeout of {timeout}s")