Tolan/frontend/src/utils/websocket.ts

import ReconnectingWebSocket from 'reconnecting-websocket';
import { v4 as uuidv4 } from 'uuid';
import type { Message } from '@/types';
import { performanceLogger } from './logger';

/**
 * WebSocket管理器
 */
export class WebSocketManager {
  private ws: ReconnectingWebSocket | null = null;
  private messageHandlers: Map<string, (message: Message) => void> = new Map();
  private pingInterval: ReturnType<typeof setInterval> | null = null;
  private readonly PING_INTERVAL = 30000; // 30秒

  connect(url: string, onConnectionChange?: (connected: boolean) => void): void {
    this.ws = new ReconnectingWebSocket(url, [], {
      maxRetries: 10,
      connectionTimeout: 5000,
      maxReconnectionDelay: 10000,
      minReconnectionDelay: 1000,
      reconnectionDelayGrowFactor: 1.3,
    });

    this.ws.onopen = () => {
      console.log('[WebSocket] 连接已建立');
      // 记录WebSocket连接建立时间
      performanceLogger.recordWebSocketConnected();
      this.startHeartbeat();
      if (onConnectionChange) {
        onConnectionChange(true);
      }
    };

    this.ws.onmessage = (event) => {
      try {
        const message: Message = JSON.parse(event.data);
        // 简化日志：只记录关键消息类型
        if (message.type === 'error') {
          console.error(`[WebSocket] [${message.requestId?.substring(0, 8) || ''}] 错误:`, message);
        }
        this.handleMessage(message);
      } catch (error) {
        console.error('[WebSocket] 解析消息失败:', error);
      }
    };

    this.ws.onerror = (error) => {
      console.error('[WebSocket] 连接错误:', error);
    };

    this.ws.onclose = (event) => {
      console.log(`[WebSocket] 连接已关闭 (code: ${event.code})`);
      this.stopHeartbeat();
      if (onConnectionChange) {
        onConnectionChange(false);
      }
    };
  }

  private handleMessage(message: Message): void {
    this.messageHandlers.forEach((handler) => {
      handler(message);
    });
  }

  onMessage(id: string, handler: (message: Message) => void): void {
    this.messageHandlers.set(id, handler);
  }

  offMessage(id: string): void {
    this.messageHandlers.delete(id);
  }

  send(message: Partial<Message>): void {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      console.warn(`[WebSocket] [${message.requestId?.substring(0, 8) || ''}] 发送失败: WebSocket未连接`, message.type);
      return;
    }

    const fullMessage = {
      ...message,
      timestamp: Date.now(),
    };

    // 简化日志：只记录错误和取消操作
    if (message.type === 'cancel' || message.type === 'error') {
      console.log(`[WebSocket] [${message.requestId?.substring(0, 8) || ''}] 发送: ${message.type}`);
    }
    
    this.ws.send(JSON.stringify(fullMessage));
  }

  sendAudioChunk(requestId: string, data: string, chunkIndex: number): void {
    this.send({
      type: 'audio_chunk',
      requestId,
      data,
      sampleRate: 16000,
      channels: 1,
      bitDepth: 16,
      chunkIndex,
    });
  }

  sendAudioEnd(requestId: string, totalChunks: number, speechDuration?: number): void {
    this.send({
      type: 'audio_end',
      requestId,
      totalChunks,
      speechDuration,
    });
  }

  sendCancel(requestId: string): void {
    this.send({
      type: 'cancel',
      requestId,
    });
  }

  private startHeartbeat(): void {
    this.pingInterval = setInterval(() => {
      this.send({
        type: 'ping',
        requestId: uuidv4(),
      });
    }, this.PING_INTERVAL);
  }

  private stopHeartbeat(): void {
    if (this.pingInterval) {
      clearInterval(this.pingInterval);
      this.pingInterval = null;
    }
  }

  disconnect(): void {
    this.stopHeartbeat();
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this.messageHandlers.clear();
  }
}