import { useEffect, useRef, useCallback } from 'react'

type EventHandler = (event: string, data: unknown) => void

export function useWebSocket(onEvent: EventHandler) {
  const wsRef = useRef<WebSocket | null>(null)
  const handleRef = useRef(onEvent)
  handleRef.current = onEvent

  const connect = useCallback(() => {
    const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:'
    const ws = new WebSocket(`${protocol}//${location.host}/ws`)

    ws.onmessage = (e) => {
      try {
        const { event, data } = JSON.parse(e.data)
        handleRef.current(event, data)
      } catch {
        // ignore malformed messages
      }
    }

    ws.onclose = () => {
      setTimeout(connect, 3000)
    }

    wsRef.current = ws
  }, [])

  useEffect(() => {
    connect()
    return () => wsRef.current?.close()
  }, [connect])
}
