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

import { useEffect, useRef, useState } from "react";
import type { Socket } from "socket.io-client";
import { io } from "socket.io-client";
import { API_BASE_URL } from "@/lib/config";

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 ConnectionError {
  message: string;
  type: string;
  description?: string;
}

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

/**
 * Check if the WebSocket URL uses secure protocol (wss://)
 * Logs a warning in production when using insecure ws://
 */
function validateWebSocketSecurity(url: string): void {
  const isProduction = process.env.NODE_ENV === "production";
  const isSecure = url.startsWith("https://") || url.startsWith("wss://");

  if (isProduction && !isSecure) {
    console.warn(
      "[Security Warning] WebSocket connection using insecure protocol (ws://). " +
        "Authentication tokens may be exposed. Use wss:// in production."
    );
  }
}

/**
 * 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, socket instance, and connection error
 */
export function useWebSocket(
  workspaceId: string,
  token: string,
  callbacks: WebSocketCallbacks = {}
): UseWebSocketReturn {
  const [socket, setSocket] = useState<Socket | null>(null);
  const [isConnected, setIsConnected] = useState<boolean>(false);
  const [connectionError, setConnectionError] = useState<ConnectionError | null>(null);

  // Use refs for callbacks to prevent stale closures and unnecessary reconnects
  // This ensures that callback changes don't trigger useEffect re-runs
  const callbacksRef = useRef<WebSocketCallbacks>(callbacks);

  // Keep the ref up-to-date with the latest callbacks
  useEffect(() => {
    callbacksRef.current = callbacks;
  }, [callbacks]);

  useEffect(() => {
    // Use WebSocket URL from central config
    const wsUrl = API_BASE_URL;

    // Validate WebSocket security - warn if using insecure connection in production
    validateWebSocketSecurity(wsUrl);

    // Clear any previous connection error
    setConnectionError(null);

    // 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);
    };

    const handleConnectError = (error: Error): void => {
      setConnectionError({
        message: error.message || "Connection failed",
        type: "connect_error",
        description: "Failed to establish WebSocket connection",
      });
      setIsConnected(false);
    };

    // Wrapper functions that always use the latest callbacks via ref
    // This prevents stale closure issues while avoiding reconnects on callback changes
    const handleTaskCreated = (task: Task): void => {
      callbacksRef.current.onTaskCreated?.(task);
    };

    const handleTaskUpdated = (task: Task): void => {
      callbacksRef.current.onTaskUpdated?.(task);
    };

    const handleTaskDeleted = (payload: DeletePayload): void => {
      callbacksRef.current.onTaskDeleted?.(payload);
    };

    const handleEventCreated = (event: Event): void => {
      callbacksRef.current.onEventCreated?.(event);
    };

    const handleEventUpdated = (event: Event): void => {
      callbacksRef.current.onEventUpdated?.(event);
    };

    const handleEventDeleted = (payload: DeletePayload): void => {
      callbacksRef.current.onEventDeleted?.(payload);
    };

    const handleProjectUpdated = (project: Project): void => {
      callbacksRef.current.onProjectUpdated?.(project);
    };

    newSocket.on("connect", handleConnect);
    newSocket.on("disconnect", handleDisconnect);
    newSocket.on("connect_error", handleConnectError);

    // Register all event handlers - they'll check the ref for actual callbacks
    newSocket.on("task:created", handleTaskCreated);
    newSocket.on("task:updated", handleTaskUpdated);
    newSocket.on("task:deleted", handleTaskDeleted);
    newSocket.on("event:created", handleEventCreated);
    newSocket.on("event:updated", handleEventUpdated);
    newSocket.on("event:deleted", handleEventDeleted);
    newSocket.on("project:updated", handleProjectUpdated);

    // Cleanup on unmount or dependency change
    return (): void => {
      newSocket.off("connect", handleConnect);
      newSocket.off("disconnect", handleDisconnect);
      newSocket.off("connect_error", handleConnectError);

      newSocket.off("task:created", handleTaskCreated);
      newSocket.off("task:updated", handleTaskUpdated);
      newSocket.off("task:deleted", handleTaskDeleted);
      newSocket.off("event:created", handleEventCreated);
      newSocket.off("event:updated", handleEventUpdated);
      newSocket.off("event:deleted", handleEventDeleted);
      newSocket.off("project:updated", handleProjectUpdated);

      newSocket.disconnect();
    };
    // Only stable values in deps - callbacks are accessed via ref
  }, [workspaceId, token]);

  return {
    isConnected,
    socket,
    connectionError,
  };
}