import asyncio
from datetime import datetime

import database as db 
import temporary as tmp
from chat import Chat
from stream import Stream

def message_json(msg):
    return {
        "id": msg.id,
        "time": msg.time,
        "sender": msg.sender,
        "kind": msg.kind,
        "subkind": msg.subkind,
        "content": msg.content,
        "threadId": msg.thread_id,
    }

async def aggregate_thread(thread_id):
    print(f"aggregate {thread_id}")
    query = db.Message.filter(thread_id=thread_id)

    # Could do all these in one call?
    message_count = await query.count()
    print(f"    count {message_count}")
    if message_count == 0: return

    first_message = await query.first()
    last_message = await query.last()
    print(f"    bounds {first_message} {last_message}")

    return await db.Thread.get(id=thread_id).update(
        first_message_id=first_message.id,
        last_message_id=last_message.id,
        message_count=message_count
    )

class ChatMixin:
    async def create_chat(self, user_id, name, description, picture, members=[], temporary=False):
        if temporary:
            chat = tmp.Chat.create(
                self,
                user_id=user_id,
                name=name,
                description=description,
                picture=picture
            )
        else:
            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

    # Gets a chat view from a given chat id
    async def get_chat(self, chat_id):
        chat = Chat(self, chat_id)
        await chat.load()
        return chat

    async def get_chats(self, user_id):
        # Direct filtering on the model class
        chats = await db.Chat.filter(user_id=user_id).all()

        tmp_chats = self.tmp.get_user_chats(user_id)
        if tmp_chats:
            chats += [self.tmp.chats[chat_id] for chat_id in tmp_chats]

        return [{
            "id": chat.id,
            "name": chat.name,
            "description": chat.description,
            "picture": chat.picture
        } for chat in chats]

    async def delete_chat(self, chat_id):
        if tmp.is_tmp_id(chat_id):
            chat = self.tmp.chats[chat_id]
            # TODO: Delete

        else:
            # 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):
        if tmp.is_tmp_id(chat_id):
            chat = self.tmp.chats[chat_id]
            
            return [message.to_json() for message in chat.messages]

        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()

        # TODO: Stop repetition.
        # There's db.Message and chat.Message, also should be combined somehow, or wrap db always.
        return [message_json(msg) for msg in messages]
    
    async def create_chat_thread(self, chat_id, parent_id=None, name=None):
        if tmp.is_tmp_id(chat_id):
            thread = tmp.Thread.create(self, chat_id, name)
        else:
            thread = await db.Thread.create(
                chat_id=chat_id,
                parent_id=parent_id,
                name=name
            )
        print(f"Created thread #{thread.id}: {thread.name}")
        return thread.id

    async def get_chat_threads(self, chat_id):
        threads = await db.Thread.filter(chat_id=chat_id).all()
        return [{
            'id': it.id,
            'parentId': it.parent_id,
            'name': it.name,
            'description': it.description,
            'firstMessageId': it.first_message_id,
            'lastMessageId': it.last_message_id,
            'messageCount': it.message_count,
        } for it in threads]

    async def get_thread_messages(self, chat_id, thread_id, limit_message_id=None):
        chat = await self.get_chat(chat_id)
        thread = chat.get_thread(thread_id)
        messages = await thread.load_page(limit_message_id)
        return [message_json(msg) for msg in messages]

    async def get_chat_thread_heads(self, chat_id, limit_message_id=None):
        chat = await self.get_chat(chat_id)
        messages = await chat.load_thread_heads(limit_message_id)
        return [message_json(msg) for msg in messages]

    async def fork_chat_thread(self, thread_id, message_id=None, name=None):
        """Move all messages from message_id onwards into a new child thread.
        If none is given, this is basically create_chat_thread with an implicit chat id."""
        thread = await db.Thread.get(id=thread_id)
        new_id = await self.create_chat_thread(thread.chat_id, parent_id=thread_id, name=name or thread.name)

        if message_id is not None:
            # Ensure all forks have the correct information
            forks = await db.Thread.filter(parent_id=thread_id).all()
            for fork in forks:
                await aggregate_thread(fork.id)

            await db.Thread.filter(parent_id=thread_id, first_message_id__gte=message_id).update(parent_id=new_id)
            await db.Message.filter(thread_id=thread_id, id__gte=message_id).update(thread_id=new_id)

        # Notify users. what a heavy operation!

        chat = await db.Chat.get(id=thread.chat_id)
        messages = await db.Message.filter(thread_id=new_id).all()
        print("Chat: " + repr(chat))
        print("Fork message list: " + repr(messages))
        for msg in messages:
            self.notify_message_to_user(chat.user_id, msg, True)

        return new_id

    async def add_agent_to_chat(self, chat_id, agent_id):
        if tmp.is_tmp_id(chat_id):
            chat = self.tmp.chats[chat_id]
            chat.participants.append(agent_id)
        else:
            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_format = format
            format = format and "steam-" + format or "stream"
            initial_content = str(content.id)
        else:
            initial_content = content

        # Set to true for streams
        updating = False
        if tmp.is_tmp_id(chat_id):
            message = Message(
                tmp.generate_tmp_id(self),
                #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
            )
            message.chat_id = chat_id
            chat = self.tmp.chats[chat_id]
            chat.messages.append(message)
            chat_owner_id = chat.user_id
        else:
            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():
            # Not all messages are visible to agents
            if message.kind == "resource":
                return

            # Fetch participant agents
            if tmp.is_tmp_id(chat_id):
                agent_ids = self.tmp.chats[chat_id].participants
            else:
                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):
                pass
            async def end_fn ():
                message.format = initial_format
                message.content = content.content

                if db.is_instance(message):
                    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):
        if tmp.is_tmp_id(message_id):
            raise Exception("Cannot update temporary messages")
        else:
            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:
            if db.is_instance(message):
                await message.save() # Saves the specific instance
            print(f"Updated message #{message.id} fields: {summary}")

    def notify_message_to_user(self, user_id, message, update=False):
        if user_id in self.user_connections:
            # TODO: Server's responsibility, and Message.to_json responsibility as well
            user_message = {
                'action': 'chat-message',
                'id': message.id,
                'chatId': message.chat_id,
                'time': message.time,
                'sender': message.sender,
                'kind': message.kind,
                'subkind': message.subkind,
                'format': message.format,
                '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))