stack/apps/web/src/hooks/useWebSocket.ts

import { useEffect, useState } from 'react';
import { io, Socket } from 'socket.io-client';

interface Task {
  id: string;
  [key: string]: unknown;
}

interface Event {
  id: string;
  [key: string]: unknown;
}

interface Project {
  id: string;
  [key: string]: unknown;
}

interface DeletePayload {
  id: string;
}

interface WebSocketCallbacks {
  onTaskCreated?: (task: Task) => void;
  onTaskUpdated?: (task: Task) => void;
  onTaskDeleted?: (payload: DeletePayload) => void;
  onEventCreated?: (event: Event) => void;
  onEventUpdated?: (event: Event) => void;
  onEventDeleted?: (payload: DeletePayload) => void;
  onProjectUpdated?: (project: Project) => void;
}

interface UseWebSocketReturn {
  isConnected: boolean;
  socket: Socket | null;
}

/**
 * Hook for managing WebSocket connections and real-time updates
 * 
 * @param workspaceId - The workspace ID to subscribe to
 * @param token - Authentication token
 * @param callbacks - Event callbacks for real-time updates
 * @returns Connection status and socket instance
 */
export function useWebSocket(
  workspaceId: string,
  token: string,
  callbacks: WebSocketCallbacks = {}
): UseWebSocketReturn {
  const [socket, setSocket] = useState<Socket | null>(null);
  const [isConnected, setIsConnected] = useState<boolean>(false);

  const {
    onTaskCreated,
    onTaskUpdated,
    onTaskDeleted,
    onEventCreated,
    onEventUpdated,
    onEventDeleted,
    onProjectUpdated,
  } = callbacks;

  useEffect(() => {
    // Get WebSocket URL from environment or default to API URL
    const wsUrl = process.env.NEXT_PUBLIC_API_URL || 'http://localhost:3001';

    // Create socket connection
    const newSocket = io(wsUrl, {
      auth: { token },
      query: { workspaceId },
    });

    setSocket(newSocket);

    // Connection event handlers
    const handleConnect = (): void => {
      setIsConnected(true);
    };

    const handleDisconnect = (): void => {
      setIsConnected(false);
    };

    newSocket.on('connect', handleConnect);
    newSocket.on('disconnect', handleDisconnect);

    // Real-time event handlers
    if (onTaskCreated) {
      newSocket.on('task:created', onTaskCreated);
    }
    if (onTaskUpdated) {
      newSocket.on('task:updated', onTaskUpdated);
    }
    if (onTaskDeleted) {
      newSocket.on('task:deleted', onTaskDeleted);
    }
    if (onEventCreated) {
      newSocket.on('event:created', onEventCreated);
    }
    if (onEventUpdated) {
      newSocket.on('event:updated', onEventUpdated);
    }
    if (onEventDeleted) {
      newSocket.on('event:deleted', onEventDeleted);
    }
    if (onProjectUpdated) {
      newSocket.on('project:updated', onProjectUpdated);
    }

    // Cleanup on unmount or dependency change
    return (): void => {
      newSocket.off('connect', handleConnect);
      newSocket.off('disconnect', handleDisconnect);
      
      if (onTaskCreated) newSocket.off('task:created', onTaskCreated);
      if (onTaskUpdated) newSocket.off('task:updated', onTaskUpdated);
      if (onTaskDeleted) newSocket.off('task:deleted', onTaskDeleted);
      if (onEventCreated) newSocket.off('event:created', onEventCreated);
      if (onEventUpdated) newSocket.off('event:updated', onEventUpdated);
      if (onEventDeleted) newSocket.off('event:deleted', onEventDeleted);
      if (onProjectUpdated) newSocket.off('project:updated', onProjectUpdated);
      
      newSocket.disconnect();
    };
  }, [
    workspaceId,
    token,
    onTaskCreated,
    onTaskUpdated,
    onTaskDeleted,
    onEventCreated,
    onEventUpdated,
    onEventDeleted,
    onProjectUpdated,
  ]);

  return {
    isConnected,
    socket,
  };
}