frontend/src/hooks/useStreamingChat.js

/**
 * JARVIS — useStreamingChat hook
 * Connects to the WebSocket streaming endpoint and manages
 * real-time token delivery + tool use events.
 *
 * Usage:
 *   const { sendMessage, isStreaming, currentTokens, toolEvents } = useStreamingChat(sessionId)
 */

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

const WS_BASE = import.meta.env.VITE_WS_URL ||
  (window.location.protocol === 'https:' ? 'wss://' : 'ws://') +
  window.location.host + '/api/ws'

export function useStreamingChat(sessionId) {
  const [isStreaming,   setIsStreaming]   = useState(false)
  const [currentTokens, setCurrentTokens] = useState('')   // accumulates during stream
  const [toolEvents,    setToolEvents]    = useState([])
  const [error,         setError]         = useState(null)

  const ws        = useRef(null)
  const onDoneRef = useRef(null)

  // ── Connect / reconnect ─────────────────────────────────
  const connect = useCallback(() => {
    if (!sessionId) return
    if (ws.current?.readyState === WebSocket.OPEN) return

    const url = `${WS_BASE}/chat/${sessionId}`
    ws.current = new WebSocket(url)

    ws.current.onopen = () => {
      setError(null)
    }

    ws.current.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data)

        switch (msg.type) {
          case 'thinking':
            setCurrentTokens('...')
            break

          case 'token':
            setCurrentTokens((prev) => prev === '...' ? msg.data : prev + msg.data)
            break

          case 'tool_use':
            setToolEvents((prev) => [...prev, msg.data])
            break

          case 'done':
            setIsStreaming(false)
            if (onDoneRef.current) {
              onDoneRef.current({
                text:       currentTokens,
                used_tools: msg.data.used_tools || [],
                memory_hits: msg.data.memory_hits || 0,
              })
            }
            setCurrentTokens('')
            setToolEvents([])
            break

          case 'error':
            setError(msg.data)
            setIsStreaming(false)
            setCurrentTokens('')
            break

          default:
            break
        }
      } catch (e) {
        console.error('WS parse error:', e)
      }
    }

    ws.current.onerror = () => {
      setError('WebSocket connection error')
      setIsStreaming(false)
    }

    ws.current.onclose = () => {
      // Auto-reconnect after 2s if not intentionally closed
      setTimeout(() => {
        if (sessionId) connect()
      }, 2000)
    }
  }, [sessionId])

  // ── Connect on mount / sessionId change ─────────────────
  useEffect(() => {
    connect()
    return () => {
      ws.current?.close()
    }
  }, [connect])

  // ── Send message ─────────────────────────────────────────
  const sendMessage = useCallback(
    (message, imageCaption = null, onDone = null) => {
      if (!ws.current || ws.current.readyState !== WebSocket.OPEN) {
        connect()
        // Retry after connection
        setTimeout(() => sendMessage(message, imageCaption, onDone), 500)
        return
      }

      setIsStreaming(true)
      setCurrentTokens('')
      setToolEvents([])
      setError(null)
      onDoneRef.current = onDone

      ws.current.send(JSON.stringify({
        message,
        image_caption: imageCaption,
      }))
    },
    [connect],
  )

  // ── Disconnect ───────────────────────────────────────────
  const disconnect = useCallback(() => {
    ws.current?.close()
  }, [])

  return {
    sendMessage,
    disconnect,
    isStreaming,
    currentTokens,
    toolEvents,
    error,
    isConnected: ws.current?.readyState === WebSocket.OPEN,
  }
}