// In useGenerationPolling.ts
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { useEffect, useMemo } from 'react';
import { deepCamelKeys } from 'string-ts';

import { useApiClient } from '@/lib/apiClient';

import { filterBatchesByCompletionStatus } from '../../../components/modal/editClipMetadata/batchFiltering';
import type {
  BatchIdWithType,
  HistoryBatchItem,
} from '../../../components/modal/editClipMetadata/generationHistoryTypes';
import {
  generationQueryKeys,
  removeCompletedGenerations,
  updateGenerationHistoryCache,
} from '../../../components/modal/editClipMetadata/queryKeys';

const POLLING_INTERVAL = 7000; // 7 seconds

/**
 * Hook to poll for generation completion across multiple clips
 *
 * @param enabled - Whether polling should be active
 * @returns Query with completed batches organized by clip and batch ID
 */
export function useGenerationPolling(enabled: boolean = true) {
  const apiClient = useApiClient();
  const queryClient = useQueryClient();

  const pendingBatchIdsQuery = useQuery({
    queryKey: generationQueryKeys.pendingBatchIds(),
    queryFn: async () => {
      const { data, error } = await apiClient.POST(
        '/api/video_gen/pending_batches',
        {
          body: {},
        }
      );

      if (error || !data) {
        throw new Error('Failed to get pending batch IDs');
      }

      return data.batch_ids;
    },
    enabled: enabled,
  });

  const pendingBatchIds =
    pendingBatchIdsQuery?.data?.filter((batch) => batch.type === 'video') ??
    ([] as BatchIdWithType[]);

  const pollBatchesQuery = useQuery({
    queryKey: generationQueryKeys.pollMultipleClips(pendingBatchIds),
    queryFn: async () => {
      if (!pendingBatchIds.length) return {};

      const { data, error } = await apiClient.POST(
        '/api/video_gen/poll_batches',
        {
          body: {
            batch_ids: pendingBatchIds,
          },
        }
      );

      if (error || !data) {
        throw new Error('Failed to poll batches');
      }

      const convertedBatches: Record<
        string,
        Record<string, HistoryBatchItem[]>
      > = {};

      for (const [batchId, items] of Object.entries(data.batches)) {
        if (items.length === 0) {
          continue;
        }
        const clipId = items[0].clip_id;
        if (!clipId) {
          continue;
        }
        if (!(clipId in convertedBatches)) {
          convertedBatches[clipId] = {};
        }
        convertedBatches[clipId][batchId] = deepCamelKeys(
          items
        ) as HistoryBatchItem[];
      }
      return convertedBatches;
    },
    enabled: enabled && pendingBatchIds.length > 0,
    refetchInterval: POLLING_INTERVAL,
    staleTime: 0,
  });

  const { finishedClipBatches = {}, someCompletedClipBatches = {} } =
    useMemo(() => {
      if (!pollBatchesQuery.data) return {};
      const finishedClipBatches: Record<
        string,
        Record<string, HistoryBatchItem[]>
      > = {};
      const someCompletedClipBatches: Record<
        string,
        Record<string, HistoryBatchItem[]>
      > = {};
      for (const [clipId, batches] of Object.entries(pollBatchesQuery.data)) {
        const { someCompletedBatches, completedBatches } =
          filterBatchesByCompletionStatus(batches);
        if (Object.keys(completedBatches).length > 0) {
          finishedClipBatches[clipId] = completedBatches;
        }
        if (Object.keys(someCompletedBatches).length > 0) {
          someCompletedClipBatches[clipId] = someCompletedBatches;
        }
      }
      return { finishedClipBatches, someCompletedClipBatches };
    }, [pollBatchesQuery.data]);

  useEffect(() => {
    for (const [clipId] of Object.entries(finishedClipBatches)) {
      // Remove completed batch IDs from pending list
      removeCompletedGenerations(
        queryClient,
        finishedClipBatches[clipId] ?? {}
      );
    }
    for (const [clipId] of Object.entries(someCompletedClipBatches)) {
      // Update the generation history cache for this clip with completed batches
      updateGenerationHistoryCache(
        queryClient,
        clipId,
        someCompletedClipBatches[clipId] ?? {}
      );
    }
  }, [someCompletedClipBatches, finishedClipBatches, queryClient]);

  return {
    finishedClipBatches,
    pendingBatchIds,

    pollBatchesQuery: pollBatchesQuery,
    pendingBatchIdsQuery: pendingBatchIdsQuery,
  };
}