Wizard/wizard_website/src/api/useWizard.ts

import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import type {
  Action,
  ClientMessage,
  GameOptions,
  GameView,
  Lobby,
  ServerMessage,
  ServerStatus,
  TrickResult,
} from './types'

// 'busy': the server refused the connection because it is full.
export type ConnectionStatus = 'connecting' | 'open' | 'closed' | 'replaced' | 'busy'

export interface Toast {
  id: number
  message: string
}

// Per tab, so several tabs can act as different players; survives reloads.
const TOKEN_KEY = 'wizard.token'
const TRICK_DISPLAY_MS = 2200
const TOAST_MS = 4000
const MAX_RETRY_DELAY_MS = 10000
// Close code the server uses when the session was opened in another tab.
const CLOSE_REPLACED = 4000
// "Try Again Later": the server is full. Retry slowly so we don't add load.
const CLOSE_TRY_AGAIN_LATER = 1013
const BUSY_RETRY_DELAYS_MS = [15000, 30000, 60000]

function websocketUrl(): string {
  const configured = import.meta.env.VITE_WS_URL as string | undefined
  if (configured) return configured
  // The API lives next to the website: /wizard/website/ -> /wizard/api/ws
  // (and /api/ws in development).
  const url = new URL('../api/ws', window.location.href)
  url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'
  url.search = ''
  return url.toString()
}

function readToken(): string | null {
  try {
    return sessionStorage.getItem(TOKEN_KEY)
  } catch {
    return null
  }
}

function storeToken(token: string) {
  try {
    sessionStorage.setItem(TOKEN_KEY, token)
  } catch {
    // Storage unavailable: reconnecting after a reload will start a new session.
  }
}

export function useWizard() {
  const [status, setStatus] = useState<ConnectionStatus>('connecting')
  const [playerId, setPlayerId] = useState<string | null>(null)
  const [lobby, setLobby] = useState<Lobby | null>(null)
  const [game, setGame] = useState<GameView | null>(null)
  // The most recently completed trick, shown for a moment before it is cleared.
  const [completedTrick, setCompletedTrick] = useState<TrickResult | null>(null)
  const [toasts, setToasts] = useState<Toast[]>([])
  const [serverStatus, setServerStatus] = useState<ServerStatus>({ load: 'normal', accepting_games: true })

  const socketRef = useRef<WebSocket | null>(null)
  const toastId = useRef(0)

  const showError = useCallback((message: string) => {
    const id = ++toastId.current
    setToasts((current) => [...current, { id, message }])
    setTimeout(() => setToasts((current) => current.filter((toast) => toast.id !== id)), TOAST_MS)
  }, [])

  const dismissToast = useCallback((id: number) => {
    setToasts((current) => current.filter((toast) => toast.id !== id))
  }, [])

  useEffect(() => {
    let disposed = false
    let attempt = 0
    let busyAttempt = 0
    let retryTimer: ReturnType<typeof setTimeout> | undefined
    let trickTimer: ReturnType<typeof setTimeout> | undefined

    const handleMessage = (message: ServerMessage) => {
      switch (message.type) {
        case 'welcome':
          storeToken(message.token)
          setPlayerId(message.player_id)
          break
        case 'state': {
          setLobby(message.lobby)
          setGame(message.game)
          const trick = message.events.find((event) => event.type === 'trick_completed')
          if (trick) {
            setCompletedTrick(trick.trick)
            clearTimeout(trickTimer)
            trickTimer = setTimeout(() => setCompletedTrick(null), TRICK_DISPLAY_MS)
          }
          if (!message.game) setCompletedTrick(null)
          break
        }
        case 'server_status':
          setServerStatus({ load: message.load, accepting_games: message.accepting_games })
          break
        case 'error':
          showError(message.message)
          break
      }
    }

    const connect = () => {
      const socket = new WebSocket(websocketUrl())
      socketRef.current = socket

      // A "busy" error before the welcome means the server refused this connection.
      let welcomed = false
      let refused = false
      socket.onopen = () => {
        attempt = 0
        setStatus('open')
        socket.send(JSON.stringify({ type: 'hello', token: readToken() } satisfies ClientMessage))
      }
      socket.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data as string) as ServerMessage
          if (message.type === 'welcome') {
            welcomed = true
            busyAttempt = 0
          }
          // The "server is full" error arrives right before the 1013 close; the
          // banner explains it, so no toast.
          if (message.type === 'error' && message.code === 'busy' && !welcomed) {
            refused = true
            return
          }
          handleMessage(message)
        } catch (error) {
          console.error('Invalid message from server', error)
        }
      }
      socket.onclose = (event) => {
        if (disposed) return
        socketRef.current = null
        if (event.code === CLOSE_REPLACED) {
          setStatus('replaced')
          return
        }
        let delay: number
        if (event.code === CLOSE_TRY_AGAIN_LATER || refused) {
          setStatus('busy')
          delay = BUSY_RETRY_DELAYS_MS[Math.min(busyAttempt, BUSY_RETRY_DELAYS_MS.length - 1)]
          busyAttempt += 1
        } else {
          setStatus('closed')
          busyAttempt = 0
          delay = Math.min(1000 * 2 ** attempt, MAX_RETRY_DELAY_MS)
          attempt += 1
        }
        retryTimer = setTimeout(() => {
          setStatus('connecting')
          connect()
        }, delay)
      }
    }

    connect()
    return () => {
      disposed = true
      clearTimeout(retryTimer)
      clearTimeout(trickTimer)
      socketRef.current?.close()
      socketRef.current = null
    }
  }, [showError])

  const send = useCallback(
    (message: ClientMessage) => {
      const socket = socketRef.current
      if (socket?.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify(message))
      } else {
        showError('Not connected to the server')
      }
    },
    [showError],
  )

  const actions = useMemo(
    () => ({
      createLobby: (name: string) => send({ type: 'create_lobby', name }),
      joinLobby: (code: string, name: string) => send({ type: 'join_lobby', code, name }),
      leaveLobby: () => send({ type: 'leave_lobby' }),
      addAI: () => send({ type: 'add_ai' }),
      removePlayer: (playerId: string) => send({ type: 'remove_player', player_id: playerId }),
      setOptions: (options: GameOptions) => send({ type: 'set_options', ...options }),
      startGame: () => send({ type: 'start_game' }),
      play: (action: Action) => send({ type: 'action', action }),
      returnToLobby: () => send({ type: 'return_to_lobby' }),
    }),
    [send],
  )

  return { status, playerId, lobby, game, completedTrick, serverStatus, toasts, dismissToast, actions }
}

export type WizardActions = ReturnType<typeof useWizard>['actions']