mono/packages/ui/src/contexts/WS_Socket.tsx

import React, { createContext, useContext, useEffect, useState, useCallback, ReactNode } from 'react';
import modbusService from '@/services/modbusService';
import logger from '@/Logger';

export type WsStatus = 'DISCONNECTED' | 'CONNECTING' | 'CONNECTED' | 'ERROR' | 'RECONNECTING';

interface WebSocketContextType {
  isConnected: boolean;
  isConnecting: boolean;
  wsStatus: WsStatus;
  apiUrl: string;
  setApiUrl: (url: string) => void;
  connectToServer: (url?: string) => Promise<boolean>;
  disconnectFromServer: () => void;
  abortConnectionAttempt: () => void;
}

const WebSocketContext = createContext<WebSocketContextType | undefined>(undefined);

export const useWebSocket = () => {
  const context = useContext(WebSocketContext);
  if (context === undefined) {
    throw new Error('useWebSocket must be used within a WebSocketProvider');
  }
  return context;
};

interface WebSocketProviderProps {
  children: ReactNode;
  url?: string;
}

export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ children, url: initialUrl = '' }) => {
  const [apiUrl, setApiUrl] = useState<string>(initialUrl);
  const [isConnected, setIsConnected] = useState<boolean>(false);
  const [isConnecting, setConnecting] = useState<boolean>(false);
  const [wsStatus, setWsStatus] = useState<WsStatus>('DISCONNECTED');
  const [connectionAborted, setConnectionAborted] = useState<boolean>(false);

  // Use a ref or simply the variable to track initialization to prevent double connect in strict mode if needed,
  // but simpler logic is often better.

  const handleWsStatusChange = useCallback((status: WsStatus) => {
    setWsStatus(status);
    setIsConnected(status === 'CONNECTED');

    if (status === 'CONNECTED') {
      setConnecting(false);
    } else {
      if (status === 'ERROR') {
        setConnecting(false);
      } else if (status === 'RECONNECTING') {
        setConnecting(true);
      } else if (status === 'DISCONNECTED') {
        setConnecting(false);
      }
    }
  }, []);

  const disconnectWebSocket = useCallback((intentional: boolean = true) => {
    modbusService.disconnect(intentional);
  }, []);

  const connectWebSocket = useCallback(async (urlToConnect: string): Promise<boolean> => {
    if (!urlToConnect) {
      logger.error('API URL not set');
      return false;
    }
    let wsUrl: string;
    try {
      const url = new URL(urlToConnect);
      // If the providing URL is http/https, switch to ws/wss. If it's already ws/wss, use as is.
      // The previous logic assumed http input. Let's be robust.
      if (url.protocol === 'http:') {
        url.protocol = 'ws:';
      } else if (url.protocol === 'https:') {
        url.protocol = 'wss:';
      }

      // Ensure /ws path
      if (!url.pathname.endsWith('/ws')) {
        url.pathname = url.pathname.replace(/\/+$/, '') + '/ws';
      }

      wsUrl = url.toString();
    } catch (error) {
      logger.logError(error, 'Invalid API URL');
      return false;
    }

    try {
      const newConnectionEstablished = await modbusService.connect(
        wsUrl,
        handleWsStatusChange,
      );
      return newConnectionEstablished;
    } catch (error) {
      logger.logError(error, '[ModbusContext] modbusService.connect() failed');
      return false;
    }
  }, [handleWsStatusChange]);

  const abortConnectionAttempt = useCallback(() => {
    setConnectionAborted(true);
    setConnecting(false);
    disconnectWebSocket(true);
  }, [disconnectWebSocket]);

  const connectToServer = useCallback(async (urlToUse?: string): Promise<boolean> => {
    const targetUrl = urlToUse || apiUrl;

    if (urlToUse && urlToUse !== apiUrl) {
      setApiUrl(targetUrl);
    }

    setConnectionAborted(false);
    setConnecting(true);
    try {
      const didConnect = await connectWebSocket(targetUrl);

      if (didConnect) {

      } else {
        if (modbusService.getConnectionStatus() === 'CONNECTED') {

        } else {

        }
      }

      const isNowConnected = modbusService.getConnectionStatus() === 'CONNECTED';
      setConnecting(false);
      return isNowConnected;

    } catch (error: any) {
      logger.logError(error, '[ModbusContext] Error in connectToServer');
      setIsConnected(false);
      setConnecting(false);
      return false;
    }
  }, [apiUrl, connectWebSocket]);

  const disconnectFromServer = useCallback((): void => {
    setConnectionAborted(true);
    setIsConnected(false);
    disconnectWebSocket(true);
  }, [disconnectWebSocket]);

  // Initial connection effect
  useEffect(() => {
    if (initialUrl && !isConnected && !isConnecting && !connectionAborted) {
      connectToServer(initialUrl);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [initialUrl]);
  // We strictly want this only on mount or if initialUrl changes significantly and we want to reconnect.
  // Including dependencies like 'connectToServer' might cause loops if not careful.

  return (
    <WebSocketContext.Provider
      value={{
        isConnected,
        isConnecting,
        wsStatus,
        apiUrl,
        setApiUrl,
        connectToServer,
        disconnectFromServer,
        abortConnectionAttempt,
      }}
    >
      {children}
    </WebSocketContext.Provider>
  );
};