import { useState, useEffect, useRef, useCallback } from 'react';
import { WebSocketState, MessageType } from '../constants';

export const useWebSocket = (url) => {
  const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
  const wsRef = useRef(null);
  const handlersRef = useRef([]);

  const connect = useCallback(() => {
    setWsState(WebSocketState.CONNECTING);
    const ws = new WebSocket(url);
    wsRef.current = ws;

    ws.onopen = () => {
      setWsState(WebSocketState.CONNECTED);
    };

    ws.onclose = () => {
      setWsState(WebSocketState.DISCONNECTED);
      setTimeout(connect, 2000);
    };

    ws.onerror = (error) => {
      console.error('WebSocket error:', error);
    };

    ws.onmessage = (event) => {
      try {
        const message = JSON.parse(event.data);
        handlersRef.current.forEach(handler => handler(message));
      } catch (error) {
        console.error('Error handling WebSocket message:', error);
      }
    };

    return ws;
  }, [url]);

  useEffect(() => {
    const ws = connect();
    return () => {
      if (ws) ws.close();
    };
  }, [connect]);

  const sendMessage = useCallback((data) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(data));
    } else {
      console.error('WebSocket not connected, state:', wsRef.current?.readyState);
    }
  }, []);

  const addMessageHandler = useCallback((handler) => {
    handlersRef.current.push(handler);
  }, []);

  return {
    wsState,
    sendMessage,
    addMessageHandler
  };
};