import {useEffect, useState} from 'react';

export function useWebSocket(path, {host, schema} = {}) {
  const [socket, setSocket] = useState();
  useEffect(() => {
    async function connect() {
      const socket = new WebSocket(`${schema || ''}//${host || location.host}${path || '/'}`);
      socket.addEventListener('open', () => {
        setSocket(socket);
      });
      socket.addEventListener('close', ({reason}) => {
        if ('hmr' === reason) {
          connect();
        }
        setSocket(undefined);
      });
    }
    connect();
    return () => {
      setSocket((socket) => {
        if (socket) {
          socket.close();
        }
        return undefined;
      });
    }
  }, [path, host, schema]);
  return socket;
}