attune/web/src/hooks/useEnforcementStream.ts

import { useCallback } from "react";
import { useQueryClient } from "@tanstack/react-query";
import {
  useEntityNotifications,
  type Notification,
} from "@/contexts/WebSocketContext";
import type { EnforcementSummary } from "@/api";

interface UseEnforcementStreamOptions {
  /**
   * Optional enforcement ID to filter updates for a specific enforcement.
   * If not provided, receives updates for all enforcements.
   */
  enforcementId?: number;

  /**
   * Whether the stream should be active.
   * Defaults to true.
   */
  enabled?: boolean;
}

/** Shape of data coming from WebSocket notifications for enforcements */
interface EnforcementNotification {
  entity_id: number;
  entity_type: string;
  notification_type: string;
  payload: Partial<EnforcementSummary> & Record<string, unknown>;
  timestamp: string;
}

/** Query params shape used in enforcement list query keys */
interface EnforcementQueryParams {
  status?: string;
  event?: number;
  rule?: number;
  triggerRef?: string;
  ruleRef?: string;
}

/** Shape of the paginated API response stored in React Query cache */
interface EnforcementListCache {
  data: EnforcementSummary[];
  pagination?: {
    total_items?: number;
    page?: number;
    page_size?: number;
  };
}

/** Shape of a single enforcement detail response stored in React Query cache */
interface EnforcementDetailCache {
  data: EnforcementSummary;
}

/**
 * Check if an enforcement matches the given query parameters
 * Only checks fields that are reliably present in WebSocket payloads
 */
function enforcementMatchesParams(
  enforcement: Partial<EnforcementSummary>,
  params: EnforcementQueryParams | undefined,
): boolean {
  if (!params) return true;

  // Check status filter
  if (params.status && enforcement.status !== params.status) {
    return false;
  }

  // Check event filter
  if (params.event !== undefined && enforcement.event !== params.event) {
    return false;
  }

  // Check rule filter
  if (params.rule !== undefined && enforcement.rule !== params.rule) {
    return false;
  }

  // Check trigger_ref filter (always present)
  if (params.triggerRef && enforcement.trigger_ref !== params.triggerRef) {
    return false;
  }

  // Note: rule_ref is NOT checked here for new enforcements because it may not be
  // present in WebSocket payloads. For this filter, we only update existing
  // enforcements, never add new ones.

  return true;
}

/**
 * Check if query params include filters not present in WebSocket payloads
 */
function hasUnsupportedFilters(
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  _params: EnforcementQueryParams | undefined,
): boolean {
  // Currently all enforcement filters are supported in WebSocket payloads
  return false;
}

/**
 * Hook to subscribe to real-time enforcement updates via WebSocket.
 *
 * Automatically reconnects on connection loss and updates React Query cache
 * when enforcement updates are received.
 *
 * @example
 * ```tsx
 * // Listen to all enforcement updates
 * useEnforcementStream();
 *
 * // Listen to updates for a specific enforcement
 * useEnforcementStream({ enforcementId: 123 });
 * ```
 */
export function useEnforcementStream(
  options: UseEnforcementStreamOptions = {},
) {
  const { enforcementId, enabled = true } = options;
  const queryClient = useQueryClient();

  const handleNotification = useCallback(
    (raw: Notification) => {
      const notification = raw as unknown as EnforcementNotification;
      // Filter by enforcement ID if specified
      if (enforcementId && notification.entity_id !== enforcementId) {
        return;
      }

      // Extract enforcement data from notification payload (flat structure)
      const enforcementData =
        notification.payload as Partial<EnforcementSummary>;

      // Update specific enforcement query if it exists
      queryClient.setQueryData(
        ["enforcements", notification.entity_id],
        (old: EnforcementDetailCache | undefined) => {
          if (!old) return old;
          return {
            ...old,
            data: {
              ...old.data,
              ...enforcementData,
            },
          };
        },
      );

      // Update enforcement list queries by modifying existing data
      // We need to iterate manually to access query keys for filtering
      const queries = queryClient
        .getQueriesData<EnforcementListCache>({
          queryKey: ["enforcements"],
          exact: false,
        })
        .filter(([, data]) => data && Array.isArray(data?.data));

      queries.forEach(([queryKey, oldData]) => {
        // Extract query params from the query key (format: ["enforcements", params])
        const queryParams = queryKey[1] as EnforcementQueryParams | undefined;

        const old = oldData as EnforcementListCache;

        // Check if enforcement already exists in the list
        const existingIndex = old.data.findIndex(
          (enf) => enf.id === notification.entity_id,
        );

        let updatedData: EnforcementSummary[];
        if (existingIndex >= 0) {
          // Always update existing enforcement in the list
          updatedData = [...old.data];
          updatedData[existingIndex] = {
            ...updatedData[existingIndex],
            ...enforcementData,
          };

          // Note: We don't remove enforcements from cache based on filters.
          // The cache represents what the API query returned.
          // Client-side filtering (in the page component) handles what's displayed.
        } else {
          // For new enforcements, be conservative with filters we can't verify
          if (hasUnsupportedFilters(queryParams)) {
            // Don't add new enforcement when using filters we can't verify
            return;
          }

          // Only add new enforcement if it matches the query parameters
          if (enforcementMatchesParams(enforcementData, queryParams)) {
            // Add to beginning and cap at 50 items to prevent performance issues
            updatedData = [
              enforcementData as EnforcementSummary,
              ...old.data,
            ].slice(0, 50);
          } else {
            // Don't modify the list if the new enforcement doesn't match the query
            return;
          }
        }

        // Update the query with the new data
        queryClient.setQueryData(queryKey, {
          ...old,
          data: updatedData,
          pagination: {
            ...old.pagination,
            total_items: (old.pagination?.total_items || 0) + 1,
          },
        });
      });

      // Also update related queries (rules and events enforcements)
      if (enforcementData.rule) {
        queryClient.invalidateQueries({
          queryKey: ["rules", enforcementData.rule, "enforcements"],
        });
      }
      if (enforcementData.event) {
        queryClient.invalidateQueries({
          queryKey: ["events", enforcementData.event, "enforcements"],
        });
      }
    },
    [enforcementId, queryClient],
  );

  const { connected } = useEntityNotifications(
    "enforcement",
    handleNotification,
    enabled,
  );

  return {
    isConnected: connected,
    error: null,
  };
}