import { QueryClient } from '@tanstack/react-query';

import logWebUserEvent from '@/logging/logWebUserEvent';

import type {
  BatchIdWithType,
  HistoryBatchItem,
  HistoryEntity,
} from './generationHistoryTypes';

export const generationQueryKeys = {
  pendingBatchIds: () => ['pending-batch-ids'] as const,
  history: (clipId: string) => ['generation-history', clipId] as const,
  videoGenerationFavorites: () => ['video-generation-favorites'] as const,
  imageBatch: (batchId: string) => ['image-batch', batchId] as const,
  pollBatches: (batchIds: Array<{ id: string; type: string }>) =>
    ['poll-batches', batchIds] as const,
  clipCover: (clipId: string) => ['clip-cover', clipId] as const,
  pollMultipleClips: (batches: { id: string; type: string }[]) =>
    ['generation', 'poll-multiple', batches] as const,
};

/**
 * Log a generation completion event.
 */
function logGenerationCompleted(
  clipId: string,
  batchId: string,
  item: HistoryBatchItem
) {
  logWebUserEvent({
    actionName: 'GenerateCoverArtGenerationCompleted',
    principalObjectType: 'clip',
    principalObjectValue: clipId,
    context: {
      batchId: batchId,
      generationType: item.type === 'image' ? 'image' : 'video',
      videoId: item.type === 'video' ? item.id : undefined,
      imageId: item.type === 'image' ? item.id : undefined,
    },
  });
}

/**
 * Log a generation failure event.
 */
function logGenerationFailed(
  clipId: string,
  batchId: string,
  item: HistoryBatchItem
) {
  logWebUserEvent({
    actionName: 'GenerateCoverArtGenerationFailed',
    principalObjectType: 'clip',
    principalObjectValue: clipId,
    context: {
      batchId: batchId,
      generationType: item.type === 'image' ? 'image' : 'video',
      videoId: item.type === 'video' ? item.id : undefined,
      imageId: item.type === 'image' ? item.id : undefined,
    },
  });
}

/**
 * Update the generation history cache for a specific clip with completed batch data.
 * This function should be called after converting the poll response from snake_case to camelCase.
 *
 * @param queryClient - TanStack Query client instance
 * @param clipId - The clip ID whose history should be updated
 * @param completedBatches - Map of batchId to array of completed items (already converted to camelCase)
 */
export function updateGenerationHistoryCache(
  queryClient: QueryClient,
  clipId: string,
  completedBatches: Record<string, HistoryBatchItem[]>
) {
  const completedBatchIds = Object.keys(completedBatches);
  let failedGeneration = false;

  if (completedBatchIds.length === 0) return;

  queryClient.setQueryData(
    generationQueryKeys.history(clipId),
    (oldHistory: HistoryEntity[] | undefined) => {
      const oldBatchIds = oldHistory?.map((entity) => entity.batchId) ?? [];

      // Log events for batches that weren't in old history
      for (const [batchId, items] of Object.entries(completedBatches)) {
        if (!oldBatchIds.includes(batchId)) {
          for (const item of items) {
            if (item.status === 'complete') {
              logGenerationCompleted(clipId, batchId, item);
            } else if (item.status === 'error') {
              logGenerationFailed(clipId, batchId, item);
              failedGeneration = true;
            }
          }
        }
      }

      if (!oldHistory) return oldHistory;

      return oldHistory.map((entity) => {
        if (completedBatchIds.includes(entity.batchId)) {
          // only log if the item is new or the old item's status was processing
          // this will prevent duplicate logs for items that have already reached
          // a termainal status
          for (const item of completedBatches[entity.batchId]) {
            const oldItem = entity.items.find(
              (oldItem) => oldItem.id === item.id
            );
            const shouldLog = !oldItem || oldItem.status === 'processing';
            if (shouldLog && item.status === 'complete') {
              logGenerationCompleted(clipId, entity.batchId, item);
            } else if (shouldLog && item.status === 'error') {
              logGenerationFailed(clipId, entity.batchId, item);
              failedGeneration = true;
            }
          }
          const updatedItems = completedBatches[entity.batchId] || [];
          const allComplete = updatedItems.every(
            (item: HistoryBatchItem) => item.status !== 'processing'
          );
          return {
            ...entity,
            items: updatedItems,
            status: allComplete ? 'complete' : 'processing',
          };
        }
        return entity;
      });
    }
  );
  if (failedGeneration) {
    // if any generations failed, refetch subscription info to update credits
    queryClient.invalidateQueries({ queryKey: ['subscriptionInfo'] });
  }
}

/**
 * Remove completed generations from the pending batch IDs list.
 *
 * @param queryClient - TanStack Query client instance
 * @param completedBatches - Map of batchId to array of completed items
 */
export function removeCompletedGenerations(
  queryClient: QueryClient,
  completedBatches: Record<string, HistoryBatchItem[]>
) {
  const completedBatchIds = Object.keys(completedBatches);

  if (completedBatchIds.length === 0) return;

  queryClient.setQueryData(
    generationQueryKeys.pendingBatchIds(),
    (oldPendingBatchIds: BatchIdWithType[] | undefined) => {
      if (!oldPendingBatchIds) return oldPendingBatchIds;

      return oldPendingBatchIds.filter(
        (pendingBatchId) => !completedBatchIds.includes(pendingBatchId.id)
      );
    }
  );
}