class BaseGoogle(LLM):
    """Base class to implement a new Google LLM

     LLM base class is extended to be used with Google Palm API.
    """

    genai: Any
    temperature: Optional[float] = 0
    top_p: Optional[float] = None
    top_k: Optional[float] = None
    max_output_tokens: Optional[int] = None

    def _configure(self, api_key: str):
        """
        Configure Google Palm API Key
        Args:
            api_key (str): A string of API keys generated from Google Cloud

        Returns:

        """

        if not api_key:
            raise APIKeyNotFoundError("Google Palm API key is required")
        try:
            # pylint: disable=import-outside-toplevel
            import google.generativeai as genai

            genai.configure(api_key=api_key)
        except ImportError as ex:
            raise ImportError(
                "Could not import google-generativeai python package. "
                "Please install it with `pip install google-generativeai`."
            ) from ex
        self.genai = genai

    def _valid_params(self):
        return ["temperature", "top_p", "top_k", "max_output_tokens"]

    def _set_params(self, **kwargs):
        """
        Set Parameters
        Args:
            **kwargs: ["temperature", "top_p", "top_k", "max_output_tokens"]

        Returns:

        """

        valid_params = self._valid_params()
        for key, value in kwargs.items():
            if key in valid_params:
                setattr(self, key, value)

    def _validate(self):
        """Validates the parameters for Google"""

        if self.temperature is not None and not 0 <= self.temperature <= 1:
            raise ValueError("temperature must be in the range [0.0, 1.0]")

        if self.top_p is not None and not 0 <= self.top_p <= 1:
            raise ValueError("top_p must be in the range [0.0, 1.0]")

        if self.top_k is not None and not 0 <= self.top_k <= 1:
            raise ValueError("top_k must be in the range [0.0, 1.0]")

        if self.max_output_tokens is not None and self.max_output_tokens <= 0:
            raise ValueError("max_output_tokens must be greater than zero")

    @abstractmethod
    def _generate_text(self, prompt: str) -> str:
        """
        Generates text for prompt, specific to implementation.

        Args:
            prompt (str): Prompt

        Returns:
            str: LLM response
        """
        raise MethodNotImplementedError("method has not been implemented")

    def call(self, instruction: Prompt, value: str, suffix: str = "") -> str:
        """
        Call the Google LLM.

        Args:
            instruction (str): Instruction to pass
            value (str): Value to pass
            suffix (str): Suffix to pass

        Returns:
            str: Response
        """
        self.last_prompt = str(instruction) + str(value)
        prompt = str(instruction) + str(value) + suffix
        return self._generate_text(prompt)

class BaseOpenAI(LLM, ABC):
    """Base class to implement a new OpenAI LLM
    LLM base class, this class is extended to be used with OpenAI API.

    """

    api_token: str
    temperature: float = 0
    max_tokens: int = 512
    top_p: float = 1
    frequency_penalty: float = 0
    presence_penalty: float = 0.6
    stop: Optional[str] = None

    def _set_params(self, **kwargs):

        """
        Set Parameters
        Args:
            **kwargs: ["model", "engine", "deployment_id", "temperature","max_tokens",
            "top_p", "frequency_penalty", "presence_penalty", "stop", ]

        Returns: None

        """

        valid_params = [
            "model",
            "engine",
            "deployment_id",
            "temperature",
            "max_tokens",
            "top_p",
            "frequency_penalty",
            "presence_penalty",
            "stop",
        ]
        for key, value in kwargs.items():
            if key in valid_params:
                setattr(self, key, value)

    @property
    def _default_params(self) -> Dict[str, Any]:

        """
        Get the default parameters for calling OpenAI API

        Returns (Dict): A dict of OpenAi API parameters

        """

        return {
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "top_p": self.top_p,
            "frequency_penalty": self.frequency_penalty,
            "presence_penalty": self.presence_penalty,
        }

    def completion(self, prompt: str) -> str:
        """
        Query the completion API

        Args:
            prompt (str): Prompt

        Returns:
            str: LLM response
        """
        params = {**self._default_params, "prompt": prompt}

        if self.stop is not None:
            params["stop"] = [self.stop]

        response = openai.Completion.create(**params)

        return response["choices"][0]["text"]

    def chat_completion(self, value: str) -> str:
        """
        Query the chat completion API

        Args:
            value (str): Prompt

        Returns:
            str: LLM response
        """
        params = {
            **self._default_params,
            "messages": [
                {
                    "role": "system",
                    "content": value,
                }
            ],
        }

        if self.stop is not None:
            params["stop"] = [self.stop]

        response = openai.ChatCompletion.create(**params)

        return response["choices"][0]["message"]["content"]

class HuggingFaceLLM(LLM):
    """Base class to implement a new Hugging Face LLM.

    LLM base class is extended to be used with HuggingFace LLM Modes APIs

    """

    last_prompt: Optional[str] = None
    api_token: str
    _api_url: str = "https://api-inference.huggingface.co/models/"
    _max_retries: int = 3

    @property
    def type(self) -> str:
        return "huggingface-llm"

    def query(self, payload):
        """
        Query the HF API
        Args:
            payload: A JSON form payload

        Returns: Generated Response

        """

        headers = {"Authorization": f"Bearer {self.api_token}"}

        response = requests.post(
            self._api_url, headers=headers, json=payload, timeout=60
        )

        return response.json()[0]["generated_text"]

    def call(self, instruction: Prompt, value: str, suffix: str = "") -> str:

        """
        A call method of HuggingFaceLLM class.
        Args:
            instruction (object): A prompt object
            value (str):
            suffix (str):

        Returns (str): A string response

        """

        prompt = str(instruction)
        payload = prompt + value + suffix

        # sometimes the API doesn't return a valid response, so we retry passing the
        # output generated from the previous call as the input
        for _i in range(self._max_retries):
            response = self.query({"inputs": payload})
            payload = response
            if response.count("<endCode>") >= 2:
                break

        # replace instruction + value from the inputs to avoid showing it in the output
        output = response.replace(prompt + value + suffix, "")
        return output

class LLM:
    """Base class to implement a new LLM."""

    last_prompt: Optional[str] = None

    @property
    def type(self) -> str:
        """
        Return type of LLM.

        Raises:
            APIKeyNotFoundError: Type has not been implemented

        Returns:
            str: Type of LLM a string
        """
        raise APIKeyNotFoundError("Type has not been implemented")

    def _polish_code(self, code: str) -> str:
        """
        Polish the code by removing the leading "python" or "py",  \
        removing the imports and removing trailing spaces and new lines.

        Args:
            code (str): Code

        Returns:
            str: Polished code
        """
        if re.match(r"^(python|py)", code):
            code = re.sub(r"^(python|py)", "", code)
        if re.match(r"^`.*`$", code):
            code = re.sub(r"^`(.*)`$", r"\1", code)
        code = code.strip()
        return code

    def _is_python_code(self, string):
        """
        Return True if it is valid python code.
        Args:
            string (str):

        Returns (bool): True if Python Code otherwise False

        """
        try:
            ast.parse(string)
            return True
        except SyntaxError:
            return False

    def _extract_code(self, response: str, separator: str = "```") -> str:
        """
        Extract the code from the response.

        Args:
            response (str): Response
            separator (str, optional): Separator. Defaults to "```".

        Raises:
            NoCodeFoundError: No code found in the response

        Returns:
            str: Extracted code from the response
        """
        code = response
        match = re.search(
            rf"{START_CODE_TAG}(.*)({END_CODE_TAG}|{END_CODE_TAG.replace('<', '</')})",
            code,
            re.DOTALL,
        )
        if match:
            code = match.group(1).strip()
        if len(code.split(separator)) > 1:
            code = code.split(separator)[1]
        code = self._polish_code(code)
        if not self._is_python_code(code):
            raise NoCodeFoundError("No code found in the response")

        return code

    @abstractmethod
    def call(self, instruction: Prompt, value: str, suffix: str = "") -> str:
        """
        Execute the LLM with given prompt.

        Args:
            instruction (Prompt): Prompt
            value (str): Value
            suffix (str, optional): Suffix. Defaults to "".

        Raises:
            MethodNotImplementedError: Call method has not been implemented
        """
        raise MethodNotImplementedError("Call method has not been implemented")

    def generate_code(self, instruction: Prompt, prompt: str) -> str:
        """
        Generate the code based on the instruction and the given prompt.

        Returns:
            str: Code
        """
        return self._extract_code(self.call(instruction, prompt, suffix="\n\nCode:\n"))

class OpenAI(BaseOpenAI):
    """OpenAI LLM using BaseOpenAI Class.

    An API call to OpenAi API is sent and response is recorded and returned.
    The default chat model is **gpt-3.5-turbo** while **text-davinci-003** is only supported
    completion model.
    The list of supported Chat models includes ["gpt-4", "gpt-4-0314", "gpt-4-32k",
     "gpt-4-32k-0314","gpt-3.5-turbo", "gpt-3.5-turbo-0301"].

    """

    _supported_chat_models = [
        "gpt-4",
        "gpt-4-0314",
        "gpt-4-32k",
        "gpt-4-32k-0314",
        "gpt-3.5-turbo",
        "gpt-3.5-turbo-0301",
    ]
    _supported_completion_models = ["text-davinci-003"]

    model: str = "gpt-3.5-turbo"

    def __init__(
        self,
        api_token: Optional[str] = None,
        **kwargs,
    ):
        """
        __init__ method of OpenAI Class
        Args:
            api_token (str): API Token fro OpenAI platform.
            **kwargs: Extended Parameters inferred from BaseOpenAI class

        Returns: Response generated from OpenAI API
        """

        self.api_token = api_token or os.getenv("OPENAI_API_KEY") or None
        if self.api_token is None:
            raise APIKeyNotFoundError("OpenAI API key is required")
        openai.api_key = self.api_token

        self._set_params(**kwargs)

    @property
    def _default_params(self) -> Dict[str, Any]:
        """Get the default parameters for calling OpenAI API"""
        return {
            **super()._default_params,
            "model": self.model,
        }

    def call(self, instruction: Prompt, value: str, suffix: str = "") -> str:
        """
        Call the OpenAI LLM.

        Args:
            instruction (Prompt): Instruction to pass
            value (str): Value to pass
            suffix (str): Suffix to pass

        Raises:
            UnsupportedOpenAIModelError: Unsupported model

        Returns:
            str: Response
        """
        self.last_prompt = str(instruction) + str(value)

        if self.model in self._supported_completion_models:
            response = self.completion(str(instruction) + str(value) + suffix)
        elif self.model in self._supported_chat_models:
            response = self.chat_completion(str(instruction) + str(value) + suffix)
        else:
            raise UnsupportedOpenAIModelError("Unsupported model")

        return response

    @property
    def type(self) -> str:
        return "openai"

class OpenAssistant(HuggingFaceLLM):
    """Open Assistant LLM
     A base HuggingFaceLLM class is extended to use OpenAssistant Model via its API.
     Currently `oasst-sft-1-pythia-12b` is supported via this module.
    """

    api_token: str
    _api_url: str = (
        "https://api-inference.huggingface.co/models/"
        "OpenAssistant/oasst-sft-1-pythia-12b"
    )
    _max_retries: int = 10

    def __init__(self, api_token: Optional[str] = None):

        """
        __init__ method of OpenAssistant Calss

        Raises:
            APIKeyNotFoundError: HuggingFace Hub API key is required

        Args:
            api_token (str): API token from Huggingface platform
        """
        self.api_token = api_token or os.getenv("HUGGINGFACE_API_KEY") or None
        if self.api_token is None:
            raise APIKeyNotFoundError("HuggingFace Hub API key is required")

    @property
    def type(self) -> str:
        return "open-assistant"

class Starcoder(HuggingFaceLLM):

    """Starcoder LLM API

    A base HuggingFaceLLM class is extended to use Starcoder model.

    """

    api_token: str
    _api_url: str = "https://api-inference.huggingface.co/models/bigcode/starcoder"
    _max_retries: int = 5

    def __init__(self, api_token: Optional[str] = None):
        """
        __init__ method of Starcoder Class
        Args:
            api_token (str): API token from Huggingface platform
        """

        self.api_token = api_token or os.getenv("HUGGINGFACE_API_KEY") or None
        if self.api_token is None:
            raise APIKeyNotFoundError("HuggingFace Hub API key is required")

    @property
    def type(self) -> str:
        return "starcoder"

class AzureOpenAI(BaseOpenAI):
    """OpenAI LLM via Microsoft Azure
    This class in using BaseOpenAI class to include Azure OpenAI Features.
    """

    api_base: str
    api_type: str = "azure"
    api_version: str
    engine: str

    def __init__(
        self,
        api_token: Optional[str] = None,
        api_base: Optional[str] = None,
        api_version: Optional[str] = None,
        deployment_name: str = None,
        **kwargs, ):

        """
        __init__ method of AzureOpenAI Class

        Args:
            api_token (str): Azure OpenAI API token.
            api_base (str): Base url of the Azure endpoint.
                It should look like the following: <https://YOUR_RESOURCE_NAME.openai.azure.com/>
            api_version (str): Version of the Azure OpenAI API. Be aware the API version may change.
            deployment_name: Custom name of the deployed model
            **kwargs: Inference Parameters
        """

        self.api_token = api_token or os.getenv("AZURE_OPENAI_KEY") or None
        self.api_base = api_base or os.getenv("AZURE_OPENAI_ENDPOINT") or None
        self.api_version = api_version or "2023-03-15-preview"
        if self.api_token is None:
            raise APIKeyNotFoundError("Azure OpenAI key is required")
        if self.api_base is None:
            raise APIKeyNotFoundError("Azure OpenAI base endpoint is required")

        openai.api_key = self.api_token
        openai.api_base = self.api_base
        openai.api_version = self.api_version
        openai.api_type = self.api_type

        if deployment_name is None:
            raise UnsupportedOpenAIModelError("Model deployment name is required.")
        try:
            model_name = openai.Deployment.retrieve(deployment_name).model
            model_capabilities = openai.Model.retrieve(model_name).capabilities
            if (
                not model_capabilities.completion
                and not model_capabilities.chat_completion
            ):
                raise UnsupportedOpenAIModelError(
                    "Model deployment name does not correspond to a "
                    "chat nor a completion model."
                )
            self.is_chat_model = model_capabilities.chat_completion
            self.engine = deployment_name
        except InvalidRequestError as ex:
            raise UnsupportedOpenAIModelError(
                "Model deployment name does not correspond to a valid model entity."
            ) from ex
        except APIConnectionError as ex:
            raise UnsupportedOpenAIModelError(
                f"Invalid Azure OpenAI Base Endpoint {api_base}"
            ) from ex

        self._set_params(**kwargs)

    @property
    def _default_params(self) -> Dict[str, Any]:

        """Get the default parameters for calling OpenAI API

        Returns: A dict of Default Params

        """
        return {**super()._default_params, "engine": self.engine}

    def call(self, instruction: str, value: str, suffix: str = "") -> str:

        """
        Call the Azure OpenAI LLM.

        Args:
            instruction (str): Instruction to pass
            value (str): Value to pass
            suffix(str): Suffix to pass

        Returns:
            str: Response
        """
        self.last_prompt = str(instruction) + str(value)

        if self.is_chat_model:
            response = self.chat_completion(str(instruction) + str(value) + suffix)
        else:
            response = self.completion(str(instruction) + str(value) + suffix)

        return response

    @property
    def type(self) -> str:
        return "azure-openai"

class GooglePalm(BaseGoogle):
    """Google Palm LLM
    BaseGoogle class is extended for Google Palm model. The default and only model support at the
    moment is models/text-bison-001.

    """

    model: str = "models/text-bison-001"

    def __init__(self, api_key: str, **kwargs):
        """
        __init__ method of GooglePalm Class
        Args:
            api_key (str): API Key
            **kwargs: Extended Parameters inferred from BaseGoogle Class
        """
        self._configure(api_key=api_key)
        self._set_params(**kwargs)

    def _valid_params(self):
        """Returns if the Parameters are valid or Not

        """
        return super()._valid_params() + ["model"]

    def _validate(self):
        """
        A method to Validate the Model

        """

        super()._validate()

        if not self.model:
            raise ValueError("model is required.")

    def _generate_text(self, prompt: str) -> str:
        """
        Generates text for prompt

        Args:
            prompt (str): Prompt

        Returns:
            str: LLM response
        """
        self._validate()
        completion = self.genai.generate_text(
            model=self.model,
            prompt=prompt,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            max_output_tokens=self.max_output_tokens,
        )
        return completion.result

    @property
    def type(self) -> str:
        return "google-palm"

class FakeLLM(LLM):
    """Fake LLM"""

    _output: str = 'print("Hello world")'

    def __init__(self, output: Optional[str] = None):
        if output is not None:
            self._output = output

    def call(self, instruction: Prompt, value: str, suffix: str = "") -> str:
        self.last_prompt = str(instruction) + str(value) + suffix
        return self._output

    @property
    def type(self) -> str:
        return "fake"