from __future__ import annotations

import itertools
import random
from typing import Any, Callable, Set, Type, get_type_hints

from pydantic import BaseModel
from tenacity import RetryError

from ragelo.evaluators.base_evaluator import BaseEvaluator
from ragelo.llm_providers.base_llm_provider import BaseLLMProvider, get_llm_provider
from ragelo.logger import logger
from ragelo.types.configurations import BaseAnswerEvaluatorConfig, PairwiseEvaluatorConfig
from ragelo.types.evaluables import AgentAnswer, Document, Evaluable, PairwiseGame
from ragelo.types.experiment import Experiment
from ragelo.types.formats import LLMInputPrompt
from ragelo.types.query import Query
from ragelo.types.results import AnswerEvaluatorResult
from ragelo.types.types import AnswerEvaluatorTypes
from ragelo.utils import call_async_fn


class BaseAnswerEvaluator(BaseEvaluator):
    config: BaseAnswerEvaluatorConfig
    evaluable_name: str = "Agent Answer"
    _warned_queries: Set[str] = set()

    def evaluate(
        self,
        query: Query | str,
        answer: AgentAnswer | str | None = None,
        answer_a: AgentAnswer | str | None = None,
        answer_b: AgentAnswer | str | None = None,
        retrieved_documents: list[Document | str] | str | None = None,
        document_metadata: list[dict[str, Any]] | None = None,
        query_metadata: dict[str, Any] | None = None,
        answer_metadata: dict[str, Any] | None = None,
        answer_a_metadata: dict[str, Any] | None = None,
        answer_b_metadata: dict[str, Any] | None = None,
    ) -> AnswerEvaluatorResult:
        """
        Evaluates a single sample using an answer evaluator.
        Args:
            query(Query|str): The query to evaluate.
                If a string is provided, a Query object will be created with the provided query_metadata.
            answer(Optional[AgentAnswer|str]): The answer generated by an agent to evaluate.
                If a string is provided, the AgentAnswer object will be created with the provided answer_metadata
            answer_a(Optional[AgentAnswer|str]): The first answer generated by an agent to evaluate.
                Only used if the evaluator is pairwise.
            answer_b(Optional[AgentAnswer|str]): The second answer generated by an agent to evaluate.
                Only used if the evaluator is pairwise.
            retrieved_documents(Optional[list[str]|list[Document]]): The documents retrieved by the agent.
            document_metadata(Optional[list[dict[str, Any]]): The metadata of the documents retrieved by the agent.
            query_metadata(Optional[dict[str, Any]]): The metadata for the query.
            answer_metadata(Optional[dict[str, Any]]): The metadata for the answer.
            answer_a_metadata(Optional[dict[str, Any]]): The metadata for the first answer.
                Only used if the evaluator is pairwise.
            answer_b_metadata(Optional[dict[str, Any]]): The metadata for the second answer.
                Only used if the evaluator is pairwise.
        Returns:
            AnswerEvaluatorResult: The result of the evaluation.
        """
        query = Query.assemble_query(query, query_metadata)

        if isinstance(retrieved_documents, str):
            retrieved_documents = [retrieved_documents]

        if retrieved_documents:
            retrieved_and_assembled_docs = Document.assemble_documents(
                retrieved_documents, query.qid, document_metadata
            )
            query.retrieved_docs = retrieved_and_assembled_docs

        agent: str | tuple[str, str]
        evaluable: AgentAnswer | PairwiseGame

        if self.config.pairwise:
            if not answer_a or not answer_b:
                raise ValueError("Pairwise evaluations require two answers")
            answer_a = AgentAnswer.assemble_answer(answer_a, query.qid, metadata=answer_a_metadata)
            answer_b = AgentAnswer.assemble_answer(answer_b, query.qid, metadata=answer_b_metadata)
            agent = (answer_a.agent, answer_b.agent)
            evaluable = PairwiseGame(
                qid=query.qid,
                agent_a_answer=answer_a,
                agent_b_answer=answer_b,
            )
        else:
            if not answer:
                raise ValueError("Pointwise evaluations require an answer")
            evaluable = AgentAnswer.assemble_answer(answer, query.qid, metadata=answer_metadata)
            agent = evaluable.agent
        result = call_async_fn(self.evaluate_async, (query, evaluable))

        if result.exception or result.raw_answer is None or result.answer is None:
            raise ValueError(
                f"Failed to evaluate qid: {query.qid} agent(s): {agent}",
                f"Exception: {result.exception}",
            )
        return result

    def evaluate_experiment(self, experiment: Experiment, n_threads: int | None = None):
        """
        Trigger the evaluation of all the queries in the experiment.
        The evaluation is done in asynchronously with the number of threads defined in config.n_processes parameter.
        This can be overwritten by the n_threads parameter.

        Args:
            experiment(Experiment): The experiment to evaluate.
            n_threads(int): The number of threads to use for the evaluation.
                If None, the number of threads defined in the config will be used.
        """
        if self.config.pairwise:
            self.__add_pairwise_games(experiment)
        super().evaluate_experiment(experiment, n_threads)

    async def evaluate_async(self, eval_sample: tuple[Query, Evaluable]) -> AnswerEvaluatorResult:
        """
        Evaluates a single sample (either an answer or a pairwise game) asynchronously.
        Args:
            eval_sample (tuple[Query, Evaluable]): The query and evaluable to evaluate.
        """
        agent: str | tuple[str, str]
        query, evaluable = eval_sample

        if not isinstance(evaluable, AgentAnswer) and not isinstance(evaluable, PairwiseGame):
            type_name = type(evaluable).__name__
            raise ValueError(f"can't evaluate a {type_name} in an Answer Evaluator")

        if evaluable.evaluation is not None and not self.config.force:
            if isinstance(evaluable, AgentAnswer):
                return AnswerEvaluatorResult(
                    qid=query.qid,
                    agent=evaluable.agent,
                    raw_answer=evaluable.evaluation.raw_answer,
                    answer=evaluable.evaluation.answer,
                    pairwise=False,
                )
            return AnswerEvaluatorResult(
                qid=query.qid,
                agent_a=evaluable.agent_a_answer.agent,
                agent_b=evaluable.agent_b_answer.agent,
                raw_answer=evaluable.evaluation.raw_answer,
                answer=evaluable.evaluation.answer,
                pairwise=True,
            )

        if isinstance(evaluable, AgentAnswer):
            agent = evaluable.agent
            prompt = self._build_message(query, evaluable)
        elif isinstance(evaluable, PairwiseGame):
            agent = (evaluable.agent_a_answer.agent, evaluable.agent_b_answer.agent)
            prompt = self._build_message_pairwise(query, evaluable)
        else:
            raise ValueError(f"Unknown evaluable type {type(evaluable)}")

        exc = None
        try:
            llm_response = await self.llm_provider.call_async(
                input=prompt,
                response_schema=self.config.llm_response_schema,
            )
            llm_response = self._process_answer(llm_response)
        except ValueError as e:
            logger.warning(
                f"Failed to PARSE answer for qid: {query.qid} agent(s): {agent}\nRaw answer: {llm_response.raw_answer}"
            )
            exc = str(e)
        except Exception as e:
            logger.warning(f"Failed to FETCH answers for qid: {query.qid}")
            logger.warning(f"agent(s): {agent}")
            if isinstance(e, RetryError):
                exc = str(e.last_attempt.exception())
            else:
                exc = str(e)
        if isinstance(evaluable, AgentAnswer):
            return AnswerEvaluatorResult(
                qid=query.qid,
                agent=evaluable.agent,
                raw_answer=llm_response.raw_answer,
                answer=llm_response.parsed_answer,
                pairwise=False,
                exception=exc,
            )
        assert isinstance(evaluable, PairwiseGame)
        return AnswerEvaluatorResult(
            qid=query.qid,
            agent_a=evaluable.agent_a_answer.agent,
            agent_b=evaluable.agent_b_answer.agent,
            raw_answer=llm_response.raw_answer,
            answer=llm_response.parsed_answer,
            pairwise=True,
            exception=exc,
        )

    def _get_tuples_to_evaluate(self, experiment: Experiment) -> list[tuple[Query, Evaluable]]:
        """
        Creates the list of pairs (query, evaluable) to evaluate
        """
        tuples_to_eval: list[tuple[Query, Evaluable]] = []
        all_tuples = 0
        missing_evaluations = 0
        for q in experiment:
            if self.config.pairwise:
                for g in q.pairwise_games:
                    all_tuples += 1
                    tuples_to_eval.append((q, g))
                    if g.evaluation is None:
                        missing_evaluations += 1

            else:
                for a in q.answers.values():
                    all_tuples += 1
                    tuples_to_eval.append((q, a))
                    if a.evaluation is None:
                        missing_evaluations += 1

        if missing_evaluations == 0 and not self.config.force:
            logger.info(
                f"All {all_tuples} answers are already evaluated.\n"
                "If you want to re-evaluate them, use the --force flag"
            )

        return tuples_to_eval

    def _build_message(self, query: Query, answer: AgentAnswer) -> LLMInputPrompt:
        """Builds the message to send to the LLM evaluator"""
        documents = self._filter_documents(query)
        context = {"query": query, "answer": answer, "documents": documents}
        user_message = self.user_prompt.render(**context)
        system_prompt = self.system_prompt.render(**context) if self.system_prompt else None
        return LLMInputPrompt(
            system_prompt=system_prompt,
            user_message=user_message,
        )

    def _build_message_pairwise(self, query: Query, game: PairwiseGame) -> LLMInputPrompt:
        """Builds the message to send to the LLM evaluator based on a pairwise game"""
        documents = self._filter_documents(query)
        context = {"query": query, "game": game, "documents": documents}
        user_message = self.user_prompt.render(**context)
        system_prompt = self.system_prompt.render(**context) if self.system_prompt else None
        return LLMInputPrompt(
            system_prompt=system_prompt,
            user_message=user_message,
        )

    @classmethod
    def from_config(cls, config: BaseAnswerEvaluatorConfig, llm_provider: BaseLLMProvider):
        return cls(config, llm_provider)

    @classmethod
    def get_config_class(cls) -> Type[BaseAnswerEvaluatorConfig]:
        return get_type_hints(cls)["config"]

    def __add_pairwise_games(self, experiment: Experiment):
        if not self.config.pairwise:
            return
        if not isinstance(self.config, PairwiseEvaluatorConfig):
            raise ValueError("Trying to add pairwise games to a non-pairwise evaluator")
        for query in experiment:
            query_agents = list(query.answers.keys())
            pairs = list(itertools.combinations(query_agents, 2))
            if self.config.bidirectional:
                pairs += [(b, a) for a, b in pairs]
            random.shuffle(pairs)

            # Filter out games that already exist
            existing_games = {(a.agent_a_answer.agent, a.agent_b_answer.agent) for a in query.pairwise_games}
            games = [g for g in pairs if g not in existing_games]

            games_to_add = min(self.config.n_games_per_query - len(existing_games), len(games))
            games = games[:games_to_add]
            for agent_a, agent_b in games:
                query.pairwise_games.append(
                    PairwiseGame(
                        qid=query.qid,
                        agent_a_answer=query.answers[agent_a],
                        agent_b_answer=query.answers[agent_b],
                    )
                )

    def _filter_documents(self, query: Query) -> list[Document]:
        documents = []
        for did, d in query.retrieved_docs.items():
            if self.config.document_relevance_threshold is not None:
                # Skip documents with relevance below the threshold
                if d.evaluation is None:
                    continue
                # check if evaluation.answer is an integer or a string that can be converted to an integer
                score = d.evaluation.answer
                if isinstance(d.evaluation.answer, dict):
                    score = d.evaluation.answer["score"]
                elif isinstance(d.evaluation.answer, BaseModel):
                    score = d.evaluation.answer.score  # type: ignore
                else:
                    score = d.evaluation.answer

                if isinstance(score, str) and score.isdigit():
                    score = int(score)
                if not isinstance(score, (int, float)):
                    continue
                if score < self.config.document_relevance_threshold:
                    continue
            if self.config.document_filter is not None and not self.config.document_filter(d):
                continue
            documents.append(d)
        if len(documents) == 0:
            if query.qid not in self._warned_queries:
                logger.warning(
                    f"No relevant documents were retrieved for the query {query.qid}. "
                    "No documents will be provided to the Answer Evaluator."
                )
                self._warned_queries.add(query.qid)
        return documents


class AnswerEvaluatorFactory:
    registry: dict[AnswerEvaluatorTypes, Type[BaseAnswerEvaluator]] = {}

    @classmethod
    def register(cls, name: AnswerEvaluatorTypes) -> Callable:
        def inner_wrapper(wrapped_class: Type[BaseAnswerEvaluator]):
            if name in cls.registry:
                logger.warning(f"Overwriting {name} in registry")
            cls.registry[name] = wrapped_class
            return wrapped_class

        return inner_wrapper

    @classmethod
    def create(
        cls,
        evaluator_name: AnswerEvaluatorTypes,
        llm_provider: BaseLLMProvider | str,
        config: BaseAnswerEvaluatorConfig | None = None,
        **kwargs,
    ) -> BaseAnswerEvaluator:
        if evaluator_name not in cls.registry:
            raise ValueError(
                f"Unknown answer evaluator {evaluator_name}\nValid options are {list(cls.registry.keys())}"
            )
        if isinstance(llm_provider, str):
            llm_provider_instance = get_llm_provider(llm_provider, **kwargs)
        else:
            llm_provider_instance = llm_provider
        if config is None:
            class_ = cls.registry[evaluator_name]
            type_config = class_.get_config_class()
            valid_keys = [field for field in type_config.model_fields]
            valid_args = {k: v for k, v in kwargs.items() if k in valid_keys}
            required_fields = [arg for arg, info in type_config.model_fields.items() if info.is_required()]
            for field in required_fields:
                if field not in valid_args:
                    raise ValueError(f"Required argument {field} for evaluator {evaluator_name} not provided")
            config = type_config(**valid_args)
        return cls.registry[evaluator_name].from_config(config, llm_provider_instance)


def get_answer_evaluator(
    evaluator_name: AnswerEvaluatorTypes | str | None = None,
    llm_provider: BaseLLMProvider | str = "openai",
    config: BaseAnswerEvaluatorConfig | None = None,
    **kwargs,
) -> BaseAnswerEvaluator:
    """Gets an answer evaluator
    Args:
        evaluator_name (AnswerEvaluatorTypes | str | None): The name of the answer evaluator to get.
        llm_provider (BaseLLMProvider | str): The LLM provider to use for the evaluation.
        config (BaseAnswerEvaluatorConfig | None): The configuration for the answer evaluator.
        **kwargs: Additional arguments to pass to the answer evaluator.
    Returns:
        BaseAnswerEvaluator: The answer evaluator.
    """
    if evaluator_name is None:
        # get the name from the config
        if config is None:
            raise ValueError("Either the evaluator_name or a config object must be provided")
        evaluator_name = config.evaluator_name
    if isinstance(evaluator_name, str):
        try:
            evaluator_name = AnswerEvaluatorTypes(evaluator_name)
        except ValueError:
            raise ValueError(f"Unknown answer evaluator {evaluator_name}")
    if evaluator_name is None:
        raise ValueError("The evaluator_name must be provided")
    return AnswerEvaluatorFactory.create(
        evaluator_name,
        llm_provider=llm_provider,
        config=config,
        **kwargs,
    )
