""" apps/business/network_chat/consumers.py WebSocket consumer for real-time chat between distributors. Handles message broadcasting, user authentication, and connection lifecycle. """ from asgiref.sync import sync_to_async from channels.generic.websocket import AsyncJsonWebsocketConsumer from django.db.models import Q from .models import ChatConversation, ChatMessage from apps.platform.settings.models import PlatformFeatureFlag class ChatConsumer(AsyncJsonWebsocketConsumer): """ WebSocket consumer for real-time chat. Message format (incoming): { "type": "chat_message", "conversation_id": "uuid", "recipient_distributor_code": "code", "body": "message text" } Broadcast format (outgoing): { "type": "chat_message", "id": "message_id", "conversation_id": "conversation_id", "sender_distributor_code": "code", "body": "text", "created_at": "timestamp", "is_read": false } """ async def connect(self): """Accept WebSocket connection if user is authenticated and add to group.""" if not PlatformFeatureFlag.is_on(PlatformFeatureFlag.CHAT_ENABLED): await self.close(code=4003) return if self.scope['user'] and self.scope['user'].is_authenticated: # Join all distributor-code groups for this user so any conversation code can receive pushes. distributor_codes = await self.get_user_distributor_codes(self.scope['user']) if not distributor_codes: await self.close(code=4004) return self.distributor_codes = distributor_codes for distributor_code in distributor_codes: await self.channel_layer.group_add( f'chat_{distributor_code}', self.channel_name ) await self.accept() else: await self.close() async def disconnect(self, close_code): """Handle disconnection and remove from group.""" if hasattr(self, 'distributor_codes'): for distributor_code in self.distributor_codes: await self.channel_layer.group_discard( f'chat_{distributor_code}', self.channel_name ) async def receive_json(self, content, **kwargs): """ Receive and process incoming messages. """ if not PlatformFeatureFlag.is_on(PlatformFeatureFlag.CHAT_ENABLED): await self.send_json({ 'type': 'error', 'message': 'Chat is currently disabled by admin.', }) return msg_type = content.get('type') if msg_type == 'chat_message': await self.handle_chat_message(content) elif msg_type == 'mark_read': await self.handle_mark_read(content) async def handle_chat_message(self, content): """ Handle incoming chat message: 1. Retrieve or create conversation 2. Create and save message 3. Broadcast to recipient 4. Broadcast unread count update to recipient """ try: user = self.scope['user'] conversation_id = content.get('conversation_id') recipient_code = content.get('recipient_distributor_code') body = content.get('body', '').strip() if not body: await self.send_json({ 'type': 'error', 'message': 'Message body cannot be empty' }) return if len(body) > 2000: await self.send_json({ 'type': 'error', 'message': 'Message body too long (max 2000 characters)' }) return # Create or get conversation conversation = await self.create_or_get_conversation( user, recipient_code, conversation_id ) if not conversation: await self.send_json({ 'type': 'error', 'message': 'Could not create conversation' }) return # Resolve message routing/sender identity from the actual conversation participants. sender_code, recipient_push_codes = await self.get_sender_and_recipient_codes(conversation, user) if not sender_code: await self.send_json({ 'type': 'error', 'message': 'Could not resolve sender distributor code' }) return if not recipient_push_codes: await self.send_json({ 'type': 'error', 'message': 'Could not resolve recipient distributor code(s)' }) return # Create message message = await self.create_message(conversation, user, sender_code, body) payload = { 'type': 'chat_message', 'id': str(message.id), 'conversation_id': str(conversation.id), 'sender_distributor_code': sender_code, 'body': message.body, 'created_at': message.created_at.isoformat(), 'is_read': message.is_read, } # Broadcast to all recipient distributor-code groups so active socket(s) receive instantly. for code in recipient_push_codes: await self.channel_layer.group_send( f'chat_{code}', payload ) # Also broadcast to sender's group so they see the message (with is_mine=true on frontend) await self.channel_layer.group_send( f'chat_{sender_code}', payload ) # Broadcast unread count update to recipient unread_payload_by_code = await self.get_unread_payload_by_codes(recipient_push_codes) for code, unread_count in unread_payload_by_code.items(): await self.channel_layer.group_send( f'chat_{code}', { 'type': 'unread_count_update', 'unread_count': unread_count, } ) # Send confirmation to sender (optional - message already in broadcast) await self.send_json({ 'type': 'message_sent', 'id': str(message.id), 'conversation_id': str(conversation.id), 'created_at': message.created_at.isoformat(), }) except Exception as e: await self.send_json({ 'type': 'error', 'message': f'Failed to send message: {str(e)}' }) async def handle_mark_read(self, content): """Handle mark conversation as read.""" try: user = self.scope['user'] conversation_id = content.get('conversation_id') if not conversation_id: return # Mark all messages in conversation as read await self.mark_conversation_read(conversation_id, user) except Exception as e: await self.send_json({ 'type': 'error', 'message': f'Failed to mark as read: {str(e)}' }) # Methods for handling WebSocket group messages async def chat_message(self, event): """ Receive a chat message from group and send to WebSocket. """ await self.send_json({ 'type': 'chat_message', 'id': event['id'], 'conversation_id': event['conversation_id'], 'sender_distributor_code': event['sender_distributor_code'], 'body': event['body'], 'created_at': event['created_at'], 'is_read': event['is_read'], }) async def unread_count_update(self, event): """ Receive unread count update and send to WebSocket. """ await self.send_json({ 'type': 'unread_count_update', 'unread_count': event['unread_count'], }) async def user_online(self, event): """Notify when a user comes online.""" await self.send_json({ 'type': 'user_online', 'distributor_code': event['distributor_code'], }) async def user_offline(self, event): """Notify when a user goes offline.""" await self.send_json({ 'type': 'user_offline', 'distributor_code': event['distributor_code'], }) # Async database helper methods @sync_to_async def get_user_distributor_codes(self, user): """Get all distributor codes for a user.""" try: from apps.business.distributor.models import DistributorID return list( DistributorID.objects .filter(user=user) .order_by('global_position') .values_list('distributor_code', flat=True) ) except Exception: return [] @sync_to_async def create_or_get_conversation(self, user, recipient_code, conversation_id=None): """Create or retrieve a conversation.""" try: from apps.business.distributor.models import DistributorID sender_code = ( DistributorID.objects .filter(user=user) .order_by('global_position') .values_list('distributor_code', flat=True) .first() ) if not sender_code: return None if conversation_id: conversation = ChatConversation.objects.get(id=conversation_id) if conversation.initiator_id != user.id and conversation.recipient_id != user.id: return None return conversation # Try to get existing conversation in either direction conversation = ChatConversation.objects.filter( Q(initiator_distributor_code=sender_code, recipient_distributor_code=recipient_code) | Q(initiator_distributor_code=recipient_code, recipient_distributor_code=sender_code) ).first() if conversation: return conversation # Get recipient user recipient_dist = DistributorID.objects.filter( distributor_code=recipient_code ).first() if not recipient_dist: return None # Create new conversation (user as initiator) conversation = ChatConversation.objects.create( initiator=user, recipient=recipient_dist.user, initiator_distributor_code=sender_code, recipient_distributor_code=recipient_code, ) return conversation except Exception as e: print(f"Error creating/getting conversation: {e}") return None @sync_to_async def get_sender_and_recipient_codes(self, conversation, user): """Resolve sender code and all recipient codes based on conversation participants.""" try: from apps.business.distributor.models import DistributorID if conversation.initiator_id == user.id: sender_code = conversation.initiator_distributor_code recipient_user_id = conversation.recipient_id elif conversation.recipient_id == user.id: sender_code = conversation.recipient_distributor_code recipient_user_id = conversation.initiator_id else: return None, [] recipient_codes = list( DistributorID.objects .filter(user_id=recipient_user_id) .values_list('distributor_code', flat=True) ) # Ensure conversation-bound code is always included for deterministic routing. if conversation.initiator_id == user.id: conv_recipient_code = conversation.recipient_distributor_code else: conv_recipient_code = conversation.initiator_distributor_code if conv_recipient_code and conv_recipient_code not in recipient_codes: recipient_codes.append(conv_recipient_code) return sender_code, recipient_codes except Exception as e: print(f"Error resolving sender/recipient codes: {e}") return None, [] @sync_to_async def get_unread_payload_by_codes(self, distributor_codes): """Return unread counts keyed by distributor code for the same recipient user(s).""" payload = {} try: for code in distributor_codes: payload[code] = self.get_unread_count_sync(code) except Exception as e: print(f"Error getting unread payloads: {e}") return payload def get_unread_count_sync(self, distributor_code): """Sync helper for unread count by distributor code.""" try: from apps.business.distributor.models import DistributorID dist_id = DistributorID.objects.filter( distributor_code=distributor_code ).first() if not dist_id: return 0 user = dist_id.user unread = ChatMessage.objects.filter( is_read=False, ).filter( Q(conversation__initiator=user) | Q(conversation__recipient=user) ).exclude( sender=user ).count() return unread except Exception as e: print(f"Error getting unread count: {e}") return 0 @sync_to_async def create_message(self, conversation, user, sender_code, body): """Create a chat message.""" return ChatMessage.objects.create( conversation=conversation, sender=user, sender_distributor_code=sender_code, body=body, ) @sync_to_async def mark_conversation_read(self, conversation_id, user): """Mark all messages in a conversation as read for the user.""" try: ChatMessage.objects.filter( conversation_id=conversation_id, is_read=False ).exclude(sender=user).update(is_read=True) except Exception as e: print(f"Error marking conversation read: {e}") @sync_to_async def get_unread_count(self, distributor_code): """Get total unread message count for a distributor.""" return self.get_unread_count_sync(distributor_code)