from concurrent.futures import ThreadPoolExecutor
import json
import re
import io
from PIL.Image import Image
import vertexai
from vertexai.generative_models import GenerativeModel, Image as GenAIImage, Part
import yaml
from pixaris.metrics.base import BaseMetric
from pixaris.metrics.utils import dict_mean


class LLMMetric(BaseMetric):
    def __init__(self, object_images: list[Image], style_images: list[Image]):
        super().__init__()
        self.object_images = object_images
        self.style_images = style_images

    def _PIL_image_to_vertex_image(self, image: Image) -> GenAIImage:
        """
        Converts a PIL image to a vertex image.
        Args:
            image (PIL.Image): The input PIL image.
        Returns:
            GenAIImage: The converted vertex image.
        """
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format="PNG")
        return GenAIImage.from_bytes(img_byte_arr.getvalue())

    def _llm_prompt(
        self,
        evaulation_image: Image,
        object_image: Image,
        style_image: Image,
    ):
        """
        Generates a prompt for rating images on various metrics and returns a JSON object with the ratings.
        Args:
            input_image_vertex (Part): The input image.
            output_image_vertex (Part): The output image.
            style_image_vertex (Part): The style image.
        Returns:
            str: A prompt that provokes a response  with a JSON with the following keys:
                - llm_reality: How real does the first picture look where 0 is not real at all and 1 is photorealistic?
                - llm_similarity: How similar is the main object in the first picture to the template in the second picture where 0 is not similar at all and 1 is identical?
                - llm_errors: How many errors can you find in the first picture?
                - llm_style: How well does the style of the first picture match the style of the third picture where 0 is not at all and 1 is identical?
        """
        json_prompt = "rate the following evluation image on the following metrics. return only a json file without newlines with the following keys:"
        reality_prompt = "llm_reality: How real does the evaluation image look where 0 is not real at all and 1 is photorealistic?"
        similarity_prompt = "llm_similarity: How similar is the main object in the evaluation image to the template in the following picture where 0 is not similar at all and 1 is identical?"
        error_prompt = (
            "llm_errors: How many errors can you find in the evaulation image?"
        )
        style_prompt = "llm_style: How well does the style of the evaluation image match the style of the following style image where 0 is not at all and 1 is identical?"

        return [
            json_prompt,
            Part.from_image(self._PIL_image_to_vertex_image(evaulation_image)),
            reality_prompt,
            similarity_prompt,
            Part.from_image(self._PIL_image_to_vertex_image(object_image)),
            error_prompt,
            style_prompt,
            Part.from_image(self._PIL_image_to_vertex_image(style_image)),
        ]

    def _postprocess_response(self, response_text: str) -> str:
        """
        If there is some sort of JSON-like structure in the response text, extract it and return it.
        Args:
        response_text (str): The response text from the model.
        Returns:
        str: The extracted JSON-like structure if found, otherwise the original response text.
        """
        pattern = r"\{.*\}"
        match = re.search(pattern, response_text)
        if match:
            extracted_string = match.group(0)
            return extracted_string
        else:
            raise ValueError("No JSON-like structure found in the llm response text.")

    def _response_to_dict(self, response_text: str) -> dict:
        """
        Converts the response text to a dictionary.
        Args:
            response_text (str): The response text from the model.
        Returns:
            dict: The response text as a dictionary.
        """
        parsed_text = self._postprocess_response(response_text)
        response_dict = json.loads(parsed_text)
        response_dict = {key: float(value) for key, value in response_dict.items()}

        if not all(
            metrics in response_dict
            for metrics in ["llm_reality", "llm_similarity", "llm_errors", "llm_style"]
        ):
            raise ValueError("Response dictionary does not contain all required keys.")
        return response_dict

    def _call_gemini(self, prompt) -> str:
        """
        Sends the prompt to Google API
        Args:
            prompt (str): The prompt for the LLM metrics. Generated by llm_prompt().
        Returns:
            string: The LLM response.
        """
        with open("pixaris/config.yaml", "r") as f:
            config = yaml.safe_load(f)

        vertexai.init(project=config["gcp_project_id"], location=config["gcp_location"])

        model = GenerativeModel("gemini-1.5-flash-preview-0514")

        responses = model.generate_content(prompt, stream=False)

        return responses.text

    def _successful_evaluation(self, prompt, max_tries: int = 3) -> dict:
        """
        Perform an evaluation by calling the `call_gemini` function with the given parameters.
        Assures that gemini returns correct json code by calling it up to max_tries times if it fails.
        Args:
            prompt: The prompt for the
            max_tries (optional): The maximum number of tries to perform the  Defaults to 3.
        Returns:
            A dictionary containing the response from the `call_gemini` function.
        Raises:
            ValueError: If an error occurs during the
        """
        for i in range(max_tries):
            try:
                return self._response_to_dict(self._call_gemini(prompt))
            except ValueError:
                pass

    def _combined_score(self, response: dict) -> float:
        """
        Calculates the combined score from the response dictionary.
        Args:
            response (dict): The response dictionary.
        Returns:
            float: The combined score. Score at the moment is just the mean of all the scores.
        """
        return sum(response.values()) / len(response.values())

    def _llm_scores_per_image(
        self,
        evaluation_image: Image,
        object_image: Image,
        style_image: Image,
        sample_size: int = 3,
    ) -> dict:
        """
        Calculates the LLM score for the generated image.
        Args:
            generated_image (Image): The generated image.
            template_image (Image): The template image.
            style_image (Image): The style image.
            sample_size (optional): The number of llm calls to evaluate. Defaults to 3
        Returns:
            dict: A dictionary containing the LLM scores.
        """
        scores = [
            self._successful_evaluation(
                self._llm_prompt(evaluation_image, object_image, style_image)
            )
            for _ in range(sample_size)
        ]

        # Calculate the average score for each metric
        average_scores_per_metric = dict_mean(scores)

        # Invert the error count to get a score
        average_scores_per_metric["llm_errors"] = 1 / (
            1 + average_scores_per_metric["llm_errors"]
        )

        # Calculate the llm_average score
        average_scores_per_metric["llm_average"] = self._combined_score(
            average_scores_per_metric
        )

        return average_scores_per_metric

    def calculate(self, x: list[Image]) -> dict:
        """
        Calculate the LLM metrics for a list of generated images.
        Args:
            x (list[Image]): A list of evaluation images.
        Returns:
            dict: A dictionary containing the mean LLM scores for each prompt.
        """
        if len(x) != len(self.object_images):
            raise ValueError(
                "There should be as many object images as generated images in the llm metric"
            )
        if len(x) != len(self.style_images):
            raise ValueError(
                "There should be as many style images as generated images in the llm metric"
            )

        with ThreadPoolExecutor(len(x)) as executor:
            llm_metrics = dict_mean(
                list(
                    executor.map(
                        self._llm_scores_per_image,
                        x,
                        self.object_images,
                        self.style_images,
                    )
                )
            )
            return llm_metrics
