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 { 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 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})`);

    // 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}`);
      });
  }

  /**
   * 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,
      );

      // Broadcast to all clients in the conversation room
      this.server
        .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);

      // Notify the other party that messages have been read
      client.to(`conversation:${data.conversationId}`).emit('messages_read', {
        conversationId: data.conversationId,
        readBy: client.userId,
        readAt: new Date(),
      });

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

  /**
   * 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.forEach((socketId) => {
        this.server.to(socketId).emit(event, data);
      });
    }
  }

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