from .. import constants,exceptions 
from websockets.asyncio.client import connect
import aiohttp

class CRWritableChat(constants.CRChat): # Expansion of the CRChat class to add methods for sending messages and deleting chats or messages.
    """
        A writable chat from ChattedRooms.
        This class is used to represent a chat in the ChattedRooms API.
        It contains the chat ID, bot ID, visibility, and messages.
    """
    def __init__(self,_client, id: int, botid: int, visibility: int, messages: list['constants.CRChatMessage']):
        super().__init__(id, botid, visibility, messages)
        self._client = _client
        self.base_url = _client.base_url
    async def send_message(self, message, streaming=False):
        """
        Send a message to the ChattedRooms API.
        This method is used to send a message to the ChattedRooms API.
        Args:
            message (str): The message to send.
        Returns:
            If streaming is False, returns the full message as a string.
            If streaming is True, returns an async iterator of message chunks (strings).
        """
        async def stream_messages():
            try:
                async with connect("wss://api.chattedrooms.com/api/v2/llm/generateStreaming") as websocket:
                    await websocket.send({
                        "togenerate": message,
                        "chat_id": self.id,
                        "char_id": self.botid,
                        "key": self._client.headers["X-Authorization-CR"]
                    })
                    async for chunk in websocket:
                        yield chunk
            except Exception as e:
                raise exceptions.MessageSentError(f"Streaming message failed: {e}")
        if streaming:
            return stream_messages()
        else:
            finalMessage = ""
            async with self._client.session.post("https://api.chattedrooms.com/api/v2/llm/generate", json={
                "togenerate": message,
                "chat_id": self.id,
                "char_id": self.botid,
            },headers=self._client.headers) as response:
                if response.status != 200:
                    raise exceptions.MessageSentError(f"Non-streaming message failed with status {response.status}")
                data = await response.json()
                finalMessage = data.get("generated", "")
            return finalMessage

class ChatMan():
    def __init__(self,client):
        self._client = client
    async def fetch_recent_chats(self):
        """
        Fetch the recent chats from the ChattedRooms API.
        This method is used to get the recent chats of the user.
        Returns:
            list: A list of recent chats.
        """
        async with self._client.session.get(f"{self._client.base_url}/api/v1/users/getChats/", headers=self._client.headers) as response:
            if response.status == 200:
                resp = await response.json()
                returned = [constants.CRChatShort(chat["id"], chat["bot_id"], 0) for chat in resp]
                return returned
            elif response.status == 403 or response.status == 401:
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise Exception(f"Failed to fetch recent chats: {response.status}")
    
    async def get_chat_from_id(self,bot_id,id):
        """
        Fetch a chat from the ChattedRooms API.
        This method is used to get a chat by its ID.
        Args:
            bot_id (int): The ID of the bot.
            id (int): The ID of the chat.
        Returns:
            CRChat: The chat object.
        """
        async with self._client.session.get(f"{self._client.base_url}/api/v1/chats/getChat/{bot_id}/{id}", headers=self._client.headers) as response:
            if response.status == 200:
                resp = await response.json()
                messages = [constants.CRChatMessage(message["content"], message["role"], message["id"]) for message in resp["messages"]]
                return CRWritableChat(id,self._client, resp["botid"], resp["visibility"], messages)
            elif response.status == 403 or response.status == 401:
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise Exception(f"Failed to fetch chat: {response.status}")
            
    async def delete_chat(self, char_id, chat_id):
        """
        Delete a chat from the ChattedRooms API.
        This method is used to delete a chat by its ID.
        Args:
            char_id (int): The ID of the character.
            chat_id (int): The ID of the chat.
        Returns:
            bool: True if the chat was deleted successfully, False otherwise.
        """
        async with self._client.session.delete(f"{self.base_url}/api/v1/chat/delchat/{char_id}/{chat_id}", headers=self._client.headers) as response:
            if response.status == 200:
                return True
            elif response.status == 403 or response.status == 401:
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise exceptions.ChatDeletionError(f"Failed to delete chat: {response.status}")

    # PyCharacterAI Polyfills Start:

    # Middleman method for PyCharacterAI compatibility.
    async def fetch_chat(self,chat_id):
        """
        Fetch a chat from the ChattedRooms API.
        This method is used to get a chat by its ID. (completely bypassing the requirement for bot_id)
        Args:
            id (int): The ID of the chat.
        Returns:
            CRWritableChat: The chat object.
        """
        async with self._client.session.get(f"{self.base_url}/api/v1/chat/getBotIDFromChatID/{chat_id}", headers=self._client.headers) as response:
            if response.status == 200:
                data = await response.json()
                bot_id = data.get("botid")
            if bot_id is None:
                raise exceptions.ChatError("botid not found in response")
            elif response.status in (401, 403):
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise exceptions.ChatError(f"Failed to fetch botid: {response.status}")
        id = chat_id
        return await self.get_chat_from_id(bot_id,id)
    
    # Middleman method for PyCharacterAI compatibility. Archiving doesn't exist in ChattedRooms, so this just deletes the chat instead.
    async def archive_chat(self,chat_id):
        """
        Archive a chat from the ChattedRooms API.
        This method is used to delete a chat by its ID.
        Args:
            id (int): The ID of the chat.
        Returns:
            bool: True if the chat was archived successfully, False otherwise.
        """
        async with self._client.session.get(f"{self.base_url}/api/v1/chat/getBotIDFromChatID/{chat_id}", headers=self._client.headers) as response:
            if response.status == 200:
                data = await response.json()
                bot_id = data.get("botid")
            if bot_id is None:
                raise exceptions.ChatError("botid not found in response")
            elif response.status in (401, 403):
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise exceptions.ChatError(f"Failed to fetch botid: {response.status}")
        return await self.delete_chat(chat_id,chat_id)
    
    # unarchive chat. This is a no-op.
    async def unarchive_chat(self,chat_id):
        """
        Unarchive a chat from the ChattedRooms API.
        This method is used to unarchive a chat by its ID.
        (WARNING! No-Op.)
        Args:
            id (int): The ID of the chat.
        Returns:
            bool: always True. This is a no-op.
        """
        chat_id = chat_id # Prevent pylance from complaining about unused variable.
        return True
    
    # Middleman method for PyCharacterAI compatibility.
    async def fetch_all_messages(self,chat_id):
        """
        Fetch all messages from a chat.
        This method is used to get all messages from a chat.
        Args:
            chat_id (int): The ID of the chat.
        Returns:
            list: A list of messages.
        """
        async with self._client.session.get(f"{self.base_url}/api/v1/chat/getBotIDFromChatID/{chat_id}", headers=self._client.headers) as response:
            if response.status == 200:
                data = await response.json()
                bot_id = data.get("botid")
            if bot_id is None:
                raise exceptions.ChatError("botid not found in response")
            elif response.status in (401, 403):
                raise exceptions.Unauthorized("Invalid current token. Maybe you have forgot to run self.authenticate()?")
            else:
                raise exceptions.ChatError(f"Failed to fetch botid: {response.status}")
        return await self.get_chat_from_id(bot_id,chat_id).messages
    
    # Polyfills End:



    async def makeWritable(self,chat:constants.CRChat):
        """
        Make a chat writable.
        This method is used to make a chat writable.
        Args:
            chat (CRChat): The chat object.
        Returns:
            CRWritableChat: The writable chat object.
        """
        return CRWritableChat(chat.id,self._client, chat.botid, chat.visibility, chat.messages)
