Atomizer/atomizer-dashboard/frontend/src/hooks/useWebSocket.ts

import { useState, useEffect } from 'react';
import useWebSocket, { ReadyState } from 'react-use-websocket';
import { WebSocketMessage } from '../types';

interface UseOptimizationWebSocketProps {
  studyId: string | null;
  onMessage?: (message: WebSocketMessage) => void;
}

export const useOptimizationWebSocket = ({ studyId, onMessage }: UseOptimizationWebSocketProps) => {
  const [socketUrl, setSocketUrl] = useState<string | null>(null);

  useEffect(() => {
    if (studyId) {
      // In development, we might need to point to localhost:8000 explicitly if not proxied
      // But assuming Vite proxy is set up correctly:
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      const host = window.location.host; // This will be localhost:3000 in dev
      // If using proxy in vite.config.ts, this works. 
      // If not, we might need to hardcode backend URL for dev:
      const backendHost = import.meta.env.DEV ? 'localhost:8000' : host;
      
      setSocketUrl(`${protocol}//${backendHost}/api/ws/optimization/${studyId}`);
    } else {
      setSocketUrl(null);
    }
  }, [studyId]);

  const { sendMessage, lastMessage, readyState } = useWebSocket(socketUrl, {
    shouldReconnect: (_closeEvent) => true,
    reconnectAttempts: 10,
    reconnectInterval: 3000,
    onOpen: () => console.log('WebSocket Connected'),
    onClose: () => console.log('WebSocket Disconnected'),
    onError: (e) => console.error('WebSocket Error:', e),
  });

  useEffect(() => {
    if (lastMessage !== null) {
      try {
        const data = JSON.parse(lastMessage.data) as WebSocketMessage;
        if (onMessage) {
          onMessage(data);
        }
      } catch (e) {
        console.error('Failed to parse WebSocket message:', e);
      }
    }
  }, [lastMessage, onMessage]);

  const connectionStatus = {
    [ReadyState.CONNECTING]: 'Connecting',
    [ReadyState.OPEN]: 'Open',
    [ReadyState.CLOSING]: 'Closing',
    [ReadyState.CLOSED]: 'Closed',
    [ReadyState.UNINSTANTIATED]: 'Uninstantiated',
  }[readyState];

  return {
    sendMessage,
    lastMessage,
    readyState,
    connectionStatus,
  };
};