trumpsignal-frontend/lib/wsContext.tsx

'use client'

/**
 * Singleton WebSocket provider.
 *
 * The entire app shares ONE connection to /ws/prices. Components subscribe
 * to specific message types via `useWsSubscribe` — no component owns the socket.
 *
 * Why: previously DashboardClient (via usePriceSocket) and SignalMonitor each
 * opened their own WebSocket, giving the server two connections per page load
 * and causing each component to receive every broadcast twice.
 */

import {
  createContext,
  useContext,
  useEffect,
  useRef,
  type ReactNode,
} from 'react'

type Handler = (msg: unknown) => void

interface WsContextValue {
  subscribe: (type: string, fn: Handler) => () => void
}

const WsContext = createContext<WsContextValue | null>(null)

// In production, NEXT_PUBLIC_WS_URL must be set (e.g. wss://api.trumpsignal.com).
// Fallback auto-upgrades ws→wss when the page is served over HTTPS to avoid
// mixed-content blocks, and falls back to ws:// for local dev.
const _wsEnv = process.env.NEXT_PUBLIC_WS_URL
const WS_URL = _wsEnv || (
  typeof window !== 'undefined' && window.location.protocol === 'https:'
    ? 'wss://localhost:8000'
    : 'ws://localhost:8000'
)

// Exponential backoff: 2s → 4s → 8s → … capped at 60s, plus ±30% jitter
// so a server restart doesn't get hit by all clients at the same instant.
const WS_RETRY_BASE_MS  = 2_000
const WS_RETRY_MAX_MS   = 60_000
function retryDelay(attempt: number): number {
  const exp = Math.min(WS_RETRY_BASE_MS * 2 ** attempt, WS_RETRY_MAX_MS)
  return exp * (0.7 + Math.random() * 0.6)   // ±30% jitter
}

export function WsProvider({ children }: { children: ReactNode }) {
  // Stable map: message-type → set of handlers. Never replaced, only mutated.
  const subs = useRef<Map<string, Set<Handler>>>(new Map())

  useEffect(() => {
    let dead = false
    let retryTimer: ReturnType<typeof setTimeout> | null = null
    let attempt = 0
    let socket: WebSocket | null = null

    function connect() {
      if (dead) return
      socket = new WebSocket(`${WS_URL}/ws/prices`)
      const ws = socket   // local alias for handler closures

      ws.onopen = () => {
        attempt = 0   // reset backoff on successful connection
      }

      ws.onmessage = (e) => {
        try {
          const msg = JSON.parse(e.data) as { type?: string }
          const t = msg.type ?? '__unknown__'
          subs.current.get(t)?.forEach((fn) => fn(msg))
        } catch {
          // ignore malformed frames
        }
      }

      ws.onclose = () => {
        if (!dead) {
          const delay = retryDelay(attempt++)
          retryTimer = setTimeout(connect, delay)
        }
      }
      ws.onerror = () => ws.close()
    }

    // Reconnect immediately when the tab comes back to the foreground
    // (the socket may have silently died while the tab was hidden).
    function onVisible() {
      if (document.visibilityState !== 'visible') return
      if (socket && socket.readyState === WebSocket.OPEN) return
      // Cancel any pending retry — we're reconnecting now.
      if (retryTimer) { clearTimeout(retryTimer); retryTimer = null }
      attempt = 0
      if (socket && socket.readyState < WebSocket.CLOSING) {
        socket.onclose = null
        socket.close()
      }
      connect()
    }

    document.addEventListener('visibilitychange', onVisible)
    connect()
    return () => {
      dead = true
      document.removeEventListener('visibilitychange', onVisible)
      if (retryTimer) clearTimeout(retryTimer)
      // Actively close any live socket so the OS-level connection releases
      // immediately on unmount instead of waiting for the next server keepalive.
      if (socket && socket.readyState <= WebSocket.OPEN) {
        socket.onclose = null   // suppress the auto-reconnect we'd otherwise queue
        socket.close()
      }
      socket = null
    }
  }, [])

  function subscribe(type: string, fn: Handler): () => void {
    if (!subs.current.has(type)) subs.current.set(type, new Set())
    subs.current.get(type)!.add(fn)
    return () => subs.current.get(type)?.delete(fn)
  }

  return <WsContext.Provider value={{ subscribe }}>{children}</WsContext.Provider>
}

/**
 * Subscribe to a specific WebSocket message type from the shared connection.
 *
 * @param type    The `msg.type` string to listen for (e.g. 'price', 'new_post', 'funding_signal')
 * @param handler Called with the full parsed message object on every matching frame.
 *                Always receives the latest handler reference — no stale-closure issues.
 *
 * Usage:
 *   useWsSubscribe('price', (msg) => { ... })
 */
export function useWsSubscribe(type: string, handler: Handler): void {
  const ctx = useContext(WsContext)
  // Keep a ref so the subscription closure never goes stale even if handler
  // is an inline function that changes every render.
  const handlerRef = useRef<Handler>(handler)
  handlerRef.current = handler

  useEffect(() => {
    if (!ctx) return
    return ctx.subscribe(type, (msg) => handlerRef.current(msg))
  }, [ctx, type]) // type is stable in practice; ctx never changes
}