useWebSocket.ts

import { useCallback, useEffect, useRef, useState } from 'react'
import type {
  AppEvent,
  MessageType,
  PresencePayload,
  Role,
  Sender,
  WsMessage,
} from './protocol'

// Dev talks straight to the Spring Boot port; production builds assume the
// page's own origin proxies /ws to the backend (see frontend/nginx.conf).
// VITE_WS_URL overrides both.
const WS_URL: string =
  import.meta.env.VITE_WS_URL ??
  (import.meta.env.DEV
    ? 'ws://localhost:8080/ws'
    : `${location.protocol === 'https:' ? 'wss' : 'ws'}://${location.host}/ws`)

export type ConnectionStatus = 'connecting' | 'open' | 'closed'

/**
 * Connects to the backend as the given role and keeps the connection alive.
 *
 * - Sends the JOIN handshake as soon as the socket opens.
 * - Tracks this client's identity (`self`, from WELCOME) and everyone
 *   connected (`roster`, from PRESENCE) — renames flow through the roster.
 * - Reconnects with capped exponential backoff if the connection drops.
 * - Delivers every inbound message to `onMessage` (kept in a ref so callers
 *   can pass a fresh closure each render without re-connecting).
 */
export function useWebSocket(role: Role, onMessage: (msg: WsMessage) => void) {
  const [status, setStatus] = useState<ConnectionStatus>('connecting')
  const [self, setSelf] = useState<Sender | null>(null)
  const [roster, setRoster] = useState<Sender[]>([])
  const socketRef = useRef<WebSocket | null>(null)
  const onMessageRef = useRef(onMessage)
  onMessageRef.current = onMessage

  useEffect(() => {
    let disposed = false
    let attempt = 0
    let reconnectTimer: ReturnType<typeof setTimeout>

    const connect = () => {
      setStatus('connecting')
      const socket = new WebSocket(WS_URL)
      socketRef.current = socket

      socket.onopen = () => {
        attempt = 0
        setStatus('open')
        socket.send(
          JSON.stringify({ type: 'JOIN', sender: null, payload: { role }, timestamp: Date.now() }),
        )
      }

      socket.onmessage = (event) => {
        const msg = JSON.parse(event.data) as WsMessage
        if (msg.type === 'WELCOME') {
          setSelf(msg.payload as Sender)
        } else if (msg.type === 'PRESENCE') {
          const { roster: latest } = msg.payload as PresencePayload
          setRoster(latest)
          // Our own rename comes back through the roster too.
          setSelf((prev) => (prev && latest.find((c) => c.id === prev.id)) ?? prev)
        }
        onMessageRef.current(msg)
      }

      socket.onclose = () => {
        if (disposed) return
        setStatus('closed')
        setRoster([])
        const delay = Math.min(500 * 2 ** attempt, 5000)
        attempt += 1
        reconnectTimer = setTimeout(connect, delay)
      }
    }

    connect()
    return () => {
      disposed = true
      clearTimeout(reconnectTimer)
      socketRef.current?.close()
    }
  }, [role])

  const send = useCallback((type: MessageType, payload: unknown) => {
    const socket = socketRef.current
    if (socket?.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify({ type, sender: null, payload, timestamp: Date.now() }))
    }
  }, [])

  const sendEvent = useCallback((payload: AppEvent) => send('EVENT', payload), [send])
  const sendBroadcast = useCallback((payload: AppEvent) => send('BROADCAST', payload), [send])
  const rename = useCallback((name: string) => send('RENAME', { name }), [send])

  return { status, self, roster, sendEvent, sendBroadcast, rename }
}