import asyncio
from datetime import datetime

import database as db 
from chat import ChatView, Message, MESSAGE_PAGE_SIZE
from stream import Stream

class ChatMixin:
    async def create_chat(self, user_id, name, description, picture, members=[]):
        # Tortoise uses .create() for one-line instantiation and saving
        chat = await db.Chat.create(
            user_id=user_id,
            name=name,
            description=description,
            picture=picture
        )
        print(f"Created chat #{chat.id}: {chat.name}")

        for agent_id in members:
            await self.add_agent_to_chat(chat.id, agent_id)

        return chat.id

    async def get_chats(self, user_id):
        # Direct filtering on the model class
        chats = await db.Chat.filter(user_id=user_id).all()
        return [{
            "id": chat.id,
            "name": chat.name,
            "description": chat.description,
            "picture": chat.picture
        } for chat in chats]

    async def delete_chat(self, chat_id):
        # Filtering and deleting is a one-liner
        await db.Membership.filter(chat_id=chat_id).delete()
        await db.Chat.filter(id=chat_id).delete()
        # TODO: Delete messages from the chat
        print(f"Deleted chat #{chat_id} and its memberships.")

    async def get_chat_messages(self, chat_id, earliest_index=None):
        query = db.Message.filter(chat_id=chat_id)
        if earliest_index is not None:
            query = query.filter(id__gt=earliest_index)

        # Order quey by time descending (id orders better)
        # then reverse in Python to get chronological order
        messages = await query.order_by("-id").all()
        messages.reverse()

        return [{
            "id": message.id,
            "time": message.time,
            "sender": message.sender,
            "kind": message.kind,
            "content": message.content,
            "threadId": message.thread_id,
        } for message in messages]
    
    async def create_chat_thread(self, chat_id, name=None):
        thread = await db.Thread.create(
            chat_id=chat_id,
            name=name
        )
        print(f"Created thread #{thread.id}: {thread.name}")
        return thread.id

    async def add_agent_to_chat(self, chat_id, agent_id):
        await db.Membership.create(
            chat_id=chat_id,
            agent_id=agent_id
        )
        print(f"Added agent #{agent_id} to chat #{chat_id}")

    async def send_chat_message(self, chat_id, content, sender='user', thread_id=None, kind=None, subkind=None, format=None):
        now = int(datetime.now().timestamp())

        if isinstance(content, Stream):
            initial_content = "..."
        else:
            initial_content = content

        # Set to true for streams
        updating = False
        message = await db.Message.create(
            chat_id=chat_id,
            time=now,
            last_edit=now,
            sender=sender,
            kind=kind or 'message',
            subkind=subkind,
            format=format,
            content=initial_content,
            thread_id=thread_id
        )

        # Fetch chat owner using values_list for efficiency
        chat = await db.Chat.get(id=chat_id).prefetch_related("user")
        chat_owner_id = chat.user_id

        # Notification order depends on whether the content is a stream

        async def notify_agents():
            # Fetch participant agents
            agent_ids = await db.Membership.filter(chat_id=chat_id).values_list("agent_id", flat=True)
            for agent_id in agent_ids:
                if str(agent_id) != str(sender):
                    await self.notify_message_to_agent(agent_id, message)
        
        def notify_users():
            nonlocal updating
            self.notify_message_to_user(chat_owner_id, message, updating)
            updating = True

        if isinstance(content, Stream):
            notify_users()

            def chunk_fn (chunk):
                message.content = content.content + "..."
                notify_users()
            async def end_fn ():
                message.content = content.content
                await message.save()
                await notify_agents()
                notify_users()

            content.listen(chunk_fn=chunk_fn, end_fn=end_fn)
        else:
            notify_users()
            await notify_agents()
        
        # TODO: Chat preview based on the message

        # TODO: return for who lmao
        return {
            'id': message.id,
            'chatId': message.chat_id,
            'time': message.time,
            'sender': message.sender,
            'kind': message.kind,
            'content': message.content,
        }

    async def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
        message = await db.Message.get_or_none(id=message_id)

        if not message:
            raise Exception(f"Message not found with id {message_id}")

        summary = []
        if content is not None:
            message.content = content
            summary.append('content')
        if kind is not None:
            message.kind = kind
            message.subkind = subkind
            summary.append('kind')
            summary.append('subkind')
        if format is not None:
            message.format = format
            summary.append('format')
        
        if summary:
            await message.save() # Saves the specific instance
            print(f"Updated message #{message.id} fields: {summary}")

    async def get_chat_view(self, chat_id):
        chat_view = ChatView(self, chat_id)
        await chat_view.init_load()
        return chat_view

    def notify_message_to_user(self, user_id, message, update=False):
        if user_id in self.user_connections:
            user_message = {
                'action': 'chat-message',
                'id': message.id,
                'chatId': message.chat_id,
                'time': message.time,
                'sender': message.sender,
                'kind': message.kind,
                'content': message.content,
                'threadId': message.thread_id,
                'update': update,
            }
            for channel in self.user_connections[user_id]:
                asyncio.create_task(channel.put(user_message))

    async def notify_message_to_agent(self, agent_id, message):
        control = await self.get_agent_control(agent_id)
        asyncio.create_task(control.received_message(message))