PROJECT-CONTORL/frontend/src/contexts/NotificationContext.tsx

import { createContext, useContext, useState, useEffect, useCallback, ReactNode, useRef } from 'react'
import { notificationsApi, Notification, NotificationListResponse } from '../services/collaboration'

interface NotificationContextType {
  notifications: Notification[]
  unreadCount: number
  loading: boolean
  error: string | null
  fetchNotifications: (isRead?: boolean) => Promise<void>
  markAsRead: (notificationId: string) => Promise<void>
  markAllAsRead: () => Promise<void>
  refreshUnreadCount: () => Promise<void>
}

const NotificationContext = createContext<NotificationContextType | undefined>(undefined)

const WS_RECONNECT_DELAY = 3000
const WS_PING_INTERVAL = 30000

export function NotificationProvider({ children }: { children: ReactNode }) {
  const [notifications, setNotifications] = useState<Notification[]>([])
  const [unreadCount, setUnreadCount] = useState(0)
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const wsRef = useRef<WebSocket | null>(null)
  const pingIntervalRef = useRef<ReturnType<typeof setInterval> | null>(null)
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const isMountedRef = useRef(true)

  const refreshUnreadCount = useCallback(async () => {
    try {
      const count = await notificationsApi.getUnreadCount()
      setUnreadCount(count)
    } catch (err) {
      console.error('Failed to fetch unread count:', err)
    }
  }, [])

  const fetchNotifications = useCallback(async (isRead?: boolean) => {
    setLoading(true)
    setError(null)
    try {
      const response: NotificationListResponse = await notificationsApi.list(isRead)
      setNotifications(response.notifications)
      setUnreadCount(response.unread_count)
    } catch (err) {
      setError('Failed to load notifications')
      console.error('Failed to fetch notifications:', err)
    } finally {
      setLoading(false)
    }
  }, [])

  const markAsRead = useCallback(async (notificationId: string) => {
    try {
      const updated = await notificationsApi.markAsRead(notificationId)
      setNotifications(prev =>
        prev.map(n => (n.id === notificationId ? updated : n))
      )
      setUnreadCount(prev => Math.max(0, prev - 1))
    } catch (err) {
      console.error('Failed to mark as read:', err)
    }
  }, [])

  const markAllAsRead = useCallback(async () => {
    try {
      await notificationsApi.markAllAsRead()
      setNotifications(prev =>
        prev.map(n => ({ ...n, is_read: true, read_at: new Date().toISOString() }))
      )
      setUnreadCount(0)
    } catch (err) {
      console.error('Failed to mark all as read:', err)
    }
  }, [])

  // WebSocket connection
  const connectWebSocket = useCallback(() => {
    const token = localStorage.getItem('token')
    if (!token) return

    // Use env var if available, otherwise derive from current location
    let wsUrl: string
    const envWsUrl = import.meta.env.VITE_WS_URL
    if (envWsUrl) {
      wsUrl = `${envWsUrl}/ws/notifications?token=${token}`
    } else {
      const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
      wsUrl = `${wsProtocol}//${window.location.host}/ws/notifications?token=${token}`
    }

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

      ws.onopen = () => {
        console.log('WebSocket connected')
        // Start ping interval
        pingIntervalRef.current = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'ping' }))
          }
        }, WS_PING_INTERVAL)
      }

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

          switch (message.type) {
            case 'connected':
              console.log('WebSocket authenticated:', message.data.message)
              break

            case 'unread_sync':
              // Merge unread notifications without removing already-loaded notifications
              setNotifications(prev => {
                const unreadNotifications = message.data.notifications || []
                const existingIds = new Set(prev.map(n => n.id))

                // Add new unread notifications that don't exist in current list
                const newNotifications = unreadNotifications.filter(
                  (n: Notification) => !existingIds.has(n.id)
                )

                // Update existing unread notifications and prepend new ones
                const updated = prev.map(existing => {
                  const fromSync = unreadNotifications.find((n: Notification) => n.id === existing.id)
                  return fromSync || existing
                })

                return [...newNotifications, ...updated]
              })
              setUnreadCount(message.data.unread_count || 0)
              break

            case 'notification':
              // Add new notification to the top
              setNotifications(prev => [message.data, ...prev])
              setUnreadCount(prev => prev + 1)
              break

            case 'unread_count':
              setUnreadCount(message.data.unread_count)
              break

            case 'ping':
              // Server ping - respond with pong
              if (ws.readyState === WebSocket.OPEN) {
                ws.send(JSON.stringify({ type: 'pong' }))
              }
              break

            case 'pong':
              // Pong received, connection is alive
              break
          }
        } catch (err) {
          console.error('Failed to parse WebSocket message:', err)
        }
      }

      ws.onclose = () => {
        console.log('WebSocket disconnected')
        // Clear ping interval
        if (pingIntervalRef.current) {
          clearInterval(pingIntervalRef.current)
          pingIntervalRef.current = null
        }

        // Only attempt to reconnect if component is still mounted
        if (isMountedRef.current) {
          reconnectTimeoutRef.current = setTimeout(() => {
            if (isMountedRef.current) {
              connectWebSocket()
            }
          }, WS_RECONNECT_DELAY)
        }
      }

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

  // Initial fetch and WebSocket connection
  useEffect(() => {
    isMountedRef.current = true
    const token = localStorage.getItem('token')
    if (token) {
      // Fetch both unread count and initial notifications
      refreshUnreadCount()
      fetchNotifications()
      // Delay WebSocket connection to avoid StrictMode race condition
      const connectTimeout = setTimeout(() => {
        if (isMountedRef.current) {
          connectWebSocket()
        }
      }, 100)

      return () => {
        clearTimeout(connectTimeout)
        isMountedRef.current = false
        // Cleanup on unmount
        if (wsRef.current) {
          wsRef.current.close()
        }
        if (pingIntervalRef.current) {
          clearInterval(pingIntervalRef.current)
        }
        if (reconnectTimeoutRef.current) {
          clearTimeout(reconnectTimeoutRef.current)
        }
      }
    }

    return () => {
      isMountedRef.current = false
    }
  }, [refreshUnreadCount, fetchNotifications, connectWebSocket])

  return (
    <NotificationContext.Provider
      value={{
        notifications,
        unreadCount,
        loading,
        error,
        fetchNotifications,
        markAsRead,
        markAllAsRead,
        refreshUnreadCount,
      }}
    >
      {children}
    </NotificationContext.Provider>
  )
}

export function useNotifications() {
  const context = useContext(NotificationContext)
  if (context === undefined) {
    throw new Error('useNotifications must be used within a NotificationProvider')
  }
  return context
}