hazardousWaste_web/src/utils/ws.js

import { getToken } from '@/utils/auth'

let ws = null
let reconnectTimer = null
let manualClose = false
let retryCount = 0
const subscribers = new Set()

const WS_PATH = '/ws'
const MAX_RECONNECT_DELAY = 30000

function getWsBaseUrl() {
  const baseApi = import.meta.env.VITE_APP_BASE_API || ''

  if (/^https?:\/\//.test(baseApi)) {
    return baseApi.replace(/^http/, 'ws').replace(/\/$/, '')
  }

  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
  return `${protocol}//${window.location.host}${baseApi}`.replace(/\/$/, '')
}

function buildWsUrl() {
  const url = new URL(`${getWsBaseUrl()}${WS_PATH}`)
  const token = getToken()

  if (token) {
    url.searchParams.set('token', token)
  }

  return url.toString()
}

export function connectWs(onMessage) {
  console.trace('connectWs called')
  if (!getToken()) {
    console.warn('[WebSocket] missing token')
    return
  }

  if (typeof onMessage === 'function') {
    subscribers.add(onMessage)
  }

  manualClose = false
  clearReconnectTimer()

  if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
    return
  }

  ws = new WebSocket(buildWsUrl())

  ws.onopen = () => {
    retryCount = 0
    console.log('[WebSocket] connected')
  }

  ws.onmessage = (event) => {
    let data = event.data

    try {
      data = JSON.parse(event.data)
    } catch (e) {
      // Keep plain text messages as-is.
    }

    subscribers.forEach((handler) => {
      try {
        handler(data, event)
      } catch (error) {
        console.error('[WebSocket] subscriber error', error)
      }
    })
  }

  ws.onclose = (event) => {
    console.log('[WebSocket] closed', { code: event.code, reason: event.reason })
    ws = null

    if (isNonRetryableClose(event)) {
      console.warn('[WebSocket] stop reconnect because close reason is not retryable', {
        code: event.code,
        reason: event.reason
      })
      retryCount = 0
      clearReconnectTimer()
      return
    }

    if (!manualClose && subscribers.size > 0) {
      reconnectTimer = setTimeout(() => {
        console.log('[WebSocket] reconnecting')
        connectWs()
      }, getReconnectDelay())
    }
  }

  ws.onerror = (error) => {
    console.error('[WebSocket] error', error)
  }
}

export function closeWs(onMessage) {
  if (typeof onMessage === 'function') {
    subscribers.delete(onMessage)
  } else {
    subscribers.clear()
  }

  if (subscribers.size > 0) {
    return
  }

  manualClose = true
  retryCount = 0
  clearReconnectTimer()
  if (ws) {
    ws.close()
    ws = null
  }
}

export function reconnectWs(onMessage) {
  if (typeof onMessage === 'function') {
    subscribers.add(onMessage)
  }

  manualClose = false
  retryCount = 0
  clearReconnectTimer()

  if (ws) {
    ws.close()
    ws = null
  }

  manualClose = false
  connectWs()
}

export function sendWs(data) {
  if (!ws || ws.readyState !== WebSocket.OPEN) {
    console.warn('[WebSocket] send failed, socket is not open')
    return false
  }

  ws.send(typeof data === 'string' ? data : JSON.stringify(data))
  return true
}

export function getWsState() {
  return ws ? ws.readyState : WebSocket.CLOSED
}

function getReconnectDelay() {
  retryCount += 1
  return Math.min(5000 * retryCount, MAX_RECONNECT_DELAY)
}

function clearReconnectTimer() {
  if (reconnectTimer) {
    clearTimeout(reconnectTimer)
    reconnectTimer = null
  }
}

function isNonRetryableClose(event) {
  const reason = (event?.reason || '').toLowerCase()
  return (
    reason.includes('invalid token') ||
    reason.includes('missing token') ||
    reason.includes('deptid is null')
  )
}