import logging

from rag_fact_checker.data import Config, FactCheckerOutput
from rag_fact_checker.model.fact_checker import FactChecker
from rag_fact_checker.pipeline import PipelineDemonstration, PipelineLLM


class LLMMultiShotSplitFactChecker(FactChecker, PipelineLLM, PipelineDemonstration):
    """
    LLMSplitFactChecker is designed to compare answer triplets with reference triplets using a language model.
    The model compares a answer triplet with the reference triplets at each request and returns the merged comparison result.


    Attributes:
        config (Config): Config data class for initializing the class.
        logger (logging.Logger): Logger object for logging.

    Methods:

        forward(answer_triplets: List[List[str]], reference_triplets: List[List[List[str]]]): -> FactCheckerOutput
            Compares answer triplets with reference triplets and returns the comparison results.

        get_model_prompt(answer_triplets: List[List[str]], reference_triplets: List[List[str]], **kwargs) -> List[Dict[str, str]]:
            Generates a model prompt for comparing answer triplets with reference triplets.

        splitted_multishot_triplet_comparison_input_formatter(answer_triplets: List[List[str]], reference_triplets: List[List[str]]) -> Dict[str, str]:
            Formats the input for the triplet comparison.

        parse_splitted_triplet_comparison_output(string_output: str, answer_triplets: List[List[str]]
    ) -> bool:
            Parses the output from the model and returns the comparison result.
    """

    def __init__(self, config: Config, logger: logging.Logger):
        FactChecker.__init__(self, config, logger)
        PipelineLLM.__init__(self, config)
        PipelineDemonstration.__init__(self, config)

    def forward(
        self,
        answer_triplets: list[list[str]],
        reference_triplets: list[list[list[str]]],
    ) -> FactCheckerOutput:
        """
        Compares all answer triplet with reference triplets using a model and returns the comparison results.
        In one request, the model compares one answer triplet with all reference triplets.

        Args:
            answer_triplets (list): The triplets generated by a model or user.
            reference_triplets (list): The ground-truth or reference triplets.

        Returns:
            FactCheckerOutput which contains the binary fact-checking results.
        """
        comparison_result = {}

        for idx, answer_triplets in enumerate(answer_triplets):
            splitted_triplet_comparison_prompt = self.get_model_prompt(
                answer_triplets=answer_triplets,
                reference_triplets=self.flatten_triplets(reference_triplets),
            )
            response = self.model.chat.completions.create(
                model=self.config.model.llm.generator_model,
                messages=splitted_triplet_comparison_prompt,
                temperature=self.config.model.llm.temperature,
            )
            match_result = response.choices[0].message.content
            parsed_output = self.parse_splitted_triplet_comparison_output(
                string_output=match_result, answer_triplets=answer_triplets
            )
            comparison_result[idx] = parsed_output

        # log only the last prompt
        if self.config.experiment_setup.log_prompts:
            self.logger.debug(splitted_triplet_comparison_prompt)

        return FactCheckerOutput(fact_check_prediction_binary=comparison_result)

    def get_model_prompt(
        self,
        answer_triplets: list[list[str]],
        reference_triplets: list[list[str]],
        **kwargs,
    ) -> list[dict[str, str]]:
        """
        Generates a model prompt based on the provided answer and reference triplets.

        Args:
            answer_triplets (list): The triplets generated by a model or user.
            reference_triplets (list): The ground-truth or reference triplets.

        Returns:
            message_list: The formatted model prompt.
        """

        template_names = self.message_list_template["n_shot_triplet_match_test"]
        return self.create_messages(
            template_names,
            **self.splitted_multishot_triplet_comparison_input_formatter(
                answer_triplets,
                reference_triplets,
            ),
        )

    def splitted_multishot_triplet_comparison_input_formatter(
        self,
        answer_triplets: list[list[str]],
        reference_triplets: list[list[str]],
    ) -> dict[str, str]:
        """
        Formats the input for comparing answer triplets with reference triplets.

        Args:
            answer_triplets (list): The triplets generated by a model or user.
            reference_triplets (list): The ground-truth or reference triplets.

        Returns:
            MultishotFactCheckerModelInput
        """
        examples = self.get_demo_data(
            demo_type="fact_checker",
        )
        return {
            "answer_triplets": f"0: {answer_triplets}",
            "reference_triplets": "\n-".join(
                [str(source_triplet) for source_triplet in reference_triplets]
            ),
            "examples": examples,
        }

    def parse_splitted_triplet_comparison_output(
        self, string_output: str, answer_triplets: list[list[str]]
    ) -> bool:
        """
        Parses the output string from a triplet comparison and extracts the final answer.

        Args:
            string_output (str): The output string containing the triplet comparison results.
            answer_triplets (list): A list of triplets representing the answer.

        Returns:
            list or bool: The parsed final answer as a list if successful, otherwise False.
        """
        try:
            answer_part = string_output.split("[FINAL ANSWER]")[-1]
            splitted_string_outputs = answer_part.split(":")[-1].strip()

            return eval(splitted_string_outputs)
        except Exception as e:
            self.logger.warning(
                "Failed to parse the splitted fact checker output : %s", str(e)
            )
            self.logger.debug("Error occured in : %s", string_output)
            self.logger.debug("Answer triplets: %s", answer_triplets)
            return False
