import base64
import litellm
import instructor
from typing import Type
from pydantic import BaseModel
from openai import ContentFilterFinishReasonError
from litellm.exceptions import ContentPolicyViolationError
from instructor.exceptions import InstructorRetryException

from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
    LLMInterface,
)
from cognee.infrastructure.llm.exceptions import (
    ContentPolicyFilterError,
    MissingSystemPromptPathError,
)
from cognee.infrastructure.files.utils.open_data_file import open_data_file
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
    rate_limit_async,
    rate_limit_sync,
    sleep_and_retry_async,
    sleep_and_retry_sync,
)
from cognee.modules.observability.get_observe import get_observe
from cognee.shared.logging_utils import get_logger

observe = get_observe()

logger = get_logger()


class OpenAIAdapter(LLMInterface):
    """
    Adapter for OpenAI's GPT-3, GPT-4 API.

    Public methods:

    - acreate_structured_output
    - create_structured_output
    - create_transcript
    - transcribe_image
    - show_prompt

    Instance variables:

    - name
    - model
    - api_key
    - api_version
    - MAX_RETRIES
    """

    name = "OpenAI"
    model: str
    api_key: str
    api_version: str

    MAX_RETRIES = 5

    """Adapter for OpenAI's GPT-3, GPT=4 API"""

    def __init__(
        self,
        api_key: str,
        endpoint: str,
        api_version: str,
        model: str,
        transcription_model: str,
        max_completion_tokens: int,
        streaming: bool = False,
        fallback_model: str = None,
        fallback_api_key: str = None,
        fallback_endpoint: str = None,
    ):
        self.aclient = instructor.from_litellm(litellm.acompletion)
        self.client = instructor.from_litellm(litellm.completion)
        self.transcription_model = transcription_model
        self.model = model
        self.api_key = api_key
        self.endpoint = endpoint
        self.api_version = api_version
        self.max_completion_tokens = max_completion_tokens
        self.streaming = streaming

        self.fallback_model = fallback_model
        self.fallback_api_key = fallback_api_key
        self.fallback_endpoint = fallback_endpoint

    @observe(as_type="generation")
    @sleep_and_retry_async()
    @rate_limit_async
    async def acreate_structured_output(
        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
    ) -> BaseModel:
        """
        Generate a response from a user query.

        This method asynchronously creates structured output by sending a request to the OpenAI
        API using the provided parameters to generate a completion based on the user input and
        system prompt.

        Parameters:
        -----------

            - text_input (str): The input text provided by the user for generating a response.
            - system_prompt (str): The system's prompt to guide the model's response.
            - response_model (Type[BaseModel]): The expected model type for the response.

        Returns:
        --------

            - BaseModel: A structured output generated by the model, returned as an instance of
              BaseModel.
        """

        try:
            return await self.aclient.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": f"""{text_input}""",
                    },
                    {
                        "role": "system",
                        "content": system_prompt,
                    },
                ],
                api_key=self.api_key,
                api_base=self.endpoint,
                api_version=self.api_version,
                response_model=response_model,
                max_retries=self.MAX_RETRIES,
                extra_body={"reasoning_effort": "minimal"},
            )
        except (
            ContentFilterFinishReasonError,
            ContentPolicyViolationError,
            InstructorRetryException,
        ) as error:
            if (
                isinstance(error, InstructorRetryException)
                and "content management policy" not in str(error).lower()
            ):
                logger.debug(
                    "LLM Model does not support reasoning_effort parameter, trying call without the parameter."
                )
                return await self.aclient.chat.completions.create(
                    model=self.model,
                    messages=[
                        {
                            "role": "user",
                            "content": f"""{text_input}""",
                        },
                        {
                            "role": "system",
                            "content": system_prompt,
                        },
                    ],
                    api_key=self.api_key,
                    api_base=self.endpoint,
                    api_version=self.api_version,
                    response_model=response_model,
                    max_retries=self.MAX_RETRIES,
                )

            if not (self.fallback_model and self.fallback_api_key):
                raise ContentPolicyFilterError(
                    f"The provided input contains content that is not aligned with our content policy: {text_input}"
                )

            try:
                return await self.aclient.chat.completions.create(
                    model=self.fallback_model,
                    messages=[
                        {
                            "role": "user",
                            "content": f"""{text_input}""",
                        },
                        {
                            "role": "system",
                            "content": system_prompt,
                        },
                    ],
                    api_key=self.fallback_api_key,
                    # api_base=self.fallback_endpoint,
                    response_model=response_model,
                    max_retries=self.MAX_RETRIES,
                )
            except (
                ContentFilterFinishReasonError,
                ContentPolicyViolationError,
                InstructorRetryException,
            ) as error:
                if (
                    isinstance(error, InstructorRetryException)
                    and "content management policy" not in str(error).lower()
                ):
                    raise error
                else:
                    raise ContentPolicyFilterError(
                        f"The provided input contains content that is not aligned with our content policy: {text_input}"
                    )

    @observe
    @sleep_and_retry_sync()
    @rate_limit_sync
    def create_structured_output(
        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
    ) -> BaseModel:
        """
        Generate a response from a user query.

        This method creates structured output by sending a synchronous request to the OpenAI API
        using the provided parameters to generate a completion based on the user input and
        system prompt.

        Parameters:
        -----------

            - text_input (str): The input text provided by the user for generating a response.
            - system_prompt (str): The system's prompt to guide the model's response.
            - response_model (Type[BaseModel]): The expected model type for the response.

        Returns:
        --------

            - BaseModel: A structured output generated by the model, returned as an instance of
              BaseModel.
        """

        return self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": f"""{text_input}""",
                },
                {
                    "role": "system",
                    "content": system_prompt,
                },
            ],
            api_key=self.api_key,
            api_base=self.endpoint,
            api_version=self.api_version,
            response_model=response_model,
            max_retries=self.MAX_RETRIES,
        )

    @rate_limit_async
    async def create_transcript(self, input):
        """
        Generate an audio transcript from a user query.

        This method creates a transcript from the specified audio file, raising a
        FileNotFoundError if the file does not exist. The audio file is processed and the
        transcription is retrieved from the API.

        Parameters:
        -----------

            - input: The path to the audio file that needs to be transcribed.

        Returns:
        --------

            The generated transcription of the audio file.
        """

        async with open_data_file(input, mode="rb") as audio_file:
            transcription = litellm.transcription(
                model=self.transcription_model,
                file=audio_file,
                api_key=self.api_key,
                api_base=self.endpoint,
                api_version=self.api_version,
                max_retries=self.MAX_RETRIES,
            )

        return transcription

    @rate_limit_async
    async def transcribe_image(self, input) -> BaseModel:
        """
        Generate a transcription of an image from a user query.

        This method encodes the image and sends a request to the OpenAI API to obtain a
        description of the contents of the image.

        Parameters:
        -----------

            - input: The path to the image file that needs to be transcribed.

        Returns:
        --------

            - BaseModel: A structured output generated by the model, returned as an instance of
              BaseModel.
        """
        async with open_data_file(input, mode="rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        return litellm.completion(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "What's in this image?",
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{encoded_image}",
                            },
                        },
                    ],
                }
            ],
            api_key=self.api_key,
            api_base=self.endpoint,
            api_version=self.api_version,
            max_completion_tokens=300,
            max_retries=self.MAX_RETRIES,
        )

    def show_prompt(self, text_input: str, system_prompt: str) -> str:
        """
        Format and display the prompt for a user query.

        This method formats the prompt using the provided user input and system prompt,
        returning a string representation. Raises MissingSystemPromptPathError if the system prompt is not
        provided.

        Parameters:
        -----------

            - text_input (str): The input text provided by the user.
            - system_prompt (str): The system's prompt to guide the model's response.

        Returns:
        --------

            - str: A formatted string representing the user input and system prompt.
        """
        if not text_input:
            text_input = "No user input provided."
        if not system_prompt:
            raise MissingSystemPromptPathError()
        system_prompt = LLMGateway.read_query_prompt(system_prompt)

        formatted_prompt = (
            f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
            if system_prompt
            else None
        )
        return formatted_prompt
