whale-town-end/src/business/notice/notice.gateway.ts

import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  MessageBody,
  ConnectedSocket,
  OnGatewayConnection,
  OnGatewayDisconnect,
} from '@nestjs/websockets';
import { Server } from 'ws';
import * as WebSocket from 'ws';
import { Logger } from '@nestjs/common';

interface AuthenticatedSocket extends WebSocket {
  userId?: number;
}

@WebSocketGateway({
  cors: {
    origin: '*',
  },
  path: '/ws/notice',
})
export class NoticeGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server: Server;

  private readonly logger = new Logger(NoticeGateway.name);
  private readonly userSockets = new Map<number, Set<AuthenticatedSocket>>();

  handleConnection(client: AuthenticatedSocket) {
    this.logger.log(`Client connected: ${client.readyState}`);
  }

  handleDisconnect(client: AuthenticatedSocket) {
    this.logger.log(`Client disconnected`);
    
    if (client.userId) {
      const userSockets = this.userSockets.get(client.userId);
      if (userSockets) {
        userSockets.delete(client);
        if (userSockets.size === 0) {
          this.userSockets.delete(client.userId);
        }
      }
    }
  }

  @SubscribeMessage('authenticate')
  handleAuthenticate(
    @MessageBody() data: { userId: number },
    @ConnectedSocket() client: AuthenticatedSocket,
  ) {
    const { userId } = data;
    
    if (!userId) {
      client.send(JSON.stringify({ error: 'User ID is required' }));
      return;
    }

    client.userId = userId;
    
    if (!this.userSockets.has(userId)) {
      this.userSockets.set(userId, new Set());
    }
    this.userSockets.get(userId)!.add(client);

    client.send(JSON.stringify({ 
      type: 'authenticated', 
      data: { userId } 
    }));

    this.logger.log(`User ${userId} authenticated`);
  }

  @SubscribeMessage('ping')
  handlePing(@ConnectedSocket() client: AuthenticatedSocket) {
    client.send(JSON.stringify({ type: 'pong' }));
  }

  // 发送消息给特定用户
  sendToUser(userId: number, message: any) {
    const userSockets = this.userSockets.get(userId);
    if (userSockets) {
      const messageStr = JSON.stringify(message);
      userSockets.forEach(socket => {
        if (socket.readyState === WebSocket.OPEN) {
          socket.send(messageStr);
        }
      });
      this.logger.log(`Message sent to user ${userId}`);
    } else {
      this.logger.warn(`User ${userId} not connected`);
    }
  }

  // 广播消息给所有连接的用户
  broadcast(message: any) {
    const messageStr = JSON.stringify(message);
    this.server.clients.forEach(client => {
      if (client.readyState === WebSocket.OPEN) {
        client.send(messageStr);
      }
    });
    this.logger.log('Message broadcasted to all clients');
  }

  // 获取在线用户数量
  getOnlineUsersCount(): number {
    return this.userSockets.size;
  }

  // 获取在线用户列表
  getOnlineUsers(): number[] {
    return Array.from(this.userSockets.keys());
  }
}