trumpsignal-frontend/lib/useRealtimeData.ts

'use client'

import { useEffect, useRef, useCallback } from 'react'

const WS_URL = process.env.NEXT_PUBLIC_WS_URL || 'ws://localhost:8000'

type PriceUpdate = { type: 'price'; asset: 'BTC' | 'ETH'; price: number }
type PostUpdate = { type: 'new_post'; post: object }
type WsMessage = PriceUpdate | PostUpdate

interface Handlers {
  onPrice?: (asset: 'BTC' | 'ETH', price: number) => void
  onNewPost?: (post: object) => void
}

export function usePriceSocket(handlers: Handlers) {
  const wsRef = useRef<WebSocket | null>(null)
  const handlersRef = useRef(handlers)
  const unmountedRef = useRef(false)
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  handlersRef.current = handlers

  const connect = useCallback(() => {
    if (unmountedRef.current) return
    const ws = new WebSocket(`${WS_URL}/ws/prices`)
    wsRef.current = ws

    ws.onmessage = (e) => {
      try {
        const msg: WsMessage = JSON.parse(e.data)
        if (msg.type === 'price' && handlersRef.current.onPrice) {
          handlersRef.current.onPrice(msg.asset, msg.price)
        } else if (msg.type === 'new_post' && handlersRef.current.onNewPost) {
          handlersRef.current.onNewPost(msg.post)
        }
      } catch (err) {
        console.warn('[WS] malformed message:', e.data, err)
      }
    }

    ws.onclose = () => {
      // Don't schedule a reconnect if the component already unmounted —
      // otherwise we leak a socket per navigation forever.
      if (unmountedRef.current) return
      reconnectTimerRef.current = setTimeout(connect, 3000)
    }

    ws.onerror = () => {
      ws.close()
    }
  }, [])

  useEffect(() => {
    unmountedRef.current = false
    connect()
    return () => {
      unmountedRef.current = true
      if (reconnectTimerRef.current) clearTimeout(reconnectTimerRef.current)
      wsRef.current?.close()
    }
  }, [connect])
}