import datetime
import logging
import os
import time
from typing import Generator, Iterator, Union

import openai

from banterbot.config import RETRY_LIMIT, RETRY_TIME
from banterbot.data.enums import EnvVar
from banterbot.models.message import Message
from banterbot.models.openai_model import OpenAIModel
from banterbot.utils.nlp import NLP

# Set the OpenAI API key
openai.api_key = os.environ.get(EnvVar.OPENAI_API_KEY.value)


class OpenAIService:
    """
    A class that handles the interaction with the OpenAI ChatCompletion API. It provides functionality to generate
    responses from the API based on the input messages. It supports generating responses as a whole or as a stream of
    response blocks.

    The main purpose of this class is to facilitate the communication with the OpenAI API and handle the responses
    generated by the API. It can be used to create chatbots or other applications that require natural language
    processing and generation.
    """

    def __init__(self, model: OpenAIModel) -> None:
        """
        Initializes an `OpenAIService` instance for a specific model.

        Args:
            model (OpenAIModel): The OpenAI model to be used. This should be an instance of the OpenAIModel class, which
            contains information about the model, such as its name and maximum token limit.
        """
        logging.debug(f"OpenAIService initialized")

        # The selected model that will be used in OpenAI ChatCompletion prompts.
        self._model = model

        # Indicates whether the current instance of `OpenAIService` is streaming.
        self._streaming = False

        # Set the interruption flag to zero: if interruptions are raised, this will be updated.
        self._interrupt: int = 0

    def count_tokens(self, string: str) -> int:
        """
        Counts the number of tokens in the provided string.

        Args:
            string (str): A string provided by the user where the number of tokens are to be counted.

        Returns:
            int: The number of tokens in the string.
        """
        return len(self._model.tokenizer.encode(string))

    def interrupt(self) -> None:
        """
        Interrupts any ongoing streaming processes. This method sets the interrupt flag to the current time, which will
        cause any streaming processes activated prior to the current time to stop.
        """
        self._interrupt: int = time.perf_counter_ns()
        logging.debug("OpenAIService stream interrupted")

    def prompt(self, messages: list[Message], split: bool = True, **kwargs) -> Union[tuple[str], str]:
        """
        Sends messages to the OpenAI ChatCompletion API and retrieves the response as a list of sentences.

        Args:
            messages (list[Message]): A list of messages. Each message should be an instance of the Message class, which
            contains the content and role (user or assistant) of the message.

            split (bool): Whether the response should be split into sentences.

            **kwargs: Additional parameters for the API request. These can include settings such as temperature, top_p,
            and frequency_penalty.

        Returns:
            Union[list[str], str]: A list of sentences forming the response from the OpenAI API. This can be used to
            display the generated response to the user or for further processing. If `split` is False, returns a string.
        """
        response = self._request(messages=messages, stream=False, **kwargs)
        logging.debug(f"OpenAIService stream processed block: `{response}`")
        sentences = NLP.segment_sentences(response) if split else response
        return sentences

    def prompt_stream(self, messages: list[Message], **kwargs) -> Generator[tuple[str], None, None]:
        """
        Sends messages to the OpenAI API and retrieves the response as a stream of blocks of sentences.

        Args:
            messages (list[Message]): A list of messages. Each message should be an instance of the Message class, which
            contains the content and role (user or assistant) of the message.

            **kwargs: Additional parameters for the API request. These can include settings such as temperature, top_p,
            and frequency_penalty.

        Yields:
            Generator[list[str], None, None]: A stream of blocks of sentences as the response from the OpenAI API. Each
            block contains one or more sentences that form a part of the generated response. This can be used to display
            the response to the user in real-time or for further processing.
        """
        # Record the time at which the request was made, in order to account for future interruptions.
        init_time = time.perf_counter_ns()

        # Obtain a response from the OpenAI ChatCompletion API
        response = self._request(messages=messages, stream=True, **kwargs)

        # Set the streaming flag to True
        self._streaming = True

        # Yield the responses as they are streamed
        for block in self._response_parse_stream(response=response, init_time=init_time):
            yield block

        # Reset the streaming flag to False
        self._streaming = False

    @property
    def model(self) -> OpenAIModel:
        """
        Return the `OpenAIModel` associated with the current instance.

        Returns:
            OpenAIModel
        """
        return self._model

    @property
    def streaming(self) -> bool:
        """
        If the current instance of OpenAIService is in the process of streaming, returns True. Otherwise, returns False.

        Args:
            bool: The streaming state of the current instance.
        """
        return self._streaming

    def _response_parse_stream(self, response: Iterator, init_time: int) -> Generator[list[str], None, None]:
        """
        Parses a streaming response from the OpenAI API and yields blocks of text as they are received.

        Args:
            response (Iterator): A streaming response from the OpenAI ChatCompletion API.
            init_time (int): The time at which the stream was initialized.

        Yields:
            Generator[list[str], None, bool]: Lists of sentences as blocks. Each block contains one or more sentences
            that form a part of the generated response.
        """
        text = ""
        logging.debug("OpenAIService stream started")

        for chunk in response:
            delta = chunk["choices"][0]["delta"]

            if self._interrupt >= init_time:
                logging.debug(f"OpenAIService interrupted")
                break

            if "content" in delta.keys():
                text += delta["content"]

            if len(sentences := NLP.segment_sentences(text)) > 1:
                text = sentences[-1]
                logging.debug(f"OpenAIService yielded sentences: {sentences[:-1]}")
                yield sentences[:-1]
        else:
            sentences = NLP.segment_sentences(text)
            logging.debug(f"OpenAIService yielded final sentences: {sentences[:-1]}")
            yield sentences

            logging.debug("OpenAIService stream stopped")

    def _request(self, messages: list[Message], stream: bool, **kwargs) -> Union[Iterator, str]:
        """
        Sends a request to the OpenAI API and generates a response based on the specified parameters.

        Args:
            messages (list[Message]): A list of messages. Each message should be an instance of the Message class, which
            contains the content and role (user or assistant) of the message.

            stream (bool): Whether the response should be returned as an iterable stream or a complete text.

            **kwargs: Additional parameters for the API request. These can include settings such as temperature, top_p,
            and frequency_penalty.

        Returns:
            Union[Iterator, str]: The response from the OpenAI API, either as a stream (Iterator) or text (str).
        """
        kwargs["model"] = self._model.model
        kwargs["n"] = 1
        kwargs["stream"] = stream
        kwargs["messages"] = [message() for message in messages]

        success = False
        for i in range(RETRY_LIMIT):
            try:
                response = openai.ChatCompletion.create(**kwargs)
                success = True
                break

            except openai.error.APIError:
                retry_time = 0.25
                retry_timestamp = datetime.datetime.now() + datetime.timedelta(seconds=retry_time)
                retry_timestamp = retry_timestamp.strptime("%H:%M:%S")
                error_message = (
                    f"OpenAIService encountered an OpenAI API Error - Attempt {i+1}/{RETRY_LIMIT}. Waiting "
                    f"{retry_time} seconds until {retry_timestamp} to retry."
                )
                logging.info(error_message)
                time.sleep(retry_time)

            except openai.error.RateLimitError:
                retry_timestamp = datetime.datetime.now() + datetime.timedelta(seconds=RETRY_TIME)
                retry_timestamp = datetime.datetime.strftime(retry_timestamp, "%H:%M:%S")
                error_message = (
                    f"OpenAIService encountered an OpenAI Rate Limiting Error - Attempt {i+1}/{RETRY_LIMIT}. Waiting "
                    f"{RETRY_TIME} seconds until {retry_timestamp} to retry."
                )
                logging.info(error_message)
                time.sleep(RETRY_TIME)

        if not success:
            logging.info(f"OpenAIService encountered too many OpenAI Errors - exiting program.")
            raise openai.error.APIError

        return response if stream else response.choices[0].message.content.strip()
