import uuid
from pymongo import MongoClient
from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory


class ChatManagement:
    def __init__(self, cluster_url, database_name, collection_name):
        """
        Initializes the ChatManagement class with MongoDB connection details.
        
        Args:
            cluster_url (str): MongoDB cluster URL.
            database_name (str): Name of the database.
            collection_name (str): Name of the collection.
        """
        self.connection_string = cluster_url
        self.database_name = database_name
        self.collection_name = collection_name
        self.chat_sessions = {}  # Dictionary to store chat history objects for each session

    def create_new_chat(self):
        """
        Creates a new chat session by initializing a MongoDBChatMessageHistory object.
        
        Returns:
            str: The unique chat ID.
        """
        # Generate a unique chat ID
        chat_id = str(uuid.uuid4())

        # Initialize MongoDBChatMessageHistory for the chat session
        chat_message_history = MongoDBChatMessageHistory(
            session_id=chat_id,
            connection_string=self.connection_string,
            database_name=self.database_name,
            collection_name=self.collection_name,
        )

        # Store the chat_message_history object in the session dictionary
        self.chat_sessions[chat_id] = chat_message_history
        return chat_id

    def get_chat_history(self, chat_id):
        """
        Retrieves the MongoDBChatMessageHistory object for a given chat session by its chat ID.
        
        Args:
            chat_id (str): The unique ID of the chat session.
        
        Returns:
            MongoDBChatMessageHistory or None: The chat history object of the chat session, or None if not found.
        """
        # Check if the chat session is already in memory
        if chat_id in self.chat_sessions:
            return self.chat_sessions[chat_id]
        
        # If not in memory, try to fetch from the database
        chat_message_history = MongoDBChatMessageHistory(
            session_id=chat_id,
            connection_string=self.connection_string,
            database_name=self.database_name,
            collection_name=self.collection_name,
        )
        if chat_message_history.messages:  # Check if the session exists in the database
            self.chat_sessions[chat_id] = chat_message_history
            return chat_message_history
        
        return None  # Chat session not found

    def initialize_chat_history(self, chat_id):
        """
        Initializes a new chat history for the given chat ID if it does not already exist.
        
        Args:
            chat_id (str): The unique ID of the chat session.
        
        Returns:
            MongoDBChatMessageHistory: The initialized chat history object.
        """
        # If the chat history already exists, return it
        if chat_id in self.chat_sessions:
            return self.chat_sessions[chat_id]

        # Otherwise, create a new chat history
        chat_message_history = MongoDBChatMessageHistory(
            session_id=chat_id,
            connection_string=self.connection_string,
            database_name=self.database_name,
            collection_name=self.collection_name,
        )

        # Save the new chat session to the session dictionary
        self.chat_sessions[chat_id] = chat_message_history
        return chat_message_history