import { useQueries, useQueryClient } from '@tanstack/react-query';

import { useStores } from '@/app/(root)/AppProviders';
import { toast } from '@/components/toast/Toast';
import { components } from '@/lib/gen';
import { sleep } from '@/utils/utils';

export type MidiInstrumentSchema =
  components['schemas']['MidiInstrumentSchema'];

export type MidiNoteSchema = components['schemas']['MidiNoteSchema'];

export const getMidiForClip = async (
  apiClient: any,
  clipId: string
): Promise<MidiInstrumentSchema[]> => {
  const fetchMidi = () => {
    return apiClient.GET('/api/gen/{clip_id}/midi', {
      params: { path: { clip_id: clipId } },
    });
  };

  const pollMidiFetch = async (
    firstRequestTime: number = Date.now()
  ): Promise<MidiInstrumentSchema[]> => {
    const result = await fetchMidi();

    // Check for permission error (403)
    if (result.error && result.response?.status === 403) {
      toast({
        title: 'No permission to get MIDI',
        description: 'MIDI transcription requires a Premier subscription',
        status: 'error',
        duration: 5000,
        isClosable: true,
      });
      throw new Error(
        'Permission denied: MIDI transcription requires a Premier subscription'
      );
    }

    if (Date.now() - firstRequestTime > 1000 * 60 * 5) {
      // stop trying after 5 minutes
      throw new Error('Timeout');
    }
    if (result.data?.state === 'running') {
      await sleep(2500);
      return pollMidiFetch(firstRequestTime);
    } else if (result.data?.state === 'complete') {
      return result.data.instruments!;
    } else {
      console.error(result.error);
      throw new Error(result.error || 'Unknown error');
    }
  };

  return await pollMidiFetch();
};

export const useMidiTranscripts = (clipIds: (string | null)[]) => {
  const { session } = useStores();
  const queryClient = useQueryClient();

  // Use useQueries to handle dynamic number of queries safely
  const queries = useQueries({
    queries: clipIds.map((clipId) => ({
      queryKey: ['midiTranscript', clipId],
      queryFn: () => getMidiForClip(session.apiClient, clipId!),
      enabled: !!clipId,
      staleTime: 5 * 60 * 1000, // 5 minutes
      gcTime: 10 * 60 * 1000, // 10 minutes
      retry: 2,
      refetchOnWindowFocus: false,
    })),
  });

  // Return an array of objects with the same structure as useMidiTranscript
  return clipIds.map((clipId, index) => {
    const query = queries[index];
    return {
      // Data and loading states
      midiTranscript: query.data || null,
      isLoading: query.isLoading,
      isPending: query.isPending,
      isError: query.isError,
      error: query.error,

      // Utility methods
      hasMidiTranscript:
        !!query.data && query.data.some((x) => x.notes.length > 0),
      isFetching: query.isFetching,

      // Control methods
      refetch: query.refetch,
      remove: () =>
        queryClient.removeQueries({ queryKey: ['midiTranscript', clipId] }),
    };
  });
};