smartbotic-crm/webui/src/contexts/WebSocketContext.tsx

// WebSocket context for real-time updates (US-036)

import { createContext, useContext, useEffect, useRef, useState, useCallback } from 'react'
import { useAuth } from './AuthContext'

const ACCESS_TOKEN_KEY = 'smartbotic_access_token'

interface WebSocketMessage {
  type: string
  [key: string]: unknown
}

interface Notification {
  id: string
  message: string
  level: 'info' | 'warn' | 'error'
  timestamp: string
  read: boolean
}

interface DocumentEvent {
  workspace_id: string
  collection: string
  action: 'create' | 'update' | 'delete'
  document_id: string
  data?: unknown
}

interface WebSocketContextType {
  isConnected: boolean
  notifications: Notification[]
  unreadCount: number
  subscribe: (workspaceId: string, collection: string) => void
  unsubscribe: (workspaceId: string, collection: string) => void
  markAsRead: (notificationId: string) => void
  markAllAsRead: () => void
  clearNotifications: () => void
  onDocumentEvent: (callback: (event: DocumentEvent) => void) => () => void
}

const WebSocketContext = createContext<WebSocketContextType | undefined>(undefined)

const WS_URL = `ws://${window.location.hostname}:18081/ws`
const RECONNECT_DELAY = 3000
const MAX_NOTIFICATIONS = 50

export function WebSocketProvider({ children }: { children: React.ReactNode }) {
  const { isAuthenticated, isLoading: isAuthLoading } = useAuth()
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined)
  const documentEventCallbacksRef = useRef<Set<(event: DocumentEvent) => void>>(new Set())

  const [isConnected, setIsConnected] = useState(false)
  const [notifications, setNotifications] = useState<Notification[]>([])
  const [subscriptions, setSubscriptions] = useState<Set<string>>(new Set())

  const unreadCount = notifications.filter((n) => !n.read).length

  // Get access token from localStorage
  const getAccessToken = useCallback(() => {
    return localStorage.getItem(ACCESS_TOKEN_KEY)
  }, [])

  // Connect to WebSocket
  const connect = useCallback(() => {
    // Don't connect while auth is still loading
    if (isAuthLoading) {
      return
    }
    const accessToken = getAccessToken()
    if (!isAuthenticated || !accessToken || wsRef.current?.readyState === WebSocket.OPEN) {
      return
    }

    try {
      const ws = new WebSocket(WS_URL)
      wsRef.current = ws

      ws.onopen = () => {
        console.log('WebSocket connected')
        // Send authentication
        ws.send(JSON.stringify({ type: 'auth', token: accessToken }))
      }

      ws.onmessage = (event) => {
        try {
          const message: WebSocketMessage = JSON.parse(event.data)
          handleMessage(message)
        } catch (err) {
          console.error('Failed to parse WebSocket message:', err)
        }
      }

      ws.onclose = () => {
        console.log('WebSocket disconnected')
        setIsConnected(false)
        wsRef.current = null

        // Attempt reconnection
        if (isAuthenticated) {
          reconnectTimeoutRef.current = setTimeout(connect, RECONNECT_DELAY)
        }
      }

      ws.onerror = (error) => {
        console.error('WebSocket error:', error)
      }
    } catch (err) {
      console.error('Failed to create WebSocket connection:', err)
    }
  }, [isAuthenticated, isAuthLoading, getAccessToken])

  // Handle incoming messages
  const handleMessage = useCallback((message: WebSocketMessage) => {
    switch (message.type) {
      case 'auth_success':
        setIsConnected(true)
        // Re-subscribe to previous subscriptions
        subscriptions.forEach((key) => {
          const [workspaceId, collection] = key.split(':')
          wsRef.current?.send(
            JSON.stringify({ type: 'subscribe', workspace_id: workspaceId, collection })
          )
        })
        break

      case 'auth_error':
        console.error('WebSocket auth error:', message.error)
        wsRef.current?.close()
        break

      case 'notification':
        addNotification({
          id: crypto.randomUUID(),
          message: String(message.message || 'New notification'),
          level: (message.level as 'info' | 'warn' | 'error') || 'info',
          timestamp: new Date().toISOString(),
          read: false,
        })
        break

      case 'document': {
        // Notify document event listeners
        const docEvent: DocumentEvent = {
          workspace_id: String(message.workspace_id || ''),
          collection: String(message.collection || ''),
          action: message.action as 'create' | 'update' | 'delete',
          document_id: String(message.document_id || ''),
          data: message.data,
        }
        documentEventCallbacksRef.current.forEach((cb) => cb(docEvent))
        break
      }

      case 'subscribed':
        console.log(`Subscribed to ${message.workspace_id}:${message.collection}`)
        break

      case 'unsubscribed':
        console.log(`Unsubscribed from ${message.workspace_id}:${message.collection}`)
        break

      case 'error':
        console.error('WebSocket error message:', message.error)
        break

      default:
        console.log('Unknown WebSocket message type:', message.type)
    }
  }, [subscriptions])

  // Add notification
  const addNotification = useCallback((notification: Notification) => {
    setNotifications((prev) => {
      const updated = [notification, ...prev]
      // Keep only the last MAX_NOTIFICATIONS
      return updated.slice(0, MAX_NOTIFICATIONS)
    })
  }, [])

  // Subscribe to collection changes
  const subscribe = useCallback((workspaceId: string, collection: string) => {
    const key = `${workspaceId}:${collection}`
    setSubscriptions((prev) => new Set(prev).add(key))

    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: 'subscribe', workspace_id: workspaceId, collection }))
    }
  }, [])

  // Unsubscribe from collection changes
  const unsubscribe = useCallback((workspaceId: string, collection: string) => {
    const key = `${workspaceId}:${collection}`
    setSubscriptions((prev) => {
      const updated = new Set(prev)
      updated.delete(key)
      return updated
    })

    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(
        JSON.stringify({ type: 'unsubscribe', workspace_id: workspaceId, collection })
      )
    }
  }, [])

  // Mark notification as read
  const markAsRead = useCallback((notificationId: string) => {
    setNotifications((prev) =>
      prev.map((n) => (n.id === notificationId ? { ...n, read: true } : n))
    )
  }, [])

  // Mark all notifications as read
  const markAllAsRead = useCallback(() => {
    setNotifications((prev) => prev.map((n) => ({ ...n, read: true })))
  }, [])

  // Clear all notifications
  const clearNotifications = useCallback(() => {
    setNotifications([])
  }, [])

  // Register document event callback
  const onDocumentEvent = useCallback((callback: (event: DocumentEvent) => void) => {
    documentEventCallbacksRef.current.add(callback)
    return () => {
      documentEventCallbacksRef.current.delete(callback)
    }
  }, [])

  // Connect on auth change
  useEffect(() => {
    // Wait for auth to finish loading
    if (isAuthLoading) {
      return
    }

    if (isAuthenticated) {
      connect()
    } else {
      // Disconnect on logout
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      setIsConnected(false)
      setSubscriptions(new Set())
    }

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
    }
  }, [isAuthenticated, isAuthLoading, connect])

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (wsRef.current) {
        wsRef.current.close()
      }
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
    }
  }, [])

  const value: WebSocketContextType = {
    isConnected,
    notifications,
    unreadCount,
    subscribe,
    unsubscribe,
    markAsRead,
    markAllAsRead,
    clearNotifications,
    onDocumentEvent,
  }

  return <WebSocketContext.Provider value={value}>{children}</WebSocketContext.Provider>
}

export function useWebSocket() {
  const context = useContext(WebSocketContext)
  if (context === undefined) {
    throw new Error('useWebSocket must be used within a WebSocketProvider')
  }
  return context
}