from openai import OpenAI
import json
from scalexi.utilities.logger import Logger
import os
import httpx
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate



#create a logger file
logger = Logger().get_logger()

class GPT:
    def __init__(self, openai_key=None, enable_timeouts= False, timeouts_options= None):
        self.openai_api_key = openai_key if openai_key is not None else os.getenv("OPENAI_API_KEY")
        if not self.openai_api_key or not self.openai_api_key.startswith("sk-"):
            raise ValueError("Invalid OpenAI API key.")
        self.client = OpenAI(api_key=self.openai_api_key, max_retries=3)
        if enable_timeouts:
            if timeouts_options is None:
                timeouts_options = {"total": 120, "read": 60.0, "write": 60.0, "connect": 10.0}
                self.client = self.client.with_options(timeout=httpx.Timeout(120.0, read=60.0, write=60.0, connect=10.0))
            else:
                self.client = self.client.with_options(timeout=httpx.Timeout(timeouts_options["total"], timeouts_options["read"], timeouts_options["write"], timeouts_options["connect"]))
        
    def ask_gpt(self, user_prompt, system_prompt, temperature=0.7, 
                model_name="gpt-3.5-turbo", max_tokens=1024, top_p=1, 
                frequency_penalty=0, presence_penalty=0, stop=["END"], preprocess=True):
        """
        Sends a user prompt to a GPT-based language model and retrieves the generated response.

        :method ask_gpt: Interacts with a GPT-based language model using OpenAI's API to generate responses for given user and system prompts.
        :type ask_gpt: method

        :param user_prompt: The input prompt provided by the user.
        :type user_prompt: str

        :param system_prompt: The initial prompt or context provided by the system, setting the stage for the user prompt.
        :type system_prompt: str

        :param temperature: Controls the randomness of the model's responses. Lower values make responses more deterministic. Defaults to 0.7.
        :type temperature: float, optional

        :param model_name: The name of the GPT model to use, e.g., 'gpt-3.5-turbo'. Defaults to 'gpt-3.5-turbo'.
        :type model_name: str, optional

        :param max_tokens: The maximum length of the model's response. Defaults to 1024.
        :type max_tokens: int, optional

        :param top_p: Controls diversity of the model's responses via nucleus sampling, with 1 being the most diverse. Defaults to 1.
        :type top_p: float, optional

        :param frequency_penalty: Decreases the likelihood of repeated content in the response. Defaults to 0.
        :type frequency_penalty: float, optional

        :param presence_penalty: Encourages the model to introduce new topics. Defaults to 0.
        :type presence_penalty: float, optional

        :param stop: A list of stop sequences that signal the end of a response. Defaults to ["END"].
        :type stop: list[str], optional

        :param preprocess: Indicates whether the prompts should be preprocessed before being sent to the model. Defaults to True.
        :type preprocess: bool, optional

        This method sends prompts to the OpenAI API, specifying parameters like temperature, token limits, and penalties. It handles the API response and returns the content of the generated message.

        :return: The content of the response generated by the language model.
        :rtype: str

        :raises Exception: If an error occurs during the API interaction.

        :example:

        ::

            >>> ask_gpt("How do I cook pasta?", "Here's some cooking advice:")
            # This will return a response from the GPT model related to cooking pasta, based on the given user and system prompts.
        """
     
        try:
            response = self.client.chat.completions.create(
                model=model_name,
                messages=[
                        {
                        "role": "system",
                        "content": system_prompt
                        },
                        {
                        "role": "user",
                        "content": user_prompt 
                        }
                    ],
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                frequency_penalty=frequency_penalty,
                presence_penalty=presence_penalty,
                stop=["END"]
            )
            
            return response.choices[0].message.content, self.extract_token_usage(response)
        except Exception as e:
            print('[ask_gpt] Error in review_paper_section:', e)
            raise 
            #return None
            
    def ask_gpt_langchain(self, user_prompt, system_prompt, temperature=0.7, 
                    model_name="gpt-3.5-turbo", max_tokens=1024, top_p=1, 
                    frequency_penalty=0, presence_penalty=0, stop=["END"], preprocess=True):
        """
        Sends a user prompt to a GPT-based language model and retrieves the generated response.

        :method ask_gpt_langchain: Communicates with a GPT-based language model to generate responses based on a given user prompt and system prompt.
        :type ask_gpt_langchain: method

        :param user_prompt: The input prompt provided by the user.
        :type user_prompt: str

        :param system_prompt: The initial prompt or context provided by the system, setting the stage for the user prompt.
        :type system_prompt: str

        :param temperature: Controls the randomness of the model's responses. Lower values make responses more deterministic. Defaults to 0.7.
        :type temperature: float, optional

        :param model_name: The name of the GPT model to use, e.g., 'gpt-3.5-turbo'. Defaults to 'gpt-3.5-turbo'.
        :type model_name: str, optional

        :param max_tokens: The maximum length of the model's response. Defaults to 1024.
        :type max_tokens: int, optional

        :param top_p: Controls diversity of the model's responses via nucleus sampling, with 1 being the most diverse. Defaults to 1.
        :type top_p: float, optional

        :param frequency_penalty: Decreases the likelihood of repeated content in the response. Defaults to 0.
        :type frequency_penalty: float, optional

        :param presence_penalty: Encourages the model to introduce new topics. Defaults to 0.
        :type presence_penalty: float, optional

        :param stop: A list of stop sequences that signal the end of a response. Defaults to ["END"].
        :type stop: list[str], optional

        :param preprocess: Indicates whether the prompts should be preprocessed before being sent to the model. Defaults to True.
        :type preprocess: bool, optional

        This method constructs a language model interaction chain using user and system prompts, then invokes the chain to get a response from the specified GPT model. The response content is then returned.

        :return: The content of the response generated by the language model.
        :rtype: str

        :example:

        ::

            >>> ask_gpt_langchain("How do I cook pasta?", "Here's some cooking advice:")
            # This will return a response from the GPT model related to cooking pasta, based on the given user and system prompts.
        """

        #messages = [
        #    SystemMessage(content=system_prompt),
        #    HumanMessage(content=user_prompt),
        ##response = llm.invoke(messages)
        
        
        try:
            prompt = ChatPromptTemplate.from_messages([
                ("system", system_prompt),
                ("user", "{input}")
            ])
            
            llm = ChatOpenAI(openai_api_key=self.openai_api_key, temperature=temperature, 
                            model_name=model_name, max_tokens=max_tokens)
            
            chain = prompt | llm
            response = chain.invoke({"input": user_prompt})      
            return response.content
        except Exception as e:
            # Here you can log the exception, re-raise it, or handle it as needed.
            # This is a simple example of logging and re-raising the exception.
            print(f"[ask_gpt_langchain] An error occurred: {str(e)}")
            raise
        
        
    def extract_token_usage(self, response):
        """
        Extracts token usage information from a given response object.

        :method extract_token_usage: Retrieves the number of tokens used for the prompt, completion, and the total tokens from the response object, if available.
        :type extract_token_usage: method

        :param response: The response object from which to extract token usage data.
        :type response: Response

        This method checks if the provided response object contains token usage information. If so, it extracts the number of tokens used for the prompt, the completion, and the total token count. If the response object does not have token usage information, it returns an error message.

        :return: A dictionary containing token usage information, or an error message if such information is not available.
        :rtype: dict

        :example:

        ::

            >>> response = some_api_call()
            >>> extract_token_usage(response)
            # This will return a dictionary with 'prompt_tokens', 'completion_tokens', and 'total_tokens', or an error message.
        """

        if hasattr(response, 'usage'):
            token_usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens
            }
        else:
            token_usage = {"error": "No token usage information available"}

        return token_usage

