Remote-Control-Center/mc_test/src/renderer/services/websocket-notifications.ts

/**
 * WebSocket Service for real-time notifications
 *
 * Events:
 * - connection_expired: Connection expired
 * - connection_deleted: Connection deleted
 * - connection_will_expire: Connection will expire soon (in 5 min)
 * - connection_extended: Connection extended
 * - jwt_will_expire: JWT will expire soon (in 5 min)
 * - jwt_expired: JWT expired
 */

import { log } from '../utils/logger';
import { API_CONFIG } from '../config/api';

export type WebSocketEventType =
  | 'connection_expired'
  | 'connection_deleted'
  | 'connection_will_expire'
  | 'connection_extended'
  | 'jwt_will_expire'
  | 'jwt_expired'
  | 'connected'
  | 'ping'
  | 'pong';

export interface WebSocketEvent {
  type: WebSocketEventType;
  timestamp: string;
  data?: any;
}

export interface WebSocketMessage {
  type: string;
  token?: string;
  timestamp?: string;
}

export interface ConnectionExpiredEvent {
  connection_id: string;
  hostname: string;
  protocol: string;
  reason: string;
}

export interface ConnectionWillExpireEvent {
  connection_id: string;
  hostname: string;
  protocol: string;
  minutes_remaining: number;
}

export interface ConnectionExtendedEvent {
  connection_id: string;
  hostname: string;
  new_expires_at: string;
  additional_minutes: number;
}

type EventHandler = (event: WebSocketEvent) => void;

const WEBSOCKET_CONFIG = {
  MAX_RECONNECT_ATTEMPTS: 10,
  INITIAL_RECONNECT_DELAY: 1000, // 1 second
  MAX_RECONNECT_DELAY: 30000, // 30 seconds
  PING_INTERVAL: 25000, // 25 seconds
  CONNECTION_TIMEOUT: 5000, // 5 seconds
  EXPONENTIAL_BACKOFF_BASE: 2,
} as const;

const WEBSOCKET_PATH = '/ws/notifications';
const ALL_EVENTS_KEY = '*';

const PROTOCOL_REPLACEMENT = {
  HTTPS: 'wss://',
  HTTP: 'ws://',
} as const;

class WebSocketNotificationService {
  private ws: WebSocket | null = null;
  private reconnectTimer: NodeJS.Timeout | null = null;
  private reconnectAttempts = 0;
  private readonly maxReconnectAttempts = WEBSOCKET_CONFIG.MAX_RECONNECT_ATTEMPTS;
  private readonly reconnectDelay = WEBSOCKET_CONFIG.INITIAL_RECONNECT_DELAY;
  private readonly maxReconnectDelay = WEBSOCKET_CONFIG.MAX_RECONNECT_DELAY;
  private isIntentionallyDisconnected = false;
  private pingTimer: NodeJS.Timeout | null = null;
  private readonly pingInterval = WEBSOCKET_CONFIG.PING_INTERVAL;
  
  private eventHandlers: Map<string, Set<EventHandler>> = new Map();
  
  /**
   * Connect to WebSocket server
   */
  connect(token: string): Promise<void> {
    return new Promise((resolve, reject) => {
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        log.info('websocket', 'Already connected');
        resolve();
        return;
      }
      
      this.isIntentionallyDisconnected = false;

      // Determine WebSocket URL
      const wsUrl = API_CONFIG.BASE_URL
        .replace('https://', PROTOCOL_REPLACEMENT.HTTPS)
        .replace('http://', PROTOCOL_REPLACEMENT.HTTP);
      const fullUrl = `${wsUrl}${WEBSOCKET_PATH}`;
      
      log.info('websocket', 'Connecting to WebSocket', { url: fullUrl });
      
      try {
        this.ws = new WebSocket(fullUrl);
        
        this.ws.onopen = () => {
          log.info('websocket', 'WebSocket connected');

          // Send authentication
          if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            this.ws.send(JSON.stringify({
              type: 'auth',
              token
            }));
            
            log.info('websocket', 'Sent authentication');

            // Wait for connected confirmation
            const connectTimeout = setTimeout(() => {
              log.error('websocket', 'Connection confirmation timeout');
              this.disconnect();
              reject(new Error('Connection confirmation timeout'));
            }, WEBSOCKET_CONFIG.CONNECTION_TIMEOUT);

            // Temporary handler for connected
            const tempOnMessage = (event: MessageEvent) => {
              try {
                const message: WebSocketEvent = JSON.parse(event.data);
                
                if (message.type === 'connected') {
                  clearTimeout(connectTimeout);
                  log.info('websocket', 'Authentication confirmed', message.data);

                  // Reset reconnect attempts
                  this.reconnectAttempts = 0;

                  // Start ping timer
                  this.startPingTimer();

                  // Set main message handler
                  if (this.ws) {
                    this.ws.onmessage = (e) => this.handleMessage(e);
                  }
                  
                  resolve();
                }
              } catch (error) {
                log.error('websocket', 'Failed to parse auth response', { error });
              }
            };
            
            this.ws.onmessage = tempOnMessage;
          }
        };
        
        this.ws.onerror = (error) => {
          log.error('websocket', 'WebSocket error', { error });
          reject(error);
        };
        
        this.ws.onclose = (event) => {
          log.info('websocket', 'WebSocket closed', {
            code: event.code,
            reason: event.reason,
            wasClean: event.wasClean
          });
          
          this.stopPingTimer();

          // Auto-reconnect if not intentionally disconnected
          if (!this.isIntentionallyDisconnected && this.reconnectAttempts < this.maxReconnectAttempts) {
            this.scheduleReconnect(token);
          }
        };
        
      } catch (error) {
        log.error('websocket', 'Failed to create WebSocket', { error });
        reject(error);
      }
    });
  }
  
  /**
   * Disconnect from WebSocket server
   */
  disconnect(): void {
    this.isIntentionallyDisconnected = true;
    
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    
    this.stopPingTimer();
    
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    
    log.info('websocket', 'Disconnected');
  }
  
  /**
   * Schedule reconnection with exponential backoff
   */
  private scheduleReconnect(token: string): void {
    if (this.reconnectTimer) {
      return; // Already scheduled
    }
    
    this.reconnectAttempts++;
    
    // Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s, 30s, ...
    const delay = Math.min(
      this.reconnectDelay * Math.pow(WEBSOCKET_CONFIG.EXPONENTIAL_BACKOFF_BASE, this.reconnectAttempts - 1),
      this.maxReconnectDelay
    );
    
    log.info('websocket', 'Scheduling reconnect', {
      attempt: this.reconnectAttempts,
      delay: `${delay}ms`
    });
    
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect(token).catch((error) => {
        log.error('websocket', 'Reconnect failed', { error });
      });
    }, delay);
  }
  
  /**
   * Start ping timer for keep-alive
   */
  private startPingTimer(): void {
    this.stopPingTimer();
    
    this.pingTimer = setInterval(() => {
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({
          type: 'ping',
          timestamp: new Date().toISOString()
        }));
        
        log.debug('websocket', 'Sent ping');
      }
    }, this.pingInterval);
  }
  
  /**
   * Stop ping timer
   */
  private stopPingTimer(): void{
    if (this.pingTimer) {
      clearInterval(this.pingTimer);
      this.pingTimer = null;
    }
  }
  
  /**
   * Handle incoming message
   */
  private handleMessage(event: MessageEvent): void {
    try {
      const message: WebSocketEvent = JSON.parse(event.data);
      
      if (message.type === 'pong') {
        log.debug('websocket', 'Received pong');
        return;
      }
      
      if (message.type === 'ping') {
        // Respond with pong
        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
          this.ws.send(JSON.stringify({
            type: 'pong',
            timestamp: new Date().toISOString()
          }));
        }
        return;
      }
      
      log.info('websocket', 'Received event', {
        type: message.type,
        timestamp: message.timestamp
      });

      // Call handlers for this event type
      const handlers = this.eventHandlers.get(message.type);
      if (handlers) {
        handlers.forEach(handler => {
          try {
            handler(message);
          } catch (error) {
            log.error('websocket', 'Event handler error', {
              type: message.type,
              error
            });
          }
        });
      }

      // Also call handlers for '*' (all events)
      const allHandlers = this.eventHandlers.get(ALL_EVENTS_KEY);
      if (allHandlers) {
        allHandlers.forEach(handler => {
          try {
            handler(message);
          } catch (error) {
            log.error('websocket', 'Global event handler error', {
              type: message.type,
              error
            });
          }
        });
      }
      
    } catch (error) {
      log.error('websocket', 'Failed to parse message', { error });
    }
  }
  
  /**
   * Subscribe to events
   */
  on(eventType: string, handler: EventHandler): () => void {
    if (!this.eventHandlers.has(eventType)) {
      this.eventHandlers.set(eventType, new Set());
    }
    
    this.eventHandlers.get(eventType)!.add(handler);
    
    log.info('websocket', 'Registered event handler', { eventType });

    // Return unsubscribe function
    return () => {
      const handlers = this.eventHandlers.get(eventType);
      if (handlers) {
        handlers.delete(handler);
        
        if (handlers.size === 0) {
          this.eventHandlers.delete(eventType);
        }
      }
    };
  }
  
  /**
   * Unsubscribe from all events
   */
  removeAllListeners(): void {
    this.eventHandlers.clear();
    log.info('websocket', 'Removed all event handlers');
  }
  
  /**
   * Check connection status
   */
  isConnected(): boolean {
    return this.ws !== null && this.ws.readyState === WebSocket.OPEN;
  }
  
  /**
   * Get number of reconnect attempts
   */
  getReconnectAttempts(): number {
    return this.reconnectAttempts;
  }
}

// Singleton instance
export const websocketNotificationService = new WebSocketNotificationService();