# import asyncio
# import socketio
# import logging

# # Configure logging
# logging.basicConfig(level=logging.DEBUG)
# logger = logging.getLogger(__name__)

# class RasaSocketIOClient:
#     def __init__(self, uri):
#         self.sio = socketio.AsyncClient(logger=True, engineio_logger=True)  # Enable internal logging
#         self.uri = uri
#         self.response = None

#         @self.sio.on('bot_uttered')  # Adjust the event name as per your Rasa setup
#         async def on_message(data):
#             logger.info("Received response: %s", data)
#             self.response = data  # Store the latest response

#     async def connect(self):
#         try:
#             await self.sio.connect(self.uri)
#             logger.info("Successfully connected to the server.")
#         except Exception as e:
#             logger.error("Failed to connect to the server: %s", e)

#     async def send_message(self, message):
#         self.response = None  # Reset previous response
#         try:
#             await self.sio.emit('user_uttered', message)  # Adjust the event name as per your Rasa setup
#             logger.debug("Message sent: %s", message)
#             await asyncio.sleep(1)  # Wait for the response to be received
#         except Exception as e:
#             logger.error("Failed to send message: %s", e)
#         return self.response

#     async def disconnect(self):
#         try:
#             await self.sio.disconnect()
#             logger.info("Disconnected from the server.")
#         except Exception as e:
#             logger.error("Failed to disconnect: %s", e)


# import asyncio
# import socketio
# import logging
# from custom_logging import IMPORTANT_LEVEL_NUM 

# logger = logging.getLogger(__name__)

# class RasaSocketIOClient:
#     def __init__(self, uri):
#         self.sio = socketio.AsyncClient(logger=False, engineio_logger=False)
#         self.response_event = asyncio.Event()
#         self.uri = uri
#         self.response = None

#         @self.sio.on('bot_uttered')  # Adjust the event name as per your Rasa setup
#         async def on_message(data):
#             # Use the new logging level
#             logger.important("Socket ID: %s Received response: %s", self.sio.sid, data)
#             self.response = data  # Store the latest response

#     async def connect(self):
#         try:
#             await self.sio.connect(self.uri)
#             # Use the new logging level
#             logger.important("Successfully connected to the server.")
#         except Exception as e:
#             # Use the new logging level for errors too, for consistency
#             logger.important("Failed to connect to the server: %s", e)
#             return "Failed to connect to the server."

#     async def send_message(self, message):
#         self.response = None  # Reset previous response
#         try:
#             await self.sio.emit('user_uttered', message)  # Adjust the event name as per your Rasa setup
#             # Since this is more of a debug level, consider if you want to keep it under IMPORTANT or not
#             logger.important("Message sent: %s", message)
#             await asyncio.sleep(1)  # Wait for the response to be received
#         except Exception as e:
#             logger.important("Failed to send message: %s", e)
#         return self.response

#     async def disconnect(self):
#         try:
#             await self.sio.disconnect()
#             logger.important("Disconnected from the server.")
#         except Exception as e:
#             logger.important("Failed to disconnect: %s", e)


import asyncio
import socketio
import logging
import custom_logging


# Obtain a logger for this module
logger = logging.getLogger(__name__)

class RasaSocketIOClient:
    def __init__(self, uri, update_chat_history_callback=None):
        self.sio = socketio.AsyncClient(logger=False, engineio_logger=False)
        self.uri = uri
        self.message_queue = asyncio.Queue()
        self.register_event_handlers()

    def register_event_handlers(self):
        @self.sio.event
        async def connect():
            logger.important("Connected to the server.")

        @self.sio.event
        async def disconnect():
            logger.important("Disconnected from the server.")

        @self.sio.on('bot_uttered')
        async def on_message(data):
            logger.debug(f"Socket ID: {self.sio.sid} Received response: {data}")
            await self.message_queue.put(data)  # Put received message into the queue

    async def connect(self):
        try:
            await self.sio.connect(self.uri)
            logger.important("Successfully connected to the server.")
        except Exception as e:
            logger.error(f"Failed to connect to the server: {e}")

    async def send_message(self, user_message):
        try:
            await self.sio.emit('user_uttered', {'message': user_message})
            logger.important(f"Message sent: {user_message}")
            timeout = 5.0  # Timeout in seconds after the last message is received
            while True:
                try:
                    data = await asyncio.wait_for(self.message_queue.get(), timeout)
                    logger.important(f"Message in queue: {user_message}")
                    yield data
                except asyncio.TimeoutError:
                    logger.important(f"send_message function timed out")
                    break
        except Exception as e:
            logger.error(f"Failed to send message: {e}")

    async def disconnect(self):
        try:
            logger.important(f"Attempting to disconnect: {id(self.sio)}")
            await self.sio.disconnect()
            logger.important("Disconnected from the server.")
        except Exception as e:
            logger.error(f"Failed to disconnect: {e}")