import uuid
from datetime import datetime
from typing import Optional

import streamlit as st

from synthetic_agents.model.agent import Agent
from synthetic_agents.model.entity.world import TextMessage


class ChatSession:
    """
    This class represents a web session and comprises different methods for storing and
    retrieving data within an online chat session.
    """

    _SESSION_ID_KEY = "session_id"
    _CHAT_MESSAGES_KEY = "chat_messages"
    _USER_KEY = "user"
    _COACH_KEY = "coach"
    _AUTOMATED_CHAT_IN_PROGRESS_KEY = "automated_chat_in_progress"
    _MAX_MESSAGES_KEY = "max_messages"
    _LAST_AGENT_MESSAGE_INDEX_KEY = "last_agent_message_index"
    _CURRENT_MANUAL_MESSAGE_START_TIME = "current_manual_message_start_time"

    ALL_KEYS = [
        _SESSION_ID_KEY,
        _CHAT_MESSAGES_KEY,
        _USER_KEY,
        _COACH_KEY,
        _AUTOMATED_CHAT_IN_PROGRESS_KEY,
        _MAX_MESSAGES_KEY,
        _LAST_AGENT_MESSAGE_INDEX_KEY,
        _CURRENT_MANUAL_MESSAGE_START_TIME,
    ]

    @staticmethod
    def reset():
        """
        Clears all information stored in the session.
        """
        for key in ChatSession.ALL_KEYS:
            del st.session_state[key]
        ChatSession.initialize()

    @staticmethod
    def initialize():
        """
        Creates a chat session and initializes the streamlit session object with the required keys.
        """
        if ChatSession._SESSION_ID_KEY not in st.session_state:
            st.session_state[ChatSession._SESSION_ID_KEY] = str(uuid.uuid4())

        none_keys = [
            ChatSession._USER_KEY,
            ChatSession._COACH_KEY,
            ChatSession._CURRENT_MANUAL_MESSAGE_START_TIME,
        ]

        for key in none_keys:
            if key not in st.session_state:
                st.session_state[key] = None

        list_keys = [ChatSession._CHAT_MESSAGES_KEY]

        for key in list_keys:
            if key not in st.session_state:
                st.session_state[key] = []

        if ChatSession._AUTOMATED_CHAT_IN_PROGRESS_KEY not in st.session_state:
            st.session_state[ChatSession._AUTOMATED_CHAT_IN_PROGRESS_KEY] = False

        if ChatSession._LAST_AGENT_MESSAGE_INDEX_KEY not in st.session_state:
            st.session_state[ChatSession._LAST_AGENT_MESSAGE_INDEX_KEY] = {}

        if ChatSession._MAX_MESSAGES_KEY not in st.session_state:
            st.session_state[ChatSession._MAX_MESSAGES_KEY] = 0

    @staticmethod
    def get_session_id() -> str:
        """
        Gets the ID of the chat session.

        :return: unique ID of the session.
        """
        return st.session_state[ChatSession._SESSION_ID_KEY]

    @staticmethod
    def store_user(user: Agent):
        """
        Stores the user agent in the session.

        :param user: user agent to be stored.
        """

        st.session_state[ChatSession._USER_KEY] = user

    @staticmethod
    def store_coach(coach: Agent):
        """
        Stores the coach agent in the session.

        :param coach: coach agent to be stored.
        """

        st.session_state[ChatSession._COACH_KEY] = coach

    @staticmethod
    def store_chat_message(agent: Agent, message: TextMessage):
        """
        Adds to the list of chat messages a new message produced by an agent.

        :param agent: agent who produced the message.
        :param message: message from the agent.
        """

        st.session_state[ChatSession._CHAT_MESSAGES_KEY].append(message)

        # Index of the last message produced by the agent in the list of chat messages.
        st.session_state[ChatSession._LAST_AGENT_MESSAGE_INDEX_KEY][agent.agent_id] = (
            len(st.session_state[ChatSession._CHAT_MESSAGES_KEY]) - 1
        )

    @staticmethod
    def store_current_manual_message_start_time(timestamp: datetime):
        """
        Stores the timestamp when the current message started to being produced. To be used by
        manual messages entered by a human agent. It assumes it starts as soon as the message from
        the other agent is delivered to the human.

        :param timestamp: timestamp when the current message started to being produced.
        """
        st.session_state[ChatSession._CURRENT_MANUAL_MESSAGE_START_TIME] = timestamp

    @staticmethod
    def get_messages() -> list[TextMessage]:
        """
        Gets all messages in the chat history.

        :return: a list of chat messages exchanged so far.
        """
        return st.session_state[ChatSession._CHAT_MESSAGES_KEY]

    @staticmethod
    def get_last_message_from_agent(agent: Agent) -> Optional[TextMessage]:
        """
        Gets the last message produced by an agent in the chat history.

        :param agent: agent that produced the message.
        :return: last message produced by the agent. None if the agent hasn't produced any message
            yet.
        """

        if agent.agent_id not in st.session_state[ChatSession._LAST_AGENT_MESSAGE_INDEX_KEY]:
            return None

        message_index = st.session_state[ChatSession._LAST_AGENT_MESSAGE_INDEX_KEY][agent.agent_id]
        return st.session_state[ChatSession._CHAT_MESSAGES_KEY][message_index]

    @staticmethod
    def is_automated_chat_in_progress() -> bool:
        """
        Gets flag indicating whether the automated chat is in progress.

        :return: true if the automated chat is in progress, false otherwise.
        """
        return st.session_state[ChatSession._AUTOMATED_CHAT_IN_PROGRESS_KEY]

    @staticmethod
    def get_user() -> Agent:
        """
        Gets the user agent.

        :return: user agent.
        """
        return st.session_state[ChatSession._USER_KEY]

    @staticmethod
    def get_coach() -> Agent:
        """
        Gets the coach agent.

        :return: coach agent.
        """
        return st.session_state[ChatSession._COACH_KEY]

    @staticmethod
    def toggle_automated_chat(chat_on: bool):
        """
        Stores a boolean indicating whether the automated chat is running.

        :param chat_on: true if the chat is running.
        """
        st.session_state[ChatSession._AUTOMATED_CHAT_IN_PROGRESS_KEY] = chat_on

    @staticmethod
    def has_messages() -> bool:
        """
        Checks whether agents have produced any message.

        :return true if any agent has produced any message.
        """
        return len(st.session_state[ChatSession._CHAT_MESSAGES_KEY]) > 0

    @staticmethod
    def get_last_speaker_type() -> Optional[str]:
        """
        Gets last agent type that delivered a chat message.

        :return last agent type that delivered a chat message.
        """
        if not ChatSession.has_messages():
            return None

        return st.session_state[ChatSession._CHAT_MESSAGES_KEY][-1].agent_type

    @staticmethod
    def get_current_manual_message_start_time() -> datetime:
        """
        Gets the timestamp when the current message started to being produced.

        :return: timestamp when the current message started to being produced.
        """
        return st.session_state[ChatSession._CURRENT_MANUAL_MESSAGE_START_TIME]
