from dataclasses import asdict, dataclass, field
from os import getenv
from pathlib import Path
from typing import TYPE_CHECKING, Callable, List, Optional, Union
from uuid import uuid4

from pydantic import BaseModel, Field

from agno.agent import Agent, RunResponse
from agno.models.base import Model
from agno.utils.log import logger, set_log_level_to_debug, set_log_level_to_info

if TYPE_CHECKING:
    from rich.console import Console


class AccuracyAgentResponse(BaseModel):
    accuracy_score: int = Field(..., description="Accuracy Score between 1 and 10 assigned to the Agent's answer.")
    accuracy_reason: str = Field(..., description="Detailed reasoning for the accuracy score.")


@dataclass
class AccuracyEvaluation:
    question: str
    answer: str
    expected_answer: str
    score: int
    reason: str

    def print_eval(self, console: Optional["Console"] = None):
        from rich.box import ROUNDED
        from rich.console import Console
        from rich.markdown import Markdown
        from rich.table import Table

        if console is None:
            console = Console()

        results_table = Table(
            box=ROUNDED,
            border_style="blue",
            show_header=False,
            title="[ Evaluation Result ]",
            title_style="bold sky_blue1",
            title_justify="center",
        )
        results_table.add_row("Question", self.question)
        results_table.add_row("Answer", self.answer)
        results_table.add_row("Expected Answer", self.expected_answer)
        results_table.add_row("Accuracy Score", f"{str(self.score)}/10")
        results_table.add_row("Accuracy Reason", Markdown(self.reason))
        console.print(results_table)


@dataclass
class AccuracyResult:
    results: List[AccuracyEvaluation] = field(default_factory=list)
    avg_score: float = field(init=False)
    mean_score: float = field(init=False)
    min_score: float = field(init=False)
    max_score: float = field(init=False)
    std_dev_score: float = field(init=False)

    def __post_init__(self):
        self.compute_stats()

    def compute_stats(self):
        import statistics

        if self.results and len(self.results) > 0:
            _results = [r.score for r in self.results]
            self.avg_score = statistics.mean(_results)
            self.mean_score = statistics.mean(_results)
            self.min_score = min(_results)
            self.max_score = max(_results)
            self.std_dev_score = statistics.stdev(_results) if len(_results) > 1 else 0

    def print_summary(self, console: Optional["Console"] = None):
        from rich.box import ROUNDED
        from rich.console import Console
        from rich.table import Table

        if console is None:
            console = Console()

        summary_table = Table(
            box=ROUNDED,
            border_style="blue",
            show_header=False,
            title="[ Evaluation Summary ]",
            title_style="bold sky_blue1",
            title_justify="center",
        )
        summary_table.add_row("Number of Runs", f"{len(self.results)}")
        summary_table.add_row("Average Score", f"{self.avg_score:.2f}")
        summary_table.add_row("Mean Score", f"{self.mean_score:.2f}")
        summary_table.add_row("Minimum Score", f"{self.min_score:.2f}")
        summary_table.add_row("Maximum Score", f"{self.max_score:.2f}")
        summary_table.add_row("Standard Deviation", f"{self.std_dev_score:.2f}")
        console.print(summary_table)

    def print_results(self, console: Optional["Console"] = None):
        from rich.box import ROUNDED
        from rich.console import Console
        from rich.table import Table

        if console is None:
            console = Console()

        results_table = Table(
            box=ROUNDED,
            border_style="blue",
            show_header=False,
            title="[ Evaluation Result ]",
            title_style="bold sky_blue1",
            title_justify="center",
        )
        for result in self.results:
            results_table.add_row("Question", result.question)
            results_table.add_row("Answer", result.answer)
            results_table.add_row("Expected Answer", result.expected_answer)
            results_table.add_row("Accuracy Score", f"{str(result.score)}/10")
            results_table.add_row("Accuracy Reason", result.reason)
        console.print(results_table)


@dataclass
class AccuracyEval:
    """Evaluate the accuracy of an agent's answer."""

    # Evaluation name
    name: Optional[str] = None
    # Evaluation UUID (autogenerated if not set)
    eval_id: Optional[str] = None

    # Model used to evaluate the answer
    model: Optional[Model] = None

    # Evaluate an Agent
    agent: Optional[Agent] = None
    # Question to evaluate (can also be provided with the run method)
    question: Optional[Union[str, Callable]] = None
    # Answer to evaluate (can also be provided with the run method)
    answer: Optional[Union[str, Callable]] = None
    # Expected Answer for the question (can also be provided with the run method)
    expected_answer: Optional[Union[str, Callable]] = None

    evaluator_agent: Optional[Agent] = None
    # Guidelines for the evaluator agent
    evaluator_guidelines: Optional[List[str]] = None
    # Additional context to the evaluator agent
    evaluator_context: Optional[str] = None

    # Number of iterations to run
    num_iterations: int = 3
    # Result of the evaluation
    result: Optional[AccuracyResult] = None

    # Print summary of results
    print_summary: bool = False
    # Print detailed results
    print_results: bool = False
    # Save the result to a file
    save_result_to_file: Optional[str] = None

    # debug_mode=True enables debug logs
    debug_mode: bool = False

    def set_eval_id(self) -> str:
        if self.eval_id is None:
            self.eval_id = str(uuid4())
        logger.debug(f"*********** Evaluation ID: {self.eval_id} ***********")
        return self.eval_id

    def set_debug_mode(self) -> None:
        if self.debug_mode or getenv("AGNO_DEBUG", "false").lower() == "true":
            self.debug_mode = True
            set_log_level_to_debug()
            logger.debug("Debug logs enabled")
        else:
            set_log_level_to_info()

    def get_evaluator_agent(self, question: str, expected_answer: str) -> Agent:
        if self.evaluator_agent is not None:
            return self.evaluator_agent

        model = self.model
        if model is None:
            try:
                from agno.models.openai import OpenAIChat

                model = OpenAIChat(id="gpt-4o-mini")
            except (ModuleNotFoundError, ImportError) as e:
                logger.exception(e)
                logger.error(
                    "Agno uses `openai` as the default model provider. Please run `pip install openai` to use the default evaluator."
                )
                exit(1)

        evaluator_guidelines = ""
        if self.evaluator_guidelines is not None and len(self.evaluator_guidelines) > 0:
            evaluator_guidelines = "\n## Guidelines for the Agent's answer:\n"
            evaluator_guidelines += "\n- ".join(self.evaluator_guidelines)
            evaluator_guidelines += "\n"

        evaluator_context = ""
        if self.evaluator_context is not None and len(self.evaluator_context) > 0:
            evaluator_context = "## Additional Context:\n"
            evaluator_context += self.evaluator_context
            evaluator_context += "\n"

        return Agent(
            model=OpenAIChat(id="gpt-4o-mini"),
            description=f"""\
You are an Agent Evaluator tasked with assessing the accuracy of an AI Agent's answer compared to an expected answer for a given question.
Your task is to provide a detailed analysis and assign a score on a scale of 1 to 10, where 10 indicates a perfect match to the expected answer.

## Question:
{question}

## Expected Answer:
{expected_answer}

## Evaluation Criteria:
1. Accuracy of information
2. Completeness of the answer
3. Relevance to the question
4. Use of key concepts and ideas
5. Overall structure and clarity of presentation
{evaluator_guidelines}{evaluator_context}
## Instructions:
1. Carefully compare the AI Agent's answer to the expected answer.
2. Provide a detailed analysis, highlighting:
   - Specific similarities and differences
   - Key points included or missed
   - Any inaccuracies or misconceptions
3. Explicitly reference the evaluation criteria and any provided guidelines in your reasoning.
4. Assign a score from 1 to 10 (use only whole numbers) based on the following scale:
   1-2: Completely incorrect or irrelevant
   3-4: Major inaccuracies or missing crucial information
   5-6: Partially correct, but with significant omissions or errors
   7-8: Mostly accurate and complete, with minor issues
   9-10: Highly accurate and complete, matching the expected answer closely

Your evaluation should be objective, thorough, and well-reasoned. Provide specific examples from both answers to support your assessment.""",
            response_model=AccuracyAgentResponse,
            structured_outputs=True,
        )

    def get_question_to_evaluate(self, question: Optional[Union[str, Callable]] = None) -> Optional[str]:
        """Get the question to evaluate."""
        try:
            # Get question from the run method
            if question is not None:
                if isinstance(question, str):
                    return question
                elif callable(question):
                    _question = question()
                    if isinstance(_question, str):
                        return _question
                    else:
                        logger.error("Question is not a string")
                else:
                    logger.error("Question is not a string or callable")

            # Get the question from the eval
            if self.question is not None:
                if isinstance(self.question, str):
                    return self.question
                elif callable(self.question):
                    _question = self.question()
                    if isinstance(_question, str):
                        return _question
                    else:
                        logger.error("Question is not a string")
                else:
                    logger.error("Question is not a string or callable")
        except Exception as e:
            logger.error(f"Failed to get question to evaluate: {e}")
        return None

    def get_answer_to_evaluate(
        self, question: str, answer: Optional[Union[str, Callable]] = None
    ) -> Optional[RunResponse]:
        """Get the answer to evaluate.

        Priority:
        1. Answer provided with the run method
        2. Answer provided with the eval
        3. Answer from the agent
        """
        try:
            # Get answer from the run method
            if answer is not None:
                if isinstance(answer, str):
                    return RunResponse(content=answer)
                elif callable(answer):
                    _answer = answer()
                    if isinstance(_answer, str):
                        return RunResponse(content=_answer)
                    else:
                        logger.error("Answer is not a string")
                else:
                    logger.error("Answer is not a string or callable")

            # Get answer from the eval
            if self.answer is not None:
                if isinstance(self.answer, str):
                    return RunResponse(content=self.answer)
                elif callable(self.answer):
                    _answer = self.answer()
                    if isinstance(_answer, str):
                        return RunResponse(content=_answer)
                    else:
                        logger.error("Answer is not a string")
                else:
                    logger.error("Answer is not a string or callable")

            # Get answer from the agent
            if self.agent is not None and question is not None:
                logger.debug("Getting answer from agent")
                return self.agent.run(question)
        except Exception as e:
            logger.error(f"Failed to get answer to evaluate: {e}")
        return None

    def get_expected_answer_to_evaluate(self, expected_answer: Optional[Union[str, Callable]] = None) -> Optional[str]:
        """Get the expected answer to evaluate."""
        try:
            # Get expected_answer from the run method
            if expected_answer is not None:
                if isinstance(expected_answer, str):
                    return expected_answer
                elif callable(expected_answer):
                    _expected_answer = expected_answer()
                    if isinstance(_expected_answer, str):
                        return _expected_answer
                    else:
                        logger.error("Expected Answer is not a string")
                else:
                    logger.error("Expected Answer is not a string or callable")

            # Get the expected_answer from the eval
            if self.expected_answer is not None:
                if isinstance(self.expected_answer, str):
                    return self.expected_answer
                elif callable(self.expected_answer):
                    _expected_answer = self.expected_answer()
                    if isinstance(_expected_answer, str):
                        return _expected_answer
                    else:
                        logger.error("Expected Answer is not a string")
                else:
                    logger.error("Expected Answer is not a string or callable")
        except Exception as e:
            logger.error(f"Failed to get expected answer to evaluate: {e}")
        return None

    def run(
        self,
        *,
        question: Optional[Union[str, Callable]] = None,
        expected_answer: Optional[Union[str, Callable]] = None,
        answer: Optional[Union[str, Callable]] = None,
        print_summary: bool = True,
        print_results: bool = True,
    ) -> Optional[AccuracyResult]:
        from rich.console import Console
        from rich.live import Live
        from rich.status import Status

        self.set_eval_id()
        self.set_debug_mode()
        self.result = AccuracyResult()
        self.print_results = print_results
        self.print_summary = print_summary

        question_to_evaluate: Optional[str] = self.get_question_to_evaluate(question=question)
        if question_to_evaluate is None:
            logger.error("No Question to evaluate.")
            return None

        expected_answer_to_evaluate: Optional[str] = self.get_expected_answer_to_evaluate(
            expected_answer=expected_answer
        )
        if expected_answer_to_evaluate is None:
            logger.error("No Expected Answer to evaluate.")
            return None

        logger.debug(f"************ Evaluation Start: {self.eval_id} ************")
        logger.debug(f"Question: {question_to_evaluate}")
        logger.debug(f"Expected Answer: {expected_answer_to_evaluate}")
        logger.debug("***********************************************************")

        evaluator_agent: Agent = self.get_evaluator_agent(
            question=question_to_evaluate, expected_answer=expected_answer_to_evaluate
        )

        # Add a spinner while running the evaluations
        console = Console()
        with Live(console=console, transient=True) as live_log:
            for i in range(self.num_iterations):
                status = Status(f"Running evaluation {i + 1}...", spinner="dots", speed=1.0, refresh_per_second=10)
                live_log.update(status)

                answer_to_evaluate: Optional[RunResponse] = self.get_answer_to_evaluate(
                    question=question_to_evaluate, answer=answer
                )
                if answer_to_evaluate is None:
                    logger.error("No Answer to evaluate.")
                    continue

                try:
                    logger.debug(f"Answer #{i + 1}: {answer_to_evaluate.content}")
                    accuracy_agent_response = evaluator_agent.run(answer_to_evaluate.content).content
                    if accuracy_agent_response is None or not isinstance(
                        accuracy_agent_response, AccuracyAgentResponse
                    ):
                        logger.error("Evaluator Agent returned an invalid response")
                        continue

                    accuracy_evaluation = AccuracyEvaluation(
                        question=question_to_evaluate,
                        answer=answer_to_evaluate.content,  # type: ignore
                        expected_answer=expected_answer_to_evaluate,
                        score=accuracy_agent_response.accuracy_score,
                        reason=accuracy_agent_response.accuracy_reason,
                    )
                    if self.print_results:
                        accuracy_evaluation.print_eval(console)
                    self.result.results.append(accuracy_evaluation)
                    self.result.compute_stats()
                    status.update(f"Running evaluation {i + 1}... Done")
                except Exception as e:
                    logger.exception(f"Failed to evaluate accuracy, run #{i + 1}: {e}")
                    return None

                status.stop()

        # -*- Save result to file if save_result_to_file is set
        if self.save_result_to_file is not None and self.result is not None:
            try:
                import json

                fn_path = Path(self.save_result_to_file.format(name=self.name, eval_id=self.eval_id))
                if not fn_path.parent.exists():
                    fn_path.parent.mkdir(parents=True, exist_ok=True)
                fn_path.write_text(json.dumps(asdict(self.result), indent=4))
            except Exception as e:
                logger.warning(f"Failed to save result to file: {e}")

        # Show results
        if self.print_summary or self.print_results:
            self.result.print_summary(console)

        logger.debug(f"*********** Evaluation End: {self.eval_id} ***********")
        return self.result
