src/utils/optimisticUpdates.ts

import { QueryClient } from '@tanstack/react-query';
import { AssetRequest } from '@/types/asset-request';

/**
 * Utility functions for optimistic updates with rollback support
 */

interface OptimisticUpdateContext<T> {
  previousData: T | undefined;
  queryKey: unknown[];
}

/**
 * Perform an optimistic update on a query
 * @param queryClient - React Query client
 * @param queryKey - Query key to update
 * @param updater - Function to update the data
 * @returns Context for rollback
 */
export const performOptimisticUpdate = <T>(
  queryClient: QueryClient,
  queryKey: unknown[],
  updater: (oldData: T | undefined) => T
): OptimisticUpdateContext<T> => {
  // Cancel any outgoing refetches
  queryClient.cancelQueries({ queryKey });

  // Snapshot the previous value
  const previousData = queryClient.getQueryData<T>(queryKey);

  // Optimistically update to the new value
  queryClient.setQueryData<T>(queryKey, updater);

  return { previousData, queryKey };
};

/**
 * Rollback an optimistic update
 * @param queryClient - React Query client
 * @param context - Context from performOptimisticUpdate
 */
export const rollbackOptimisticUpdate = <T>(
  queryClient: QueryClient,
  context: OptimisticUpdateContext<T>
): void => {
  queryClient.setQueryData(context.queryKey, context.previousData);
};

/**
 * Optimistically update a request status
 * @param queryClient - React Query client
 * @param requestId - Request ID to update
 * @param newStatus - New status
 * @param additionalUpdates - Additional fields to update
 * @returns Context for rollback
 */
export const optimisticallyUpdateRequestStatus = (
  queryClient: QueryClient,
  requestId: number,
  newStatus: AssetRequest['status'],
  additionalUpdates?: Partial<AssetRequest>
): OptimisticUpdateContext<AssetRequest[]>[] => {
  const contexts: OptimisticUpdateContext<AssetRequest[]>[] = [];

  // Update all relevant query keys
  const queryKeys = [
    ['assetRequests', 'pending'],
    ['assetRequests', 'all'],
  ];

  queryKeys.forEach((queryKey) => {
    const context = performOptimisticUpdate<AssetRequest[]>(
      queryClient,
      queryKey,
      (oldData) => {
        if (!oldData) return [];
        
        return oldData.map((request) =>
          request.requestID === requestId
            ? {
                ...request,
                status: newStatus,
                ...additionalUpdates,
              }
            : request
        );
      }
    );
    contexts.push(context);
  });

  return contexts;
};

/**
 * Optimistically add a new request
 * @param queryClient - React Query client
 * @param newRequest - New request to add
 * @param employeeId - Employee ID for employee-specific queries
 * @returns Context for rollback
 */
export const optimisticallyAddRequest = (
  queryClient: QueryClient,
  newRequest: AssetRequest,
  employeeId: number
): OptimisticUpdateContext<AssetRequest[]>[] => {
  const contexts: OptimisticUpdateContext<AssetRequest[]>[] = [];

  // Update employee-specific query
  const employeeContext = performOptimisticUpdate<AssetRequest[]>(
    queryClient,
    ['assetRequests', 'employee', employeeId],
    (oldData) => {
      if (!oldData) return [newRequest];
      return [newRequest, ...oldData];
    }
  );
  contexts.push(employeeContext);

  // Update all requests query
  const allContext = performOptimisticUpdate<AssetRequest[]>(
    queryClient,
    ['assetRequests', 'all'],
    (oldData) => {
      if (!oldData) return [newRequest];
      return [newRequest, ...oldData];
    }
  );
  contexts.push(allContext);

  // Update pending requests query
  if (newRequest.status === 'pending') {
    const pendingContext = performOptimisticUpdate<AssetRequest[]>(
      queryClient,
      ['assetRequests', 'pending'],
      (oldData) => {
        if (!oldData) return [newRequest];
        return [newRequest, ...oldData];
      }
    );
    contexts.push(pendingContext);
  }

  return contexts;
};

/**
 * Optimistically remove a request from a specific status list
 * @param queryClient - React Query client
 * @param requestId - Request ID to remove
 * @param fromStatus - Status list to remove from
 * @returns Context for rollback
 */
export const optimisticallyRemoveFromStatus = (
  queryClient: QueryClient,
  requestId: number,
  fromStatus: 'pending' | 'approved'
): OptimisticUpdateContext<AssetRequest[]> => {
  const queryKey = ['assetRequests', fromStatus];
  
  return performOptimisticUpdate<AssetRequest[]>(
    queryClient,
    queryKey,
    (oldData) => {
      if (!oldData) return [];
      return oldData.filter((request) => request.requestID !== requestId);
    }
  );
};

/**
 * Rollback multiple optimistic updates
 * @param queryClient - React Query client
 * @param contexts - Array of contexts to rollback
 */
export const rollbackMultipleUpdates = <T>(
  queryClient: QueryClient,
  contexts: OptimisticUpdateContext<T>[]
): void => {
  contexts.forEach((context) => {
    rollbackOptimisticUpdate(queryClient, context);
  });
};