1  'use client'
  2  
  3  import { useEffect, useRef } from 'react'
  4  import { subscribeWs, unsubscribeWs, isWsConnected, onWsStateChange, offWsStateChange } from '@/lib/ws-client'
  5  import { hmrSingleton } from '@/lib/shared-utils'
  6  import { usePageActive } from './use-page-active'
  7  
  8  /** Shared fallback intervals keyed by topic — multiple useWs instances share one interval. */
  9  const sharedFallbacks = hmrSingleton('useWs_sharedFallbacks', () => new Map<string, {
 10    interval: ReturnType<typeof setInterval> | null
 11    handlers: Set<() => void>
 12    ms: number
 13  }>())
 14  
 15  function runAllHandlers(topic: string): void {
 16    const entry = sharedFallbacks.get(topic)
 17    if (!entry) return
 18    for (const h of entry.handlers) h()
 19  }
 20  
 21  function acquireFallback(topic: string, ms: number, handler: () => void): void {
 22    const existing = sharedFallbacks.get(topic)
 23    if (existing) {
 24      existing.handlers.add(handler)
 25      return
 26    }
 27    const handlers = new Set<() => void>([handler])
 28    const entry = { interval: null as ReturnType<typeof setInterval> | null, handlers, ms }
 29    sharedFallbacks.set(topic, entry)
 30    if (!isWsConnected()) {
 31      entry.interval = setInterval(() => runAllHandlers(topic), ms)
 32    }
 33  }
 34  
 35  function releaseFallback(topic: string, handler: () => void): void {
 36    const entry = sharedFallbacks.get(topic)
 37    if (!entry) return
 38    entry.handlers.delete(handler)
 39    if (entry.handlers.size <= 0) {
 40      if (entry.interval) clearInterval(entry.interval)
 41      sharedFallbacks.delete(topic)
 42    }
 43  }
 44  
 45  function syncFallbacks(): void {
 46    const connected = isWsConnected()
 47    for (const [topic, entry] of sharedFallbacks) {
 48      if (connected && entry.interval) {
 49        clearInterval(entry.interval)
 50        entry.interval = null
 51      } else if (!connected && !entry.interval) {
 52        entry.interval = setInterval(() => runAllHandlers(topic), entry.ms)
 53      }
 54    }
 55  }
 56  
 57  /**
 58   * Subscribe to a WebSocket topic. Calls `handler` on push events.
 59   * Falls back to polling at `fallbackMs` when WS is disconnected.
 60   */
 61  export function useWs(topic: string, handler: () => void | Promise<void>, fallbackMs?: number) {
 62    const isActive = usePageActive()
 63    const handlerRef = useRef(handler)
 64    const fallbackMsRef = useRef(fallbackMs)
 65    const inFlightRef = useRef<Promise<void> | null>(null)
 66    const wasActiveRef = useRef(isActive)
 67  
 68    useEffect(() => {
 69      handlerRef.current = handler
 70      fallbackMsRef.current = fallbackMs
 71    }, [handler, fallbackMs])
 72  
 73    const runHandler = () => {
 74      if (inFlightRef.current) return
 75      try {
 76        const result = handlerRef.current()
 77        if (result && typeof (result as PromiseLike<void>).then === 'function') {
 78          const promise = Promise.resolve(result)
 79            .catch(() => {})
 80            .finally(() => {
 81              if (inFlightRef.current === promise) {
 82                inFlightRef.current = null
 83              }
 84            })
 85          inFlightRef.current = promise
 86        }
 87      } catch {
 88        // Individual handlers already own their error reporting
 89      }
 90    }
 91  
 92    // WS subscription — only re-runs when topic changes
 93    useEffect(() => {
 94      if (!topic) return
 95  
 96      const cb = () => runHandler()
 97      subscribeWs(topic, cb)
 98      return () => { unsubscribeWs(topic, cb) }
 99    }, [topic])
100  
101    // Stable handler ref for fallback — identity stays the same across renders
102    const fallbackHandlerRef = useRef(() => runHandler())
103    useEffect(() => {
104      fallbackHandlerRef.current = () => runHandler()
105    })
106  
107    // Fallback polling with shared intervals and connection state notifications
108    useEffect(() => {
109      if (!topic) return
110  
111      const becameActive = !wasActiveRef.current && isActive
112      wasActiveRef.current = isActive
113  
114      // When page becomes visible again, fire an immediate refresh for data-fetch topics
115      if (becameActive && fallbackMsRef.current && fallbackMsRef.current > 0 && !isWsConnected()) {
116        runHandler()
117      }
118  
119      // Don't run polling while the tab is hidden
120      if (!isActive) return
121  
122      const ms = fallbackMsRef.current
123      if (!ms || ms <= 0) return
124  
125      // Subscribe to connection state changes to start/stop fallback
126      const stateHandler = () => syncFallbacks()
127      onWsStateChange(stateHandler)
128      // Use a stable wrapper that delegates to the current handler ref
129      const stableFallbackHandler = () => fallbackHandlerRef.current()
130      acquireFallback(topic, ms, stableFallbackHandler)
131  
132      return () => {
133        offWsStateChange(stateHandler)
134        releaseFallback(topic, stableFallbackHandler)
135      }
136    }, [topic, isActive])
137  }