import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { ConfigService } from '@nestjs/config';
import { WsException } from '@nestjs/websockets';

@Injectable()
export class WsJwtGuard implements CanActivate {
  constructor(
    private readonly jwtService: JwtService,
    private readonly configService: ConfigService,
  ) {}

  async canActivate(context: ExecutionContext): Promise<boolean> {
    try {
      const client = context.switchToWs().getClient();
      const token = this.extractToken(client);

      if (!token) {
        throw new WsException('未提供认证令牌');
      }

      const payload = this.jwtService.verify(token, {
        secret: this.configService.get('jwt.secret'),
      });

      // 将用户信息附加到客户端
      client.user = payload;
      
      return true;
    } catch (error) {
      throw new WsException('无效的认证令牌');
    }
  }

  private extractToken(client: any): string | null {
    // 从handshake auth中获取token
    const auth = client.handshake?.auth?.token;
    if (auth) return auth.replace('Bearer ', '');
    
    // 从query中获取token
    const token = client.handshake?.query?.token;
    if (token) return token as string;
    
    return null;
  }
}
