import openai
import time
import sys

from syphus.prompts import prompts
import syphus.data_generator.gpt_params_settings as gpt_params_settings
import syphus.data_generator.openai_settings as openai_settings

from typing import Optional, List, Any


class GPTManager(object):
    """
    A class that manages interactions with the OpenAI GPT engine.

    Attributes:
        gpt_info_path (str, optional): Path to a YAML file containing OpenAI API and GPT parameters settings.
        openai_api (openai_settings.OpenAISettings, optional): An instance of OpenAISettings containing OpenAI API settings.
        gpt_params (gpt_params_settings.GPTParamsSettings, optional): An instance of GPTParamsSettings containing GPT parameters settings.

    Raises:
        ValueError: If neither gpt_info_path nor openai_api is provided during initialization.

    """

    def __init__(
        self,
        *,
        gpt_info_path: Optional[str] = None,
        openai_api: Optional[openai_settings.OpenAISettings] = None,
        gpt_params: Optional[gpt_params_settings.GPTParamsSettings] = None,
    ):
        """
        Initialize the GPTManager instance.

        Args:
            gpt_info_path (str, optional): Path to a YAML file containing OpenAI API and GPT parameters settings.
            openai_api (openai_settings.OpenAISettings, optional): An instance of OpenAISettings containing OpenAI API settings.
            gpt_params (gpt_params_settings.GPTParamsSettings, optional): An instance of GPTParamsSettings containing GPT-3 parameters settings.

        """
        if gpt_info_path:
            if openai_api or gpt_params:
                print(
                    "Warning: gpt_info_path overrides openai_api and gpt_params",
                    flush=True,
                    file=sys.stderr,
                )
            self.openai_api = openai_settings.read_yaml(gpt_info_path)
            self.gpt_params = gpt_params_settings.read_yaml(gpt_info_path)
        elif openai_api:
            self.openai_api = openai_api
            if gpt_params:
                self.gpt_params = gpt_params
            else:
                self.gpt_params = gpt_params_settings.GPTParamsSettings()
        else:
            raise ValueError("Must provide either gpt_info_path or openai_api")
        openai.api_key = self.openai_api.key

    def set_gpt_params(self, gpt_params: gpt_params_settings.GPTParamsSettings):
        """
        Set GPT parameters for generating responses.

        Args:
            gpt_params (gpt_params_settings.GPTParamsSettings): An instance of GPTParamsSettings containing GPT parameters settings.

        """
        self.gpt_params = gpt_params

    def query_gpt(self, prompt: List[Any]):
        """
        Generate a response from the GPT-3 engine based on the provided prompt.

        Args:
            prompt (List[Any]): An instance of Prompts containing the conversation prompt and messages.

        Returns:
            dict: A dictionary containing the response generated by the GPT-3 engine.

        """
        success = False
        while not success:
            try:
                response = openai.ChatCompletion.create(
                    model=self.openai_api.engine,
                    messages=prompt,
                    temperature=self.gpt_params.temperature,
                    max_tokens=self.gpt_params.max_tokens,
                    top_p=self.gpt_params.top_p,
                    frequency_penalty=self.gpt_params.frequency_penalty,
                    presence_penalty=self.gpt_params.presence_penalty,
                    stop=self.gpt_params.stop,
                )
                success = True
            except Exception as e:
                print(f"Error: {e}")
                if "rate limit" in str(e) or "Rate limit" in str(e):
                    print("Sleeping for 3 seconds")
                    time.sleep(3)
                else:
                    raise e
        return response
