plateforme-tfjm2/chat/consumers.py

# Copyright (C) 2024 by Animath
# SPDX-License-Identifier: GPL-3.0-or-later

from channels.generic.websocket import AsyncJsonWebsocketConsumer
from django.contrib.auth.models import User
from registration.models import Registration

from .models import Channel, Message


class ChatConsumer(AsyncJsonWebsocketConsumer):
    """
    This consumer manages the websocket of the chat interface.
    """
    async def connect(self) -> None:
        """
        This function is called when a new websocket is trying to connect to the server.
        We accept only if this is a user of a team of the associated tournament, or a volunteer
        of the tournament.
        """
        if '_fake_user_id' in self.scope['session']:
            self.scope['user'] = await User.objects.aget(pk=self.scope['session']['_fake_user_id'])

        # Fetch the registration of the current user
        user = self.scope['user']
        if user.is_anonymous:
            # User is not authenticated
            await self.close()
            return

        reg = await Registration.objects.aget(user_id=user.id)
        self.registration = reg

        # Accept the connection
        await self.accept()

        channels = await Channel.get_accessible_channels(user, 'read')
        async for channel in channels.all():
            await self.channel_layer.group_add(f"chat-{channel.id}", self.channel_name)
        await self.channel_layer.group_add(f"user-{user.id}", self.channel_name)

    async def disconnect(self, close_code) -> None:
        """
        Called when the websocket got disconnected, for any reason.
        :param close_code: The error code.
        """
        if self.scope['user'].is_anonymous:
            # User is not authenticated
            return

        channels = await Channel.get_accessible_channels(self.scope['user'], 'read')
        async for channel in channels.all():
            await self.channel_layer.group_discard(f"chat-{channel.id}", self.channel_name)
        await self.channel_layer.group_discard(f"user-{self.scope['user'].id}", self.channel_name)

    async def receive_json(self, content, **kwargs):
        """
        Called when the client sends us some data, parsed as JSON.
        :param content: The sent data, decoded from JSON text. Must content a `type` field.
        """
        match content['type']:
            case 'fetch_channels':
                await self.fetch_channels()
            case 'send_message':
                await self.receive_message(content)
            case 'fetch_messages':
                await self.fetch_messages(**content)
            case 'start_private_chat':
                await self.start_private_chat(content['user_id'])
            case unknown:
                print("Unknown message type:", unknown)

    async def fetch_channels(self) -> None:
        user = self.scope['user']

        read_channels = await Channel.get_accessible_channels(user, 'read')
        write_channels = await Channel.get_accessible_channels(user, 'write')
        message = {
            'type': 'fetch_channels',
            'channels': [
                {
                    'id': channel.id,
                    'name': channel.get_visible_name(user),
                    'category': channel.category,
                    'read_access': True,
                    'write_access': await write_channels.acontains(channel),
                }
                async for channel in read_channels.prefetch_related('invited').all()
            ]
        }
        await self.send_json(message)

    async def receive_message(self, message: dict) -> None:
        user = self.scope['user']
        channel = await Channel.objects.prefetch_related('tournament__pools__juries', 'pool', 'team', 'invited') \
            .aget(id=message['channel_id'])
        write_channels = await Channel.get_accessible_channels(user, 'write')
        if not await write_channels.acontains(channel):
            return

        message = await Message.objects.acreate(
            author=user,
            channel=channel,
            content=message['content'],
        )

        await self.channel_layer.group_send(f'chat-{channel.id}', {
            'type': 'chat.send_message',
            'id': message.id,
            'channel_id': channel.id,
            'timestamp': message.created_at.isoformat(),
            'author_id': message.author_id,
            'author': await message.aget_author_name(),
            'content': message.content,
        })

    async def fetch_messages(self, channel_id: int, offset: int = 0, limit: int = 50, **_kwargs) -> None:
        channel = await Channel.objects.aget(id=channel_id)
        read_channels = await Channel.get_accessible_channels(self.scope['user'], 'read')
        if not await read_channels.acontains(channel):
            return

        limit = min(limit, 200)  # Fetch only maximum 200 messages at the time

        messages = Message.objects.filter(channel=channel).order_by('-created_at')[offset:offset + limit].all()
        await self.send_json({
            'type': 'fetch_messages',
            'channel_id': channel_id,
            'messages': list(reversed([
                {
                    'id': message.id,
                    'timestamp': message.created_at.isoformat(),
                    'author_id': message.author_id,
                    'author': await message.aget_author_name(),
                    'content': message.content,
                }
                async for message in messages
            ]))
        })

    async def start_private_chat(self, user_id: int) -> None:
        user = self.scope['user']
        other_user = await User.objects.aget(id=user_id)
        channel_qs = Channel.objects.filter(private=True).filter(invited=user).filter(invited=other_user)
        if not await channel_qs.aexists():
            channel = await Channel.objects.acreate(
                name=f"{user.first_name} {user.last_name}, {other_user.first_name} {other_user.last_name}",
                category=Channel.ChannelCategory.PRIVATE,
                private=True,
            )
            await channel.invited.aset([user, other_user])
            await channel.asave()

            await self.channel_layer.group_add(f"chat-{channel.id}", self.channel_name)

            if user != other_user:
                await self.channel_layer.group_send(f"user-{other_user.id}", {
                    'type': 'chat.start_private_chat',
                    'channel': {
                        'id': channel.id,
                        'name': f"{user.first_name} {user.last_name}",
                        'category': channel.category,
                        'read_access': True,
                        'write_access': True,
                    }
                })
        else:
            channel = await channel_qs.afirst()

        await self.channel_layer.group_send(f"user-{user.id}", {
            'type': 'chat.start_private_chat',
            'channel': {
                'id': channel.id,
                'name': f"{other_user.first_name} {other_user.last_name}",
                'category': channel.category,
                'read_access': True,
                'write_access': True,
            }
        })

    async def chat_send_message(self, message) -> None:
        await self.send_json({'type': 'send_message', 'id': message['id'], 'channel_id': message['channel_id'],
                              'timestamp': message['timestamp'], 'author': message['author'],
                              'content': message['content']})

    async def chat_start_private_chat(self, message) -> None:
        await self.channel_layer.group_add(f"chat-{message['channel']['id']}", self.channel_name)
        await self.send_json({'type': 'start_private_chat', 'channel': message['channel']})