"""
Benchmark runners for Local Deep Research.

This module provides the main functions for running benchmarks using LDR.
"""

import json
from loguru import logger
import os
from pathlib import Path
import time
from typing import Any, Callable, Dict, Optional

from ..api import quick_summary
from .datasets import DEFAULT_DATASET_URLS, load_dataset
from .datasets.base import DatasetRegistry
from .graders import extract_answer_from_response, grade_results
from .metrics import calculate_metrics, generate_report
from .templates import BROWSECOMP_QUERY_TEMPLATE


def format_query(question: str, dataset_type: str = "simpleqa") -> str:
    """
    Format query based on dataset type.

    Args:
        question: Original question
        dataset_type: Type of dataset

    Returns:
        Formatted query for LDR
    """
    if dataset_type.lower() == "browsecomp":
        # BrowseComp requires specific formatting
        return BROWSECOMP_QUERY_TEMPLATE.format(question=question)

    # Simple format for SimpleQA
    return question


def run_benchmark(
    dataset_type: str,
    dataset_path: Optional[str] = None,
    num_examples: Optional[int] = None,
    output_dir: str = "benchmark_results",
    run_evaluation: bool = True,
    evaluation_config: Optional[Dict[str, Any]] = None,
    search_config: Optional[Dict[str, Any]] = None,
    human_evaluation: bool = False,
    progress_callback: Optional[Callable[[str, int, Dict], None]] = None,
    seed: int = 42,
) -> Dict[str, Any]:
    """
    Run a benchmark on the specified dataset.

    Args:
        dataset_type: Type of dataset ("simpleqa" or "browsecomp")
        dataset_path: Optional custom dataset path
        num_examples: Number of examples to use
        output_dir: Directory to save results
        run_evaluation: Whether to evaluate results
        evaluation_config: Custom LLM config for evaluation
        search_config: Custom search parameters
        human_evaluation: Whether to use human evaluation
        progress_callback: Optional callback for progress updates
        seed: Random seed for reproducibility

    Returns:
        Dictionary with benchmark results and metrics
    """
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Default search configuration
    if not search_config:
        search_config = {
            "iterations": 3,
            "questions_per_iteration": 3,
            "search_tool": "searxng",
        }

    # Load dataset using the class-based approach
    try:
        # Create the dataset instance from registry
        dataset_instance = DatasetRegistry.create_dataset(
            dataset_id=dataset_type.lower(),
            dataset_path=dataset_path,
            num_examples=num_examples,
            seed=seed,
        )
        # Load the examples
        dataset = dataset_instance.load()

        logger.info(
            f"Loaded {len(dataset)} examples using dataset class {type(dataset_instance).__name__}"
        )
    except Exception as e:
        # Fallback to legacy function if there's any issue
        logger.warning(
            f"Error using dataset class: {e}. Falling back to legacy function."
        )
        dataset = load_dataset(
            dataset_type=dataset_type,
            dataset_path=dataset_path,
            num_examples=num_examples,
            seed=seed,
        )

    # Set up output files
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    results_file = str(
        Path(output_dir) / f"{dataset_type}_{timestamp}_results.jsonl"
    )
    evaluation_file = str(
        Path(output_dir) / f"{dataset_type}_{timestamp}_evaluation.jsonl"
    )
    report_file = str(
        Path(output_dir) / f"{dataset_type}_{timestamp}_report.md"
    )

    # Make sure output files don't exist
    for file in [results_file, evaluation_file, report_file]:
        file_path = Path(file)
        if file_path.exists():
            file_path.unlink()

    # Progress tracking
    total_examples = len(dataset)

    if progress_callback:
        progress_callback(
            "Starting benchmark",
            0,
            {
                "status": "started",
                "dataset_type": dataset_type,
                "total_examples": total_examples,
            },
        )

    # Process each example
    results = []

    for i, example in enumerate(dataset):
        # Extract question and answer in a way that uses the dataset class when available
        if "dataset_instance" in locals() and isinstance(
            dataset_instance,
            DatasetRegistry.get_dataset_class(dataset_type.lower()),
        ):
            # Use the dataset class methods to extract question and answer
            question = dataset_instance.get_question(example)
            correct_answer = dataset_instance.get_answer(example)
            logger.debug(
                "Using dataset class methods to extract question and answer"
            )
        else:
            # Fallback to the legacy approach
            if dataset_type.lower() == "simpleqa":
                question = example.get("problem", "")
                correct_answer = example.get("answer", "")
            else:  # browsecomp
                question = example.get("problem", "")
                # For BrowseComp, the answer should be in "correct_answer" after decryption
                correct_answer = example.get("correct_answer", "")
                if not correct_answer and "answer" in example:
                    # Fallback to "answer" field if "correct_answer" is not available
                    correct_answer = example.get("answer", "")

        # Update progress
        if progress_callback:
            progress_callback(
                f"Processing example {i + 1}/{total_examples}",
                int(i / total_examples * 50),
                {
                    "status": "processing",
                    "current": i + 1,
                    "total": total_examples,
                    "question": (
                        question[:50] + "..."
                        if len(question) > 50
                        else question
                    ),
                },
            )

        logger.info(f"Processing {i + 1}/{total_examples}: {question[:50]}...")

        try:
            # Format query based on dataset type
            formatted_query = format_query(question, dataset_type)

            # Time the search
            start_time = time.time()

            # Get response from LDR
            search_result = quick_summary(
                query=formatted_query,
                iterations=search_config.get("iterations", 3),
                questions_per_iteration=search_config.get(
                    "questions_per_iteration", 3
                ),
                search_tool=search_config.get("search_tool", "searxng"),
            )

            end_time = time.time()
            processing_time = end_time - start_time

            # Extract response and search info
            response = search_result.get("summary", "")

            # Extract structured information
            extracted = extract_answer_from_response(response, dataset_type)

            # Format result
            result = {
                "id": example.get("id", f"example_{i}"),
                "problem": question,
                "correct_answer": correct_answer,
                "response": response,
                "extracted_answer": extracted["extracted_answer"],
                "confidence": extracted["confidence"],
                "processing_time": processing_time,
                "sources": search_result.get("sources", []),
                "search_config": search_config,
            }

            # Add to results list
            results.append(result)

            # Write result to file
            with open(results_file, "a") as f:
                f.write(json.dumps(result) + "\n")

            # Update progress
            if progress_callback:
                progress_callback(
                    f"Completed example {i + 1}/{total_examples}",
                    int((i + 0.5) / total_examples * 50),
                    {
                        "status": "completed_example",
                        "current": i + 1,
                        "total": total_examples,
                        "result": result,
                    },
                )

        except Exception as e:
            logger.exception(f"Error processing example {i + 1}: {e!s}")

            # Create error result
            error_result = {
                "id": example.get("id", f"example_{i}"),
                "problem": question,
                "correct_answer": correct_answer,
                "error": str(e),
                "processing_time": (
                    time.time() - start_time if "start_time" in locals() else 0
                ),
            }

            # Add to results list
            results.append(error_result)

            # Write error result to file
            with open(results_file, "a") as f:
                f.write(json.dumps(error_result) + "\n")

            # Update progress
            if progress_callback:
                progress_callback(
                    f"Error processing example {i + 1}/{total_examples}",
                    int((i + 0.5) / total_examples * 50),
                    {
                        "status": "error",
                        "current": i + 1,
                        "total": total_examples,
                        "error": str(e),
                        "result": error_result,
                    },
                )

    logger.info(f"Completed processing {total_examples} examples")

    # Run evaluation if requested
    if run_evaluation:
        if progress_callback:
            progress_callback(
                "Starting evaluation",
                50,
                {"status": "evaluating", "results_file": results_file},
            )

        if human_evaluation:
            from .graders import human_evaluation as evaluate

            logger.info("Running human evaluation...")
            evaluation_results = evaluate(
                results_file=results_file,
                output_file=evaluation_file,
                interactive=True,
            )
        else:
            logger.info("Running automated evaluation...")
            try:
                evaluation_results = grade_results(
                    results_file=results_file,
                    output_file=evaluation_file,
                    dataset_type=dataset_type,
                    evaluation_config=evaluation_config,
                    progress_callback=lambda current, total, meta: (
                        progress_callback(
                            f"Evaluating {current + 1}/{total}",
                            50 + int((current + 0.5) / total * 40),
                            {**meta, "status": "evaluating"},
                        )
                        if progress_callback
                        else None
                    ),
                )
            except Exception as e:
                logger.exception(f"Automated evaluation failed: {e!s}")

                if progress_callback:
                    progress_callback(
                        "Automated evaluation failed. Falling back to human evaluation.",
                        60,
                        {"status": "evaluation_fallback", "error": str(e)},
                    )

                # Ask if user wants to fall back to human evaluation
                fallback_to_human = False
                print("\nAutomated evaluation failed with error:", str(e))
                response = input(
                    "Do you want to fall back to human evaluation? (y/n): "
                )
                fallback_to_human = response.strip().lower() == "y"

                if fallback_to_human:
                    logger.info("Falling back to human evaluation...")
                    from .graders import human_evaluation as evaluate

                    evaluation_results = evaluate(
                        results_file=results_file,
                        output_file=evaluation_file,
                        interactive=True,
                    )
                else:
                    logger.info("Skipping evaluation due to error.")
                    # Create an empty evaluation file to prevent issues
                    with open(evaluation_file, "w") as f:
                        f.write("")

                    return {
                        "status": "evaluation_error",
                        "dataset_type": dataset_type,
                        "results_path": results_file,
                        "evaluation_error": str(e),
                        "total_examples": total_examples,
                    }

        # Calculate metrics
        if progress_callback:
            progress_callback(
                "Calculating metrics", 90, {"status": "calculating_metrics"}
            )

        metrics = calculate_metrics(evaluation_file)

        # Generate report
        if progress_callback:
            progress_callback(
                "Generating report", 95, {"status": "generating_report"}
            )

        dataset_name = dataset_type.capitalize()
        report_path = generate_report(
            metrics=metrics,
            results_file=evaluation_file,
            output_file=report_file,
            dataset_name=dataset_name,
            config_info={
                "Dataset": dataset_path
                or DEFAULT_DATASET_URLS.get(dataset_type, "Unknown"),
                "Examples": total_examples,
                "Iterations": search_config.get("iterations", 3),
                "Questions per iteration": search_config.get(
                    "questions_per_iteration", 3
                ),
                "Search tool": search_config.get("search_tool", "searxng"),
                "Evaluation method": "Human"
                if human_evaluation
                else "Automated",
            },
        )

        # Mark as complete
        if progress_callback:
            progress_callback(
                "Benchmark complete",
                100,
                {
                    "status": "complete",
                    "metrics": metrics,
                    "report_path": report_path,
                },
            )

        return {
            "status": "complete",
            "dataset_type": dataset_type,
            "results_path": results_file,
            "evaluation_path": evaluation_file,
            "report_path": report_path,
            "metrics": metrics,
            "total_examples": total_examples,
            "accuracy": metrics.get("accuracy", 0),
        }

    else:
        # No evaluation, just return results
        if progress_callback:
            progress_callback(
                "Benchmark complete (no evaluation)",
                100,
                {"status": "complete_no_eval", "results_path": results_file},
            )

        return {
            "status": "complete_no_eval",
            "dataset_type": dataset_type,
            "results_path": results_file,
            "total_examples": total_examples,
        }


def run_simpleqa_benchmark(num_examples: int = 100, **kwargs) -> Dict[str, Any]:
    """
    Run SimpleQA benchmark with default settings.

    Args:
        num_examples: Number of examples to process
        **kwargs: Additional arguments to pass to run_benchmark

    Returns:
        Dictionary with benchmark results
    """
    return run_benchmark(
        dataset_type="simpleqa", num_examples=num_examples, **kwargs
    )


def run_browsecomp_benchmark(
    num_examples: int = 100, **kwargs
) -> Dict[str, Any]:
    """
    Run BrowseComp benchmark with default settings.

    Args:
        num_examples: Number of examples to process
        **kwargs: Additional arguments to pass to run_benchmark

    Returns:
        Dictionary with benchmark results
    """
    return run_benchmark(
        dataset_type="browsecomp", num_examples=num_examples, **kwargs
    )
