ironnotify-js/src/websocket.ts

import type { Notification, ConnectionState, NotifyEventHandlers } from './types';

/**
 * WebSocket manager for real-time notifications.
 */
export class WebSocketManager {
  private ws: WebSocket | null = null;
  private readonly wsUrl: string;
  private readonly apiKey: string;
  private readonly debug: boolean;
  private readonly autoReconnect: boolean;
  private readonly maxReconnectAttempts: number;
  private readonly reconnectDelay: number;

  private reconnectAttempts = 0;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private handlers: NotifyEventHandlers = {};
  private subscriptions: Set<string> = new Set();
  private _state: ConnectionState = 'disconnected';

  constructor(
    wsUrl: string,
    apiKey: string,
    debug: boolean,
    autoReconnect: boolean,
    maxReconnectAttempts: number,
    reconnectDelay: number
  ) {
    this.wsUrl = wsUrl;
    this.apiKey = apiKey;
    this.debug = debug;
    this.autoReconnect = autoReconnect;
    this.maxReconnectAttempts = maxReconnectAttempts;
    this.reconnectDelay = reconnectDelay;
  }

  /**
   * Get connection state.
   */
  get state(): ConnectionState {
    return this._state;
  }

  /**
   * Set event handlers.
   */
  setHandlers(handlers: NotifyEventHandlers): void {
    this.handlers = { ...this.handlers, ...handlers };
  }

  /**
   * Connect to the WebSocket server.
   */
  connect(): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      return;
    }

    if (typeof WebSocket === 'undefined') {
      if (this.debug) {
        console.warn('[IronNotify] WebSocket not available in this environment');
      }
      return;
    }

    this.setState('connecting');

    try {
      const url = `${this.wsUrl}?token=${encodeURIComponent(this.apiKey)}`;
      this.ws = new WebSocket(url);

      this.ws.onopen = () => {
        this.reconnectAttempts = 0;
        this.setState('connected');

        if (this.debug) {
          console.log('[IronNotify] WebSocket connected');
        }

        // Re-subscribe to any previous subscriptions
        this.subscriptions.forEach((sub) => {
          this.sendMessage({ type: 'subscribe', channel: sub });
        });
      };

      this.ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          this.handleMessage(data);
        } catch (err) {
          if (this.debug) {
            console.error('[IronNotify] Failed to parse message:', err);
          }
        }
      };

      this.ws.onerror = (event) => {
        if (this.debug) {
          console.error('[IronNotify] WebSocket error:', event);
        }
        this.handlers.onError?.(new Error('WebSocket error'));
      };

      this.ws.onclose = () => {
        this.setState('disconnected');

        if (this.debug) {
          console.log('[IronNotify] WebSocket disconnected');
        }

        if (this.autoReconnect && this.reconnectAttempts < this.maxReconnectAttempts) {
          this.scheduleReconnect();
        }
      };
    } catch (err) {
      if (this.debug) {
        console.error('[IronNotify] Failed to connect:', err);
      }
      this.handlers.onError?.(err instanceof Error ? err : new Error(String(err)));
    }
  }

  /**
   * Disconnect from the WebSocket server.
   */
  disconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }

    this.reconnectAttempts = this.maxReconnectAttempts; // Prevent reconnection

    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }

    this.setState('disconnected');
  }

  /**
   * Subscribe to a user's notifications.
   */
  subscribeToUser(userId: string): void {
    const channel = `user:${userId}`;
    this.subscriptions.add(channel);

    if (this.ws?.readyState === WebSocket.OPEN) {
      this.sendMessage({ type: 'subscribe', channel });
    }
  }

  /**
   * Subscribe to app-wide notifications.
   */
  subscribeToApp(): void {
    const channel = 'app';
    this.subscriptions.add(channel);

    if (this.ws?.readyState === WebSocket.OPEN) {
      this.sendMessage({ type: 'subscribe', channel });
    }
  }

  /**
   * Unsubscribe from a channel.
   */
  unsubscribe(channel: string): void {
    this.subscriptions.delete(channel);

    if (this.ws?.readyState === WebSocket.OPEN) {
      this.sendMessage({ type: 'unsubscribe', channel });
    }
  }

  private setState(state: ConnectionState): void {
    this._state = state;
    this.handlers.onConnectionStateChange?.(state);
  }

  private scheduleReconnect(): void {
    this.setState('reconnecting');
    this.reconnectAttempts++;

    const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);

    if (this.debug) {
      console.log(`[IronNotify] Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})`);
    }

    this.reconnectTimer = setTimeout(() => {
      this.connect();
    }, delay);
  }

  private sendMessage(message: object): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
    }
  }

  private handleMessage(data: any): void {
    switch (data.type) {
      case 'notification':
        this.handlers.onNotification?.(data.notification as Notification);
        break;

      case 'unread_count':
        this.handlers.onUnreadCountChange?.(data.count);
        break;

      case 'error':
        this.handlers.onError?.(new Error(data.message));
        break;

      default:
        if (this.debug) {
          console.log('[IronNotify] Unknown message type:', data.type);
        }
    }
  }
}