# Provides TeLLMgramBot a way to store and retrieve conversation messages by the OpenAI model
import os
import re
import sys

from .tokenGPT import TokenGPT
from .utils import generate_filename, generate_file_path, generate_error_path
from .utils import get_safe_username, read_reverse_order, append_text_file_path, log_error

# noinspection RegExpRedundantEscape,PyTypeChecker
class Conversation:
    """
    Defines a user-assistant interaction based on the assistant's Generative AI model and prompt to contextualize
    with message history. Like memory, if reaching the maximum amount of tokens, prunes earliest messages under limit.
    """
    def __init__(self, user_name: str, assist_name: str, system_content: str, system_model="gpt-4o-mini"):
        self.error_log = generate_error_path()

        # Define who's talking for the two-way conversation
        self.user_name = get_safe_username(user_name)
        self.assist_name = assist_name
        self.names_print = f"User {self.user_name} & Assistant {self.assist_name}"

        # System defines how the assistant will respond to the user by prompt and GPT model
        self.system_content = system_content
        self.system_model = TokenGPT(system_model)
        self.messages = [{"role": "system", "content": system_content}]

        # Define the directory where to store this conversation with log file name
        # Variables set for data privacy only being reported on error messages
        env_var = 'TELLMGRAMBOT_SESSIONLOGS_PATH'
        if not os.environ.get(env_var):
            sys.exit(f"{env_var} must exist during conversation! Exiting...")
        self._interaction_dir = os.environ[env_var]
        self._interaction_log = generate_filename(self.assist_name, self.user_name)
        self._interaction_path = generate_file_path(self._interaction_dir, self._interaction_log, "Conversation")

        # Private boolean determines if conversation includes past messages
        # This will be true if user calls out function get_past_interaction()
        self._has_past_interaction = False

    def set_system_content(self, new_content: str):
        """Replace the system content (e.g. bot's prompt) as the conversation's first message."""
        self.system_content = new_content
        self.messages[0] = {"role": "system", "content": new_content}

    def add_user_message(self, content: str):
        """Add user message and write to this conversation's log file."""
        self.messages.append({"role": "user", "content": content})
        self.write_interaction(role="user", content=content)

    def add_assistant_message(self, content: str):
        """Add assistant message and write to this conversation's log file."""
        self.messages.append({"role": "assistant", "content": content})
        self.write_interaction(role="assistant", content=content)

    def get_message_token_count(self):
        """Get current token count of the user-assistant message list."""
        return self.system_model.num_tokens_from_messages(self.messages)

    def write_interaction(self, role: str, content: str):
        """
        Write the current content by role (user/assistant) to a conversation log file.
        This addresses UTF-8 characters, especially on symbols like integrals.
        """
        try:
            text = f'<{role}>: {content}'
            if not os.path.exists(self._interaction_path):
                generate_file_path(self._interaction_dir, self._interaction_log, "Conversation", text)
            else:
                append_text_file_path(self._interaction_path, text)
        except Exception as e:
            log_error(e, f"{type(e).__name__} on '{self._interaction_path}'", self.error_log)

    def prune_conversation(self, token_limit: int):
        """
        Remove earliest user-assistant messages (except the first system content message)
        until the token size is below threshold. Similar to forgetting the oldest memories.
        """
        try:
            self.messages.pop(1)
            if self.get_message_token_count() > token_limit:
                self.prune_conversation(token_limit)
        except Exception as e:
            log_error(e, f"{type(e).__name__} pruning under {token_limit} tokens for {self.names_print}", self.error_log)

    def _get_interaction_files(self):
        """Get all conversation log files of the current user-assistant sorted by timestamp."""
        files = []
        for file in sorted(os.listdir(self._interaction_dir)):
            # Check if full names of user and assistant are in the log name
            if file.startswith(f"{self.assist_name}-{self.user_name}-"):
                files.append(file)
        return files

    def get_past_interaction(self, token_limit: int) -> bool:
        """
        Read historical conversation logs to store current user-assistant messages up to a token limit.
        Useful to remember conversations since any maintenance may require restart on the assistant.
        Returns boolean if there are past interactions, even if it was already checked.
        """
        if not self._has_past_interaction:
            # Start looping each log file with timestamp in descending order
            token_limit_reached = 0
            message_count = len(self.messages)
            token_count = 0
            for log_file in sorted(self._get_interaction_files(), reverse=True):
                # Skip if on the current conversation log file
                if log_file == self._interaction_log:
                    continue

                # Iterate through the log file backwards by line and populate content for each role
                # Role and content should only be filed if the line contains a user/assistant field
                role = ''
                content = ''
                for line in read_reverse_order(os.path.join(self._interaction_dir, log_file)):
                    # Append content continuously until the role name is defined by either format:
                    #   <user>: ...
                    #   <assistant>: ...
                    content = '\n' + content.strip()
                    result = re.search(r'^\<(user|assistant)\>: (.*)', line)
                    if result is not None:
                        role = result.group(1)
                        content = result.group(2) + content
                    else:
                        content = line + content

                    # If both role and content are defined, insert as dictionary
                    if role and content:
                        # Set message dictionary after the first message of system content going backwards
                        self.messages.insert(1, {"role": role, "content": content})
                        self._has_past_interaction = True

                        # Reset message parameters for next round if any
                        role = ''
                        content = ''

                        # See if the series of messages does actually fit by token limit
                        token_count = self.get_message_token_count()
                        if token_count > token_limit:
                            self.messages.pop(1)  # Remove message inserted recently
                            token_count = self.get_message_token_count()
                            token_limit_reached = 1
                            break

                # Stop going through logs if the token limit has reached
                if token_limit_reached:
                    break

            # Print in terminal if there were conversations via session logs
            if message_count != len(self.messages):
                print(
                    f"{self.names_print} had past conversations "
                    f"{'above' if token_limit_reached else 'within'} the "
                    f"token limit {token_limit}, storing {token_count} token(s)"
                )
            else:
                print(f"{self.names_print} are in their first conversation")
        else:
            print(f"{self.names_print} had their past conversations defined")

        return self._has_past_interaction

    def clear_interaction(self):
        """
        Clear out all user-assistant interactions by:
        - Deleting list of all messages except the first system content.
        - Removing past conversation logs and reset on an empty log file.
        """
        self.messages = [self.messages[0]]
        self._has_past_interaction = False
        for file in self._get_interaction_files():
            log_path = os.path.join(self._interaction_dir, file)
            if os.path.isfile(log_path):
                os.remove(log_path)
            else:
                log_error("File not found to delete", f"Error on '{log_path}'", self.error_log)
