# 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 django.db.models import Count, F, Q 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() self.read_channels = await Channel.get_accessible_channels(user, 'read') self.write_channels = await Channel.get_accessible_channels(user, 'write') async for channel in self.read_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 async for channel in self.read_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 'edit_message': await self.edit_message(**content) case 'delete_message': await self.delete_message(**content) case 'fetch_messages': await self.fetch_messages(**content) case 'mark_read': await self.mark_read(**content) case 'start_private_chat': await self.start_private_chat(**content) case unknown: print("Unknown message type:", unknown) async def fetch_channels(self) -> None: user = self.scope['user'] message = { 'type': 'fetch_channels', 'channels': [ { 'id': channel.id, 'name': channel.get_visible_name(user), 'category': channel.category, 'read_access': True, 'write_access': await self.write_channels.acontains(channel), 'unread_messages': channel.unread_messages, } async for channel in self.read_channels.prefetch_related('invited') .annotate(total_messages=Count('messages', distinct=True)) .annotate(read_messages=Count('messages', filter=Q(messages__users_read=user), distinct=True)) .annotate(unread_messages=F('total_messages') - F('read_messages')).all() ] } await self.send_json(message) async def receive_message(self, channel_id: int, content: str, **kwargs) -> None: user = self.scope['user'] channel = await Channel.objects.prefetch_related('tournament__pools__juries', 'pool', 'team', 'invited') \ .aget(id=channel_id) if not await self.write_channels.acontains(channel): return message = await Message.objects.acreate( author=user, channel=channel, content=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 edit_message(self, message_id: int, content: str, **kwargs) -> None: message = await Message.objects.aget(id=message_id) user = self.scope['user'] if user.id != message.author_id and not user.is_superuser: return message.content = content await message.asave() await self.channel_layer.group_send(f'chat-{message.channel_id}', { 'type': 'chat.edit_message', 'id': message_id, 'channel_id': message.channel_id, 'content': content, }) async def delete_message(self, message_id: int, **kwargs) -> None: message = await Message.objects.aget(id=message_id) user = self.scope['user'] if user.id != message.author_id and not user.is_superuser: return await message.adelete() await self.channel_layer.group_send(f'chat-{message.channel_id}', { 'type': 'chat.delete_message', 'id': message_id, 'channel_id': message.channel_id, }) async def fetch_messages(self, channel_id: int, offset: int = 0, limit: int = 50, **_kwargs) -> None: channel = await Channel.objects.aget(id=channel_id) if not await self.read_channels.acontains(channel): return limit = min(limit, 200) # Fetch only maximum 200 messages at the time messages = Message.objects \ .filter(channel=channel) \ .annotate(read=Count('users_read', filter=Q(users_read=self.scope['user']))) \ .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, 'read': message.read > 0, } async for message in messages ])) }) async def mark_read(self, message_ids: list[int], **_kwargs) -> None: messages = Message.objects.filter(id__in=message_ids) async for message in messages.all(): await message.users_read.aadd(self.scope['user']) unread_messages_by_channel = Message.objects.exclude(users_read=self.scope['user']).values('channel_id') \ .annotate(unread_messages=Count('channel_id')) await self.send_json({ 'type': 'mark_read', 'messages': [{'id': message.id, 'channel_id': message.channel_id} async for message in messages.all()], 'unread_messages': {group['channel_id']: group['unread_messages'] async for group in unread_messages_by_channel.all()}, }) async def start_private_chat(self, user_id: int, **kwargs) -> 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 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_edit_message(self, message) -> None: await self.send_json({'type': 'edit_message', 'id': message['id'], 'channel_id': message['channel_id'], 'content': message['content']}) async def chat_delete_message(self, message) -> None: await self.send_json({'type': 'delete_message', 'id': message['id'], 'channel_id': message['channel_id']}) 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']})