from scalexi.utilities.logger import Logger
import os
from langchain_google_genai import ChatGoogleGenerativeAI

# Create a logger
logger = Logger().get_logger()

class Gemini:
    def __init__(self, google_key=None):
        """
        Initializes a Gemini object, setting up the Google API key and configuring timeout options.

        :method __init__: Constructs a Gemini object with specified Google API key and timeout settings.
        :type __init__: constructor

        :param google_key: The Google API key to be used for requests. If not provided, the method attempts to fetch it from an environment variable named 'GOOGLE_API_KEY'.
        :type google_key: str, optional

        :param enable_timeouts: Flag to determine if timeouts are enabled for API requests. Defaults to False.

        This method initializes the Gemini object by setting the Google API key, either from the provided parameter or from an environment variable. It also sets up timeout configurations based on the provided arguments. If the Google API key is not provided or found, it raises a ValueError.

        :raises ValueError: If the Google API key is not provided and cannot be found in the environment variables.

        :return: None.
        :rtype: None

        :example:

        ::

            >>> gemini = Gemini(google_key="your_api_key")
            # Initializes a Gemini object with the specified Google API key.
        """
        # Method implementation...
        self.google_api_key = google_key if google_key is not None else os.getenv("GOOGLE_API_KEY")
        if not self.google_api_key:
            raise ValueError("Invalid Google API key.")
        # Additional initialization can be added here if required

    def ask_gemini(self, prompt, temperature=0.7, 
                   max_output_tokens=1024, top_p=1, 
                   frequency_penalty=0, presence_penalty=0, 
                   max_retries=6, n=1, verbose=False, 
                   preprocess=True, request_type="DEFAULT", 
                   model_name="gemini-pro"):
        """
        Sends a request to the Gemini AI model and retrieves the generated response.

        :param request_params: Parameters for the request.
        :param temperature: Controls the randomness of the model's responses. Defaults to 0.7.
        :param max_output_tokens: The maximum length of the model's response. Defaults to 1024.
        :param top_p: Controls diversity of the model's responses via nucleus sampling. Defaults to 1.
        :param frequency_penalty: Decreases the likelihood of repeated content in the response. Defaults to 0.
        :param presence_penalty: Encourages the model to introduce new topics. Defaults to 0.
        :param max_retries: Maximum number of retries for the API call. Defaults to 6.
        :param n: Number of responses to generate. Defaults to 1.
        :param verbose: Enables verbose logging. Defaults to False.
        :param preprocess: Indicates whether the prompts should be preprocessed. Defaults to True.
        :param model_name: The name of the Gemini model to use. Defaults to 'gemini-pro'.

        :return: The content of the response generated by the Gemini AI model.
        :rtype: str
        """
        try:

            logger.debug(prompt)

            llm = ChatGoogleGenerativeAI(model=model_name, max_output_tokens = max_output_tokens,
                            top_p=1, 
                            temperature = temperature,
                            frequency_penalty=frequency_penalty, 
                            presence_penalty=presence_penalty, 
                            max_retries=max_retries,
                            n=n,                                    
                            verbose=verbose,
                            google_api_key=os.getenv("GOOGLE_API_KEY"))

            response = llm.invoke(prompt)
            return response.content
        except Exception as e:
            logger.error(f"[Gemini] An error occurred: {str(e)}")
            raise

# Usage example
#gemini = Gemini(google_key=os.getenv("GOOGLE_API_KEY"))
#response = gemini.ask_gemini(prompt="What is the latest advancement in AI technology?")
#print(response)
