static/js/wsClient.js

/**
 * WebSocket Client for Real-time Communication
 * Manages WebSocket connections with automatic reconnection and exponential backoff
 * Supports message routing to type-specific subscribers
 */
class WSClient {
    constructor() {
        this.socket = null;
        this.status = 'disconnected';
        this.statusSubscribers = new Set();
        this.globalSubscribers = new Set();
        this.typeSubscribers = new Map();
        this.eventLog = [];
        this.backoff = 1000; // Initial backoff delay in ms
        this.maxBackoff = 16000; // Maximum backoff delay in ms
        this.shouldReconnect = true;
        this.reconnectAttempts = 0;
        this.connectionStartTime = null;
    }

    /**
     * Automatically determine WebSocket URL based on current window location
     * Always uses the current origin to avoid hardcoded URLs
     */
    get url() {
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const host = window.location.host;
        return `${protocol}//${host}/ws`;
    }

    /**
     * Log WebSocket events for debugging and monitoring
     * Maintains a rolling buffer of the last 100 events
     * @param {Object} event - Event object to log
     */
    logEvent(event) {
        const entry = { 
            ...event, 
            time: new Date().toISOString(),
            attempt: this.reconnectAttempts 
        };
        this.eventLog.push(entry);
        // Keep only last 100 events
        if (this.eventLog.length > 100) {
            this.eventLog = this.eventLog.slice(-100);
        }
        console.log('[WSClient]', entry);
    }

    /**
     * Subscribe to connection status changes
     * @param {Function} callback - Called with new status ('connecting', 'connected', 'disconnected', 'error')
     * @returns {Function} Unsubscribe function
     */
    onStatusChange(callback) {
        if (typeof callback !== 'function') {
            throw new Error('Callback must be a function');
        }
        this.statusSubscribers.add(callback);
        // Immediately call with current status
        callback(this.status);
        return () => this.statusSubscribers.delete(callback);
    }

    /**
     * Subscribe to all WebSocket messages
     * @param {Function} callback - Called with parsed message data
     * @returns {Function} Unsubscribe function
     */
    onMessage(callback) {
        if (typeof callback !== 'function') {
            throw new Error('Callback must be a function');
        }
        this.globalSubscribers.add(callback);
        return () => this.globalSubscribers.delete(callback);
    }

    /**
     * Subscribe to specific message types
     * @param {string} type - Message type to subscribe to (e.g., 'market_update', 'news_update')
     * @param {Function} callback - Called with messages of the specified type
     * @returns {Function} Unsubscribe function
     */
    subscribe(type, callback) {
        if (typeof callback !== 'function') {
            throw new Error('Callback must be a function');
        }
        if (!this.typeSubscribers.has(type)) {
            this.typeSubscribers.set(type, new Set());
        }
        const set = this.typeSubscribers.get(type);
        set.add(callback);
        return () => set.delete(callback);
    }

    /**
     * Update connection status and notify all subscribers
     * @param {string} newStatus - New status value
     */
    updateStatus(newStatus) {
        if (this.status !== newStatus) {
            const oldStatus = this.status;
            this.status = newStatus;
            this.logEvent({ 
                type: 'status_change', 
                from: oldStatus, 
                to: newStatus 
            });
            this.statusSubscribers.forEach(cb => {
                try {
                    cb(newStatus);
                } catch (error) {
                    console.error('[WSClient] Error in status subscriber:', error);
                }
            });
        }
    }

    /**
     * Establish WebSocket connection with automatic reconnection
     * Implements exponential backoff for reconnection attempts
     */
    connect() {
        // Prevent multiple simultaneous connection attempts
        if (this.socket && (this.socket.readyState === WebSocket.CONNECTING || this.socket.readyState === WebSocket.OPEN)) {
            console.log('[WSClient] Already connected or connecting');
            return;
        }

        this.connectionStartTime = Date.now();
        this.updateStatus('connecting');
        
        try {
            this.socket = new WebSocket(this.url);
            this.logEvent({ 
                type: 'connection_attempt', 
                url: this.url,
                attempt: this.reconnectAttempts + 1
            });

            this.socket.onopen = () => {
                const connectionTime = Date.now() - this.connectionStartTime;
                this.backoff = 1000; // Reset backoff on successful connection
                this.reconnectAttempts = 0;
                this.updateStatus('connected');
                this.logEvent({ 
                    type: 'connection_established', 
                    connectionTime: `${connectionTime}ms`
                });
                console.log(`[WSClient] Connected to ${this.url} in ${connectionTime}ms`);
            };

            this.socket.onmessage = (event) => {
                try {
                    const data = JSON.parse(event.data);
                    this.logEvent({ 
                        type: 'message_received', 
                        messageType: data.type || 'unknown',
                        size: event.data.length 
                    });
                    
                    // Notify global subscribers
                    this.globalSubscribers.forEach(cb => {
                        try {
                            cb(data);
                        } catch (error) {
                            console.error('[WSClient] Error in global subscriber:', error);
                        }
                    });
                    
                    // Notify type-specific subscribers
                    if (data.type && this.typeSubscribers.has(data.type)) {
                        this.typeSubscribers.get(data.type).forEach(cb => {
                            try {
                                cb(data);
                            } catch (error) {
                                console.error(`[WSClient] Error in ${data.type} subscriber:`, error);
                            }
                        });
                    }
                } catch (error) {
                    console.error('[WSClient] Message parse error:', error);
                    this.logEvent({ 
                        type: 'parse_error', 
                        error: error.message,
                        rawData: event.data.substring(0, 100) 
                    });
                }
            };

            this.socket.onclose = (event) => {
                const wasConnected = this.status === 'connected';
                this.updateStatus('disconnected');
                this.logEvent({ 
                    type: 'connection_closed', 
                    code: event.code,
                    reason: event.reason || 'No reason provided',
                    wasClean: event.wasClean
                });
                
                // Attempt reconnection if enabled
                if (this.shouldReconnect) {
                    this.reconnectAttempts++;
                    const delay = this.backoff;
                    this.backoff = Math.min(this.backoff * 2, this.maxBackoff);
                    
                    console.log(`[WSClient] Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})...`);
                    this.logEvent({ 
                        type: 'reconnect_scheduled', 
                        delay: `${delay}ms`,
                        nextBackoff: `${this.backoff}ms`
                    });
                    
                    setTimeout(() => this.connect(), delay);
                }
            };

            this.socket.onerror = (error) => {
                console.error('[WSClient] WebSocket error:', error);
                this.updateStatus('error');
                this.logEvent({ 
                    type: 'connection_error', 
                    error: error.message || 'Unknown error',
                    readyState: this.socket ? this.socket.readyState : 'null'
                });
            };
        } catch (error) {
            console.error('[WSClient] Failed to create WebSocket:', error);
            this.updateStatus('error');
            this.logEvent({ 
                type: 'creation_error', 
                error: error.message 
            });
            
            // Retry connection if enabled
            if (this.shouldReconnect) {
                this.reconnectAttempts++;
                const delay = this.backoff;
                this.backoff = Math.min(this.backoff * 2, this.maxBackoff);
                setTimeout(() => this.connect(), delay);
            }
        }
    }

    /**
     * Gracefully disconnect WebSocket and disable automatic reconnection
     */
    disconnect() {
        this.shouldReconnect = false;
        if (this.socket) {
            this.logEvent({ type: 'manual_disconnect' });
            this.socket.close(1000, 'Client disconnect');
            this.socket = null;
        }
    }

    /**
     * Manually trigger reconnection (useful for testing or recovery)
     */
    reconnect() {
        this.disconnect();
        this.shouldReconnect = true;
        this.backoff = 1000; // Reset backoff
        this.reconnectAttempts = 0;
        this.connect();
    }

    /**
     * Send a message through the WebSocket connection
     * @param {Object} data - Data to send (will be JSON stringified)
     * @returns {boolean} True if sent successfully, false otherwise
     */
    send(data) {
        if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
            console.error('[WSClient] Cannot send message: not connected');
            this.logEvent({ 
                type: 'send_failed', 
                reason: 'not_connected',
                readyState: this.socket ? this.socket.readyState : 'null'
            });
            return false;
        }

        try {
            const message = JSON.stringify(data);
            this.socket.send(message);
            this.logEvent({ 
                type: 'message_sent', 
                messageType: data.type || 'unknown',
                size: message.length 
            });
            return true;
        } catch (error) {
            console.error('[WSClient] Failed to send message:', error);
            this.logEvent({ 
                type: 'send_error', 
                error: error.message 
            });
            return false;
        }
    }

    /**
     * Get a copy of the event log
     * @returns {Array} Array of logged events
     */
    getEvents() {
        return [...this.eventLog];
    }

    /**
     * Get current connection statistics
     * @returns {Object} Connection statistics
     */
    getStats() {
        return {
            status: this.status,
            reconnectAttempts: this.reconnectAttempts,
            currentBackoff: this.backoff,
            maxBackoff: this.maxBackoff,
            shouldReconnect: this.shouldReconnect,
            subscriberCounts: {
                status: this.statusSubscribers.size,
                global: this.globalSubscribers.size,
                typed: Array.from(this.typeSubscribers.entries()).map(([type, subs]) => ({
                    type,
                    count: subs.size
                }))
            },
            eventLogSize: this.eventLog.length,
            url: this.url
        };
    }

    /**
     * Check if WebSocket is currently connected
     * @returns {boolean} True if connected
     */
    isConnected() {
        return this.socket && this.socket.readyState === WebSocket.OPEN;
    }

    /**
     * Clear all subscribers (useful for cleanup)
     */
    clearSubscribers() {
        this.statusSubscribers.clear();
        this.globalSubscribers.clear();
        this.typeSubscribers.clear();
        this.logEvent({ type: 'subscribers_cleared' });
    }
}

// Create singleton instance
const wsClient = new WSClient();

// Auto-connect on module load
wsClient.connect();

// Export singleton instance
export default wsClient;