backend/src/messages/messages.gateway.ts

import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
  ConnectedSocket,
  MessageBody,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { JwtService } from '@nestjs/jwt';
import { ConfigService } from '@nestjs/config';
import { MessagesService } from './messages.service';
import { SupportChatService } from '../support-chat/support-chat.service';
import { CreateMessageDto } from './dto';
import { Logger } from '@nestjs/common';

interface AuthenticatedSocket extends Socket {
  userId?: string;
  userRole?: string;
}

@WebSocketGateway({
  cors: {
    origin: '*', // Configure in production
    credentials: true,
  },
  transports: ['websocket', 'polling'],
})
export class MessagesGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server: Server;

  private readonly logger = new Logger(MessagesGateway.name);
  private userSocketMap = new Map<string, Set<string>>(); // userId -> Set of socketIds

  constructor(
    private readonly messagesService: MessagesService,
    private readonly supportChatService: SupportChatService,
    private readonly jwtService: JwtService,
    private readonly configService: ConfigService,
  ) {}

  /**
   * Handle new WebSocket connections
   */
  async handleConnection(client: AuthenticatedSocket) {
    this.logger.log(`New connection attempt: ${client.id}`);

    // Extract token from handshake
    const token =
      client.handshake.auth?.token ||
      client.handshake.headers?.authorization?.replace('Bearer ', '');

    if (!token) {
      this.logger.warn(`Connection rejected: No token provided for ${client.id}`);
      client.disconnect(true);
      return;
    }

    // Verify JWT token
    let payload;
    try {
      payload = await this.jwtService.verifyAsync(token, {
        secret: this.configService.get<string>('JWT_SECRET'),
      });
      this.logger.log(`Token verified for ${client.id}, user: ${payload.sub}`);
    } catch (jwtError) {
      this.logger.warn(`Connection rejected: Invalid token for ${client.id} - ${jwtError.message}`);
      client.disconnect(true);
      return;
    }

    // Attach user info to socket
    const userId = payload.sub as string;
    client.userId = userId;
    client.userRole = payload.role;

    // Track socket connection
    if (!this.userSocketMap.has(userId)) {
      this.userSocketMap.set(userId, new Set());
    }
    this.userSocketMap.get(userId)!.add(client.id);

    this.logger.log(`Client connected successfully: ${client.id} (User: ${userId}, Role: ${payload.role})`);

    // Auto-join admin users to admin_room for support chat notifications
    if (payload.role === 'ADMIN' || payload.role === 'SUPER_ADMIN') {
      client.join('admin_room');
      this.logger.log(`Admin ${userId} joined admin_room`);
    }

    // Update user online status in background (don't await)
    this.messagesService.updateOnlineStatus(userId, true)
      .then(() => {
        this.broadcastUserStatus(userId, true);
        this.logger.log(`Online status updated for ${userId}`);
      })
      .catch((err) => {
        this.logger.warn(`Failed to update online status for ${userId}: ${err.message}`);
      });

    // Mark pending SENT messages as DELIVERED now that this user is online
    this.messagesService.markPendingMessagesAsDelivered(userId)
      .then((results) => {
        for (const result of results) {
          // Notify each sender that their messages were delivered
          const deliveredEvent = {
            messageIds: result.messageIds,
            conversationId: result.conversationId,
            deliveredAt: result.deliveredAt,
          };
          this.sendToUser(result.senderId, 'message_delivered', deliveredEvent);
        }
        if (results.length > 0) {
          this.logger.log(`Marked pending messages as delivered for ${userId} across ${results.length} conversation(s)`);
        }
      })
      .catch((err) => {
        this.logger.warn(`Failed to mark pending messages as delivered for ${userId}: ${err.message}`);
      });
  }

  /**
   * Handle WebSocket disconnections
   */
  async handleDisconnect(client: AuthenticatedSocket) {
    if (client.userId) {
      // Remove socket from tracking
      const sockets = this.userSocketMap.get(client.userId);
      if (sockets) {
        sockets.delete(client.id);
        // Only mark offline if no more sockets connected
        if (sockets.size === 0) {
          this.userSocketMap.delete(client.userId);
          await this.messagesService.updateOnlineStatus(client.userId, false);
          this.broadcastUserStatus(client.userId, false);
        }
      }

      this.logger.log(`Client disconnected: ${client.id} (User: ${client.userId})`);
    }
  }

  /**
   * Join a conversation room
   */
  @SubscribeMessage('join_conversation')
  async handleJoinConversation(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string },
  ) {
    if (!client.userId) {
      return { error: 'Unauthorized' };
    }

    try {
      // Verify user is part of the conversation
      await this.messagesService.getConversationById(data.conversationId, client.userId);

      // Join the room
      client.join(`conversation:${data.conversationId}`);

      this.logger.log(`User ${client.userId} joined conversation ${data.conversationId}`);

      return { success: true };
    } catch (error) {
      return { error: error.message };
    }
  }

  /**
   * Leave a conversation room
   */
  @SubscribeMessage('leave_conversation')
  handleLeaveConversation(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string },
  ) {
    client.leave(`conversation:${data.conversationId}`);
    this.logger.log(`User ${client.userId} left conversation ${data.conversationId}`);
    return { success: true };
  }

  /**
   * Send a message via WebSocket
   */
  @SubscribeMessage('send_message')
  async handleSendMessage(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string; message: CreateMessageDto },
  ) {
    if (!client.userId) {
      return { error: 'Unauthorized' };
    }

    try {
      const message = await this.messagesService.createMessage(
        data.conversationId,
        client.userId,
        data.message,
      );

      // Send directly to the other participant's socket(s)
      // This is the primary delivery mechanism — works even if they haven't joined the room
      try {
        const participants = await this.messagesService.getConversationParticipants(data.conversationId);
        if (participants) {
          const otherUserId = client.userId === participants.userId
            ? participants.agentUserId
            : participants.userId;
          this.sendToUser(otherUserId, 'new_message', message);

          // Auto-mark as DELIVERED if recipient is online
          if (this.isUserOnline(otherUserId)) {
            await this.messagesService.markMessageDelivered(message.id);
            const deliveredEvent = {
              messageId: message.id,
              conversationId: data.conversationId,
              deliveredAt: new Date(),
            };
            // Notify sender that message was delivered
            this.sendToUser(client.userId, 'message_delivered', deliveredEvent);
          }
        }
      } catch (participantError) {
        this.logger.warn(`Failed to get participants for direct delivery in ${data.conversationId}: ${participantError.message}`);
      }

      // Also broadcast to the conversation room as a fallback
      // (client.to excludes sender; recipients dedup by message ID)
      client
        .to(`conversation:${data.conversationId}`)
        .emit('new_message', message);

      return { success: true, message };
    } catch (error) {
      return { error: error.message };
    }
  }

  /**
   * Handle typing indicator start
   */
  @SubscribeMessage('typing_start')
  handleTypingStart(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string },
  ) {
    if (!client.userId) return;

    // Broadcast to others in the room
    client.to(`conversation:${data.conversationId}`).emit('typing_start', {
      userId: client.userId,
      conversationId: data.conversationId,
    });
  }

  /**
   * Handle typing indicator stop
   */
  @SubscribeMessage('typing_stop')
  handleTypingStop(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string },
  ) {
    if (!client.userId) return;

    client.to(`conversation:${data.conversationId}`).emit('typing_stop', {
      userId: client.userId,
      conversationId: data.conversationId,
    });
  }

  /**
   * Mark messages as read
   */
  @SubscribeMessage('mark_read')
  async handleMarkRead(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { conversationId: string },
  ) {
    if (!client.userId) {
      return { error: 'Unauthorized' };
    }

    try {
      await this.messagesService.markMessagesAsRead(data.conversationId, client.userId);

      const readEvent = {
        conversationId: data.conversationId,
        readBy: client.userId,
        readAt: new Date(),
      };

      // Notify via room (for clients viewing this conversation)
      client.to(`conversation:${data.conversationId}`).emit('messages_read', readEvent);

      // Also notify the other participant directly
      const participants = await this.messagesService.getConversationParticipants(data.conversationId);
      if (participants) {
        const otherUserId = client.userId === participants.userId
          ? participants.agentUserId
          : participants.userId;
        this.sendToUser(otherUserId, 'messages_read', readEvent);
      }

      return { success: true };
    } catch (error) {
      return { error: error.message };
    }
  }

  // ==========================================
  // SUPPORT CHAT EVENTS
  // ==========================================

  /**
   * Join a support chat room
   */
  @SubscribeMessage('support_join')
  async handleSupportJoin(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { chatId: string },
  ) {
    if (!client.userId) {
      return { error: 'Unauthorized' };
    }

    try {
      const isAdmin = client.userRole === 'ADMIN';
      await this.supportChatService.getChat(data.chatId, client.userId, isAdmin);
      client.join(`support:${data.chatId}`);
      this.logger.log(`User ${client.userId} joined support chat ${data.chatId}`);
      return { success: true };
    } catch (error) {
      return { error: error.message };
    }
  }

  /**
   * Leave a support chat room
   */
  @SubscribeMessage('support_leave')
  handleSupportLeave(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { chatId: string },
  ) {
    client.leave(`support:${data.chatId}`);
    this.logger.log(`User ${client.userId} left support chat ${data.chatId}`);
    return { success: true };
  }

  /**
   * Send a message in a support chat via WebSocket
   */
  @SubscribeMessage('support_send_message')
  async handleSupportSendMessage(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { chatId: string; content: string },
  ) {
    if (!client.userId) {
      return { error: 'Unauthorized' };
    }

    try {
      const senderRole = client.userRole === 'ADMIN' ? 'ADMIN' : 'USER';
      const message = await this.supportChatService.sendMessage(
        data.chatId,
        client.userId,
        senderRole,
        data.content,
      );

      // Broadcast to all clients in the support chat room
      this.server.to(`support:${data.chatId}`).emit('support_new_message', message);

      // Notify admins about new unread count (for sidebar badge)
      if (senderRole !== 'ADMIN') {
        const unreadTotal = await this.supportChatService.getAdminUnreadTotal();
        this.server.to('admin_room').emit('support_unread_update', { unreadTotal });
      }

      return { success: true, message };
    } catch (error) {
      return { error: error.message };
    }
  }

  /**
   * Support chat typing indicators
   */
  @SubscribeMessage('support_typing_start')
  handleSupportTypingStart(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { chatId: string },
  ) {
    if (!client.userId) return;
    client.to(`support:${data.chatId}`).emit('support_typing_start', {
      userId: client.userId,
      chatId: data.chatId,
    });
  }

  @SubscribeMessage('support_typing_stop')
  handleSupportTypingStop(
    @ConnectedSocket() client: AuthenticatedSocket,
    @MessageBody() data: { chatId: string },
  ) {
    if (!client.userId) return;
    client.to(`support:${data.chatId}`).emit('support_typing_stop', {
      userId: client.userId,
      chatId: data.chatId,
    });
  }

  /**
   * Broadcast user online/offline status
   */
  private broadcastUserStatus(userId: string, isOnline: boolean) {
    this.server.emit('user_status_change', {
      userId,
      isOnline,
      lastSeenAt: isOnline ? null : new Date(),
    });
  }

  /**
   * Send message to a specific user (by userId)
   */
  sendToUser(userId: string, event: string, data: any) {
    const sockets = this.userSocketMap.get(userId);
    if (sockets && sockets.size > 0) {
      sockets.forEach((socketId) => {
        this.server.to(socketId).emit(event, data);
      });
      this.logger.debug(`Sent ${event} to user ${userId} via ${sockets.size} socket(s)`);
    } else {
      this.logger.debug(`User ${userId} has no active sockets for ${event} - message saved in DB`);
    }
  }

  /**
   * Check if a user is online
   */
  isUserOnline(userId: string): boolean {
    const sockets = this.userSocketMap.get(userId);
    return sockets ? sockets.size > 0 : false;
  }
}