frontend/src/lib/lcmLive.ts

import { get, writable } from 'svelte/store';

export enum LCMLiveStatus {
  CONNECTED = 'connected',
  DISCONNECTED = 'disconnected',
  CONNECTING = 'connecting',
  WAIT = 'wait',
  SEND_FRAME = 'send_frame',
  TIMEOUT = 'timeout',
  ERROR = 'error'
}

const initStatus: LCMLiveStatus = LCMLiveStatus.DISCONNECTED;

export const lcmLiveStatus = writable<LCMLiveStatus>(initStatus);
export const streamId = writable<string | null>(null);

// WebSocket connection
let websocket: WebSocket | null = null;
// Flag to track intentional connection closure
let intentionalClosure = false;

// Register browser unload event listener to properly close WebSockets
if (typeof window !== 'undefined') {
  window.addEventListener('beforeunload', () => {
    // Mark any closure during page unload as intentional
    intentionalClosure = true;
    // Close the WebSocket properly if it exists
    if (websocket && websocket.readyState === WebSocket.OPEN) {
      websocket.close(1000, 'Page unload');
    }
  });
}
export const lcmLiveActions = {
  async start(getSreamdata: () => any[]) {
    return new Promise((resolve, reject) => {
      try {
        // Set connecting status immediately
        lcmLiveStatus.set(LCMLiveStatus.CONNECTING);

        const userId = crypto.randomUUID();
        const websocketURL = `${
          window.location.protocol === 'https:' ? 'wss' : 'ws'
        }:${window.location.host}/api/ws/${userId}`;

        // Close any existing connection first
        if (websocket && websocket.readyState !== WebSocket.CLOSED) {
          websocket.close();
        }

        websocket = new WebSocket(websocketURL);

        // Set a connection timeout
        const connectionTimeout = setTimeout(() => {
          if (websocket && websocket.readyState !== WebSocket.OPEN) {
            console.error('WebSocket connection timeout');
            lcmLiveStatus.set(LCMLiveStatus.ERROR);
            streamId.set(null);
            reject(new Error('Connection timeout. Please try again.'));
            websocket.close();
          }
        }, 10000); // 10 second timeout

        websocket.onopen = () => {
          clearTimeout(connectionTimeout);
          console.log('Connected to websocket');
        };

        websocket.onclose = (event) => {
          clearTimeout(connectionTimeout);
          console.log(`Disconnected from websocket: ${event.code} ${event.reason}`);

          // Only change status if we're not in ERROR state (which would mean we already handled the error)
          if (get(lcmLiveStatus) !== LCMLiveStatus.ERROR) {
            lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
          }

          // If connection was never established (close without open)
          if (event.code === 1006 && get(streamId) === null) {
            reject(new Error('Cannot connect to server. Please try again later.'));
          }
        };

        websocket.onerror = (err) => {
          clearTimeout(connectionTimeout);
          console.error('WebSocket error:', err);
          lcmLiveStatus.set(LCMLiveStatus.ERROR);
          streamId.set(null);
          reject(new Error('Connection error. Please try again.'));
        };

        websocket.onmessage = (event) => {
          try {
            const data = JSON.parse(event.data);
            switch (data.status) {
              case 'connected':
                lcmLiveStatus.set(LCMLiveStatus.CONNECTED);
                streamId.set(userId);
                resolve({ status: 'connected', userId });
                break;
              case 'send_frame':
                lcmLiveStatus.set(LCMLiveStatus.SEND_FRAME);
                try {
                  const streamData = getSreamdata();
                  // Send as an object, not a string, to use the proper handling in the send method
                  this.send({ status: 'next_frame' });
                  for (const d of streamData) {
                    this.send(d);
                  }
                } catch (error) {
                  console.error('Error sending frame data:', error);
                }
                break;
              case 'wait':
                lcmLiveStatus.set(LCMLiveStatus.WAIT);
                break;
              case 'timeout':
                console.log('Session timeout');
                lcmLiveStatus.set(LCMLiveStatus.TIMEOUT);
                streamId.set(null);
                reject(new Error('Session timeout. Please restart.'));
                break;
              case 'error':
                console.error('Server error:', data.message);
                lcmLiveStatus.set(LCMLiveStatus.ERROR);
                streamId.set(null);
                reject(new Error(data.message || 'Server error occurred'));
                break;
              default:
                console.log('Unknown message status:', data.status);
            }
          } catch (error) {
            console.error('Error handling websocket message:', error);
          }
        };
      } catch (err) {
        console.error('Error initializing websocket:', err);
        lcmLiveStatus.set(LCMLiveStatus.ERROR);
        streamId.set(null);
        reject(err);
      }
    });
  },
  send(data: Blob | { [key: string]: any }) {
    try {
      if (websocket && websocket.readyState === WebSocket.OPEN) {
        if (data instanceof Blob) {
          websocket.send(data);
        } else {
          websocket.send(JSON.stringify(data));
        }
      } else {
        const readyStateText = websocket
          ? ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED'][websocket.readyState]
          : 'null';
        console.warn(`WebSocket not ready for sending: ${readyStateText}`);

        // If WebSocket is closed unexpectedly, set status to disconnected
        if (!websocket || websocket.readyState === WebSocket.CLOSED) {
          lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
          streamId.set(null);
        }
      }
    } catch (error) {
      console.error('Error sending data through WebSocket:', error);
      // Handle WebSocket error by forcing disconnection
      this.stop();
    }
  },

  async reconnect(getSreamdata: () => any[]) {
    try {
      await this.stop();
      // Small delay to ensure clean disconnection before reconnecting
      await new Promise((resolve) => setTimeout(resolve, 500));
      return await this.start(getSreamdata);
    } catch (error) {
      console.error('Reconnection failed:', error);
      throw error;
    }
  },

  async stop() {
    lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
    try {
      if (websocket) {
        // Only attempt to close if not already closed
        if (websocket.readyState !== WebSocket.CLOSED) {
          // Set up onclose handler to clean up only
          websocket.onclose = () => {
            console.log('WebSocket closed cleanly during stop()');
          };

          // Set up onerror to be silent during intentional closure
          websocket.onerror = () => {};

          websocket.close(1000, 'Client initiated disconnect');
        }
      }
    } catch (error) {
      console.error('Error during WebSocket closure:', error);
    } finally {
      // Always clean up references
      websocket = null;
      streamId.set(null);
    }
  }
};