"""
Evaluation and grading functionality.

This module provides tools for evaluating model outputs against reference answers.
"""

import json
from loguru import logger
from pathlib import Path
import re
from typing import Any, Callable, Dict, List, Optional

from langchain_core.messages.human import HumanMessage

from ..config.llm_config import get_llm
from .templates import BROWSECOMP_GRADER_TEMPLATE, SIMPLEQA_GRADER_TEMPLATE


# Default evaluation configuration using Claude 3.7 Sonnet via OpenRouter
DEFAULT_EVALUATION_CONFIG = {
    "model_name": "anthropic/claude-3.7-sonnet",  # Correct model ID for OpenRouter
    "provider": "openai_endpoint",  # Use OpenRouter
    "openai_endpoint_url": "https://openrouter.ai/api/v1",  # OpenRouter URL
    "temperature": 0,  # Zero temp for consistent evaluation
    # Note: max_tokens removed as it's not supported by LDR's get_llm()
}


def get_evaluation_llm(
    custom_config: Optional[Dict[str, Any]] = None,
    settings_snapshot: Optional[Dict[str, Any]] = None,
):
    """
    Get an LLM for evaluation purposes using Claude 3.7 Sonnet via OpenRouter
    by default, which can be overridden with custom settings.

    Args:
        custom_config: Optional custom configuration that overrides defaults
        settings_snapshot: Optional settings snapshot for thread-safe access

    Returns:
        An LLM instance for evaluation
    """
    # Start with default config (Claude 3.7 Sonnet via OpenRouter)
    config = DEFAULT_EVALUATION_CONFIG.copy()

    # Override with any custom settings
    if custom_config:
        config.update(custom_config)

    logger.info(
        f"Getting evaluation LLM with provider={config['provider']}, model={config['model_name']}"
    )

    # Remove any parameters that LDR's get_llm doesn't support
    # This ensures compatibility with LDR's implementation
    ldr_supported_params = {
        "model_name",
        "temperature",
        "provider",
        "openai_endpoint_url",
        "api_key",
    }

    filtered_config = {
        k: v for k, v in config.items() if k in ldr_supported_params
    }

    # Check if we're using openai_endpoint but don't have an API key configured
    if filtered_config.get("provider") == "openai_endpoint":
        # Try to get API key from settings snapshot or environment
        api_key = None

        if settings_snapshot:
            # Get from settings snapshot for thread safety
            api_key_setting = settings_snapshot.get(
                "llm.openai_endpoint.api_key"
            )
            if api_key_setting:
                api_key = (
                    api_key_setting.get("value")
                    if isinstance(api_key_setting, dict)
                    else api_key_setting
                )
        else:
            # No settings snapshot available
            logger.warning(
                "No settings snapshot provided for benchmark grader. "
                "API key must be provided via settings_snapshot for thread safety."
            )

        if not api_key:
            logger.warning(
                "Using openai_endpoint provider but no API key found. "
                "Set the llm.openai_endpoint.api_key setting in the database or "
                "LDR_LLM_OPENAI_ENDPOINT_API_KEY environment variable."
            )
            # Try to fall back to LDR's config if API key not explicitly provided
            # The get_llm function will handle this case

    # Get the LLM using LDR's existing function
    return get_llm(**filtered_config)


def extract_answer_from_response(
    response: str, dataset_type: str = "simpleqa"
) -> Dict[str, str]:
    """
    Extract structured information from LDR's response.

    Args:
        response: Response from LDR
        dataset_type: Type of dataset

    Returns:
        Dictionary with extracted answer and confidence
    """
    # Clean up citations
    response = re.sub(r"\[\d+\]", "", response)

    # Extract differently based on dataset type
    if dataset_type.lower() == "browsecomp":
        # Extract the final answer from structured response
        answer_match = re.search(r"Exact Answer:\s*(.*?)(?:\n|$)", response)
        exact_answer = answer_match.group(1).strip() if answer_match else "None"

        # Extract confidence
        confidence_match = re.search(r"Confidence:\s*(\d+)%", response)
        confidence = confidence_match.group(1) if confidence_match else "100"

        return {"extracted_answer": exact_answer, "confidence": confidence}

    # For SimpleQA, return the whole response as the answer
    return {
        "extracted_answer": response,
        "confidence": "100",  # SimpleQA doesn't have confidence scores
    }


def grade_single_result(
    result_data: Dict[str, Any],
    dataset_type: str = "simpleqa",
    evaluation_config: Optional[Dict[str, Any]] = None,
    settings_snapshot: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Grade a single benchmark result using LLM.

    Args:
        result_data: Dictionary containing result data with keys: id, problem, correct_answer, response, extracted_answer
        dataset_type: Type of dataset
        evaluation_config: Optional custom config for evaluation LLM
        settings_snapshot: Optional settings snapshot for thread-safe access

    Returns:
        Dictionary with grading results
    """
    # Get evaluation LLM
    evaluation_llm = get_evaluation_llm(evaluation_config, settings_snapshot)

    try:
        # Select appropriate template
        template = (
            BROWSECOMP_GRADER_TEMPLATE
            if dataset_type.lower() == "browsecomp"
            else SIMPLEQA_GRADER_TEMPLATE
        )

        question = result_data.get("problem", "")
        correct_answer = result_data.get("correct_answer", "")
        response = result_data.get("response", "")

        logger.info(f"Grading single result: {question[:50]}...")

        # Format grading prompt
        grading_prompt = template.format(
            question=question, correct_answer=correct_answer, response=response
        )

        import time

        eval_llm_start = time.time()
        logger.info(
            f"Starting grading LLM call (prompt length: {len(grading_prompt)} chars)..."
        )

        # Grade using LLM
        if hasattr(evaluation_llm, "invoke") and callable(
            evaluation_llm.invoke
        ):
            if hasattr(evaluation_llm, "chat_messages"):
                # Handle ChatOpenAI and similar models that use messages
                grading_response = evaluation_llm.invoke(
                    [HumanMessage(content=grading_prompt)]
                ).content
            else:
                # Handle other LLM types
                grading_response = evaluation_llm.invoke(grading_prompt)
                if hasattr(grading_response, "content"):
                    grading_response = grading_response.content
        else:
            # Fallback for other LLM interfaces
            grading_response = str(evaluation_llm(grading_prompt))

        eval_llm_elapsed = time.time() - eval_llm_start
        logger.info(f"Grading LLM call completed in {eval_llm_elapsed:.2f}s")

        # Extract grading information using regex
        if dataset_type.lower() == "browsecomp":
            # BrowseComp-specific extraction
            extracted_answer_match = re.search(
                r"extracted_final_answer:\s*(.*?)(?:\n|$)", grading_response
            )
            extracted_answer = (
                extracted_answer_match.group(1).strip()
                if extracted_answer_match
                else "None"
            )

            reasoning_match = re.search(
                r"reasoning:\s*(.*?)(?:\n\n|\ncorrect:|\Z)",
                grading_response,
                re.DOTALL,
            )
            reasoning = (
                reasoning_match.group(1).strip() if reasoning_match else ""
            )

            correct_match = re.search(
                r"correct:\s*(yes|no)", grading_response, re.IGNORECASE
            )
            is_correct = (
                (correct_match.group(1).lower() == "yes")
                if correct_match
                else False
            )

            confidence_match = re.search(
                r"confidence:\s*(\d+)", grading_response
            )
            confidence = (
                confidence_match.group(1) if confidence_match else "100"
            )
        else:
            # SimpleQA extraction
            extracted_answer_match = re.search(
                r"Extracted Answer:\s*(.*?)(?:\n|$)", grading_response
            )
            extracted_answer = (
                extracted_answer_match.group(1).strip()
                if extracted_answer_match
                else "None"
            )

            reasoning_match = re.search(
                r"Reasoning:\s*(.*?)(?:\nCorrect:|\Z)",
                grading_response,
                re.DOTALL,
            )
            reasoning = (
                reasoning_match.group(1).strip() if reasoning_match else ""
            )

            correct_match = re.search(
                r"Correct:\s*(yes|no)", grading_response, re.IGNORECASE
            )
            is_correct = (
                (correct_match.group(1).lower() == "yes")
                if correct_match
                else False
            )

            confidence = "100"  # SimpleQA doesn't have confidence

        # Format graded result
        return {
            "extracted_by_grader": extracted_answer,
            "reasoning": reasoning,
            "is_correct": is_correct,
            "graded_confidence": confidence,
            "grader_response": grading_response,
        }

    except Exception as e:
        logger.exception("Error grading single result")
        return {
            "grading_error": str(e),
            "is_correct": False,
            "graded_confidence": "0",
            "grader_response": f"Grading failed: {e!s}",
        }
    finally:
        from ..utilities.resource_utils import safe_close

        safe_close(evaluation_llm, "grader LLM")


def grade_results(
    results_file: str,
    output_file: str,
    dataset_type: str = "simpleqa",
    evaluation_config: Optional[Dict[str, Any]] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
) -> List[Dict[str, Any]]:
    """
    Grade benchmark results using LLM.

    Args:
        results_file: Path to results file
        output_file: Path to save graded results
        dataset_type: Type of dataset
        evaluation_config: Optional custom config for evaluation LLM
        progress_callback: Optional callback for progress updates

    Returns:
        List of graded results
    """
    # Get evaluation LLM
    evaluation_llm = get_evaluation_llm(evaluation_config)

    try:
        return _grade_results_inner(
            evaluation_llm,
            results_file,
            output_file,
            dataset_type,
            progress_callback,
        )
    finally:
        from ..utilities.resource_utils import safe_close

        safe_close(evaluation_llm, "grader LLM")


def _grade_results_inner(
    evaluation_llm,
    results_file: str,
    output_file: str,
    dataset_type: str,
    progress_callback: Optional[Callable[[int, int, Dict], None]],
) -> List[Dict[str, Any]]:
    """Inner implementation of grade_results, separated for cleanup."""
    # Select appropriate template
    template = (
        BROWSECOMP_GRADER_TEMPLATE
        if dataset_type.lower() == "browsecomp"
        else SIMPLEQA_GRADER_TEMPLATE
    )

    # Load results
    results = []
    with open(results_file, "r") as f:
        for line in f:
            if line.strip():
                results.append(json.loads(line))

    # Remove output file if it exists
    output_path = Path(output_file)
    if output_path.exists():
        output_path.unlink()

    graded_results = []
    correct_count = 0

    # Process each result
    for idx, result in enumerate(results):
        question = result.get("problem", "")
        correct_answer = result.get("correct_answer", "")
        response = result.get("response", "")

        # Call progress callback if provided
        if progress_callback:
            progress_callback(
                idx,
                len(results),
                {"status": "grading", "index": idx, "total": len(results)},
            )

        logger.info(f"Grading {idx + 1}/{len(results)}: {question[:50]}...")

        # Format grading prompt
        grading_prompt = template.format(
            question=question, correct_answer=correct_answer, response=response
        )

        try:
            # Grade using LLM
            if hasattr(evaluation_llm, "invoke") and callable(
                evaluation_llm.invoke
            ):
                if hasattr(evaluation_llm, "chat_messages"):
                    # Handle ChatOpenAI and similar models that use messages
                    grading_response = evaluation_llm.invoke(
                        [HumanMessage(content=grading_prompt)]
                    ).content
                else:
                    # Handle other LLM types
                    grading_response = evaluation_llm.invoke(grading_prompt)
                    if hasattr(grading_response, "content"):
                        grading_response = grading_response.content
            else:
                # Fallback for other LLM interfaces
                grading_response = str(evaluation_llm(grading_prompt))

            # Extract grading information using regex
            if dataset_type.lower() == "browsecomp":
                # BrowseComp-specific extraction
                extracted_answer_match = re.search(
                    r"extracted_final_answer:\s*(.*?)(?:\n|$)", grading_response
                )
                extracted_answer = (
                    extracted_answer_match.group(1).strip()
                    if extracted_answer_match
                    else "None"
                )

                reasoning_match = re.search(
                    r"reasoning:\s*(.*?)(?:\n\n|\ncorrect:|\Z)",
                    grading_response,
                    re.DOTALL,
                )
                reasoning = (
                    reasoning_match.group(1).strip() if reasoning_match else ""
                )

                correct_match = re.search(
                    r"correct:\s*(yes|no)", grading_response, re.IGNORECASE
                )
                is_correct = (
                    (correct_match.group(1).lower() == "yes")
                    if correct_match
                    else False
                )

                confidence_match = re.search(
                    r"confidence:\s*(\d+)", grading_response
                )
                confidence = (
                    confidence_match.group(1) if confidence_match else "100"
                )
            else:
                # SimpleQA extraction
                extracted_answer_match = re.search(
                    r"Extracted Answer:\s*(.*?)(?:\n|$)", grading_response
                )
                extracted_answer = (
                    extracted_answer_match.group(1).strip()
                    if extracted_answer_match
                    else "None"
                )

                reasoning_match = re.search(
                    r"Reasoning:\s*(.*?)(?:\nCorrect:|\Z)",
                    grading_response,
                    re.DOTALL,
                )
                reasoning = (
                    reasoning_match.group(1).strip() if reasoning_match else ""
                )

                correct_match = re.search(
                    r"Correct:\s*(yes|no)", grading_response, re.IGNORECASE
                )
                is_correct = (
                    (correct_match.group(1).lower() == "yes")
                    if correct_match
                    else False
                )

                confidence = "100"  # SimpleQA doesn't have confidence

            if is_correct:
                correct_count += 1

            # Format graded result
            graded_result = result.copy()
            graded_result.update(
                {
                    "extracted_by_grader": extracted_answer,
                    "reasoning": reasoning,
                    "is_correct": is_correct,
                    "graded_confidence": confidence,
                    "grader_response": grading_response,
                }
            )

            graded_results.append(graded_result)

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

            # Call progress callback if provided
            if progress_callback:
                progress_callback(
                    idx,
                    len(results),
                    {
                        "status": "graded",
                        "is_correct": is_correct,
                        "result": graded_result,
                    },
                )

        except Exception as e:
            logger.exception(f"Error grading result {idx + 1}")

            # Handle error
            error_result = result.copy()
            error_result["grading_error"] = str(e)

            with open(output_file, "a") as f:
                f.write(json.dumps(error_result) + "\n")

            graded_results.append(error_result)

            # Call progress callback if provided
            if progress_callback:
                progress_callback(
                    idx,
                    len(results),
                    {
                        "status": "error",
                        "error": str(e),
                        "result": error_result,
                    },
                )

    accuracy = correct_count / len(results) if results else 0
    logger.info(f"Grading complete. Accuracy: {accuracy:.3f}")
    logger.info(f"Correct: {correct_count}/{len(results)}")

    return graded_results


def human_evaluation(
    results_file: str, output_file: str, interactive: bool = True
) -> List[Dict[str, Any]]:
    """
    Allow for human evaluation of results.

    Args:
        results_file: Path to results file
        output_file: Path to save human-graded results
        interactive: Whether to run in interactive console mode

    Returns:
        List of human-graded results
    """
    # Load results
    results = []
    with open(results_file, "r") as f:
        for line in f:
            if line.strip():
                results.append(json.loads(line))

    # Remove output file if it exists
    output_path = Path(output_file)
    if output_path.exists():
        output_path.unlink()

    human_graded_results = []
    correct_count = 0

    if interactive:
        logger.info(f"Human evaluation: {len(results)} examples to grade")
        print(f"Human evaluation: {len(results)} examples to grade")
        print(
            "For each example, you'll see the question, correct answer, and model's response."
        )
        print("You'll be asked to judge if the model's answer is correct.")

    for idx, result in enumerate(results):
        question = result.get("problem", "")
        correct_answer = result.get("correct_answer", "")
        response = result.get("response", "")
        extracted_answer = result.get("extracted_answer", "")

        if interactive:
            print(f"\n\n===== Example {idx + 1}/{len(results)} =====")
            print(f"Question: {question}")
            print(f"\nCorrect Answer: {correct_answer}")
            print(f"\nModel Response: {response}")
            print(f"\nExtracted Answer: {extracted_answer}")

            # Get human judgment
            while True:
                judgment = (
                    input("\nIs the model's answer correct? (y/n): ")
                    .strip()
                    .lower()
                )
                if judgment in ["y", "n"]:
                    break
                print("Please enter 'y' or 'n'")

            is_correct = judgment == "y"

            # Get reasoning
            reasoning = input(
                "Please provide reasoning for your judgment: "
            ).strip()
        else:
            # Non-interactive mode - placeholder for API/UI implementation
            # In a real implementation, this would be filled by UI actions
            is_correct = False
            reasoning = "Non-interactive evaluation"

        if is_correct:
            correct_count += 1

        # Update result with human judgment
        human_result = result.copy()
        human_result.update(
            {
                "is_correct": is_correct,
                "reasoning": reasoning,
                "human_evaluation": True,
            }
        )

        human_graded_results.append(human_result)

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

    accuracy = correct_count / len(results) if results else 0
    logger.info(f"Human evaluation complete. Accuracy: {accuracy:.3f}")
    if interactive:
        print(f"\nHuman evaluation complete. Accuracy: {accuracy:.3f}")
        print(f"Correct: {correct_count}/{len(results)}")

    return human_graded_results
