import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { useCallback, useEffect, useState } from 'react';

import { useApiClient } from '@/lib/apiClient';
import type { components } from '@/lib/gen';

import { generationQueryKeys } from './queryKeys';
import { useGenerateTextPrompt } from './useGenerateTextPrompt';
import { useStartImageGeneration } from './useStartImageGeneration';

// Helper function to extract S3 ID from image URL
function extractS3IdFromUrl(url: string | null): string | null {
  if (!url) return null;

  try {
    const filename = new URL(url).pathname.split('/').pop() ?? '';
    const dotIndex = filename.lastIndexOf('.');
    return dotIndex > 0 ? filename.substring(0, dotIndex) : filename;
  } catch {
    return null;
  }
}

interface ImageGenerationOptions {
  maxRetries?: number;
  timeout?: number;
  pollingInterval?: number;
  clipId?: string; // needed to fetch existing cover art
  shouldUseCover?: boolean; // controls whether to show/use cover image data
}

interface UploadImageResult {
  uploadId: string;
  imageUrl: string;
}

interface GenerateFromSongParams {
  clipId: string;
  quantity?: number;
  models?: ('sdxl-lightning' | 'openai_image_gpt_1' | 'fal_nano_banana')[];
}

type ImageBatchResponse = components['schemas']['ImageBatchResponse'];

export function useImageGenerationPolling(
  options: ImageGenerationOptions = {}
) {
  const {
    maxRetries = 10,
    pollingInterval = 3000, // 3 seconds
    clipId,
    shouldUseCover = false,
  } = options;

  const apiClient = useApiClient();
  const [batchId, setBatchId] = useState<string | null>(null);
  const [retryCount, setRetryCount] = useState(0);
  const queryClient = useQueryClient();

  // Query to fetch existing cover art metadata
  // Fetches automatically when clipId is available, caches the result
  const existingCoverQuery = useQuery({
    queryKey: generationQueryKeys.clipCover(clipId || ''),
    queryFn: async () => {
      if (!clipId) {
        throw new Error('Clip ID is required');
      }

      const { data, error } = await apiClient.GET('/api/clip/{clip_id}', {
        params: {
          path: { clip_id: clipId },
        },
      });

      if (error || !data) {
        throw new Error('Failed to fetch clip cover art');
      }

      return data;
    },
    enabled: !!clipId, // Fetch automatically when clipId is available
    staleTime: Infinity, // Don't refetch - cover art doesn't change during modal session
  });

  const { error: existingCoverError, refetch: refetchExistingCover } =
    existingCoverQuery;

  // Extract cover data - only return values if shouldUseCover is true
  // This preserves the cached data without exposing it when we don't want to use it
  const coverImageUrl = shouldUseCover
    ? (existingCoverQuery.data?.image_url ?? null)
    : null;

  // Extract S3 ID from the URL path (last segment, removing file extension if present)
  const coverImageId = shouldUseCover
    ? extractS3IdFromUrl(coverImageUrl)
    : null;

  // Get mutation functions from existing hooks
  const generateTextPrompt = useGenerateTextPrompt();
  const startImageGeneration = useStartImageGeneration();

  // Mutation to generate from song (combines text generation + image generation)
  const generateFromSongMutation = useMutation({
    mutationFn: async (params: GenerateFromSongParams) => {
      // Step 1: Generate text prompt
      const textData = await generateTextPrompt({
        clipId: params.clipId,
        target: 'image',
        prompt: null,
      });

      // Step 2: Start image generation
      const imageData = await startImageGeneration({
        generatedTextId: textData.id,
        prompt: textData.prompt,
        quantity: params.quantity ?? 2,
        models: params.models ?? ['sdxl-lightning'],
      });

      return {
        batchId: imageData.batch_id,
      };
    },
    onSuccess: (data) => {
      // Cancel any existing polling
      if (batchId) {
        queryClient.cancelQueries({
          queryKey: generationQueryKeys.imageBatch(batchId),
        });
      }
      setBatchId(data.batchId);
      setRetryCount(0);
    },
  });

  // Mutation to upload image from device
  const uploadImageMutation = useMutation({
    mutationFn: async (file: File): Promise<UploadImageResult> => {
      // Step 1: Get upload parameters from /api/uploads/image/
      const extension = file.name.split('.').pop() || 'jpeg';

      const { data: uploadData, error: uploadError } = await apiClient.POST(
        '/api/uploads/image/',
        {
          body: {
            extension,
          },
        }
      );

      if (!uploadData || uploadError) {
        throw new Error('Failed to upload');
      }

      const uploadId = uploadData.id;
      const uploadUrl = uploadData.url;
      const uploadFields = uploadData.fields as Record<string, string>;

      // Step 2: Upload file to S3 using form data
      const formData = new FormData();
      Object.entries(uploadFields).forEach(([key, value]) => {
        formData.append(key, value);
      });
      formData.append('file', file);

      const s3Response = await fetch(uploadUrl, {
        method: 'POST',
        body: formData,
      });

      if (!s3Response.ok) {
        throw new Error('Failed to upload file to S3');
      }

      // Step 3: Call upload-finish to moderate and move to permanent bucket
      const { data: finishData, error: finishError } = await apiClient.POST(
        '/api/uploads/image/{upload_id}/upload-finish/',
        {
          params: {
            path: { upload_id: uploadId },
          },
        }
      );

      if (!finishData || finishError) {
        throw new Error('Failed to finish image upload');
      }

      if (finishData.moderation_status === 'rejected') {
        throw new Error('Image was rejected by moderation');
      }

      // Construct CDN URL from upload ID
      const imageUrl = `https://cdn1.suno.ai/image_${uploadId}.${extension}`;

      return {
        uploadId,
        imageUrl,
      };
    },
  });

  // Query to poll image generation status
  const imageBatchQuery = useQuery<ImageBatchResponse>({
    queryKey: generationQueryKeys.imageBatch(batchId || ''),
    queryFn: async () => {
      if (!batchId) throw new Error('No batch ID');

      const { data, error } = await apiClient.GET(
        '/api/video_gen/image/generate/{batch_id}',
        {
          params: {
            path: { batch_id: batchId },
          },
        }
      );

      if (error || !data) {
        throw new Error('Failed to fetch image generation status');
      }

      return data;
    },
    enabled: !!batchId,
    refetchInterval: (query) => {
      const data = query.state.data;

      if (data?.status === 'complete' || data?.status === 'error') {
        return false;
      }

      if (retryCount >= maxRetries) {
        return false;
      }

      return pollingInterval;
    },
    staleTime: 0,
  });

  // Track retry count
  useEffect(() => {
    if (imageBatchQuery.dataUpdatedAt > 0) {
      setRetryCount((prev) => prev + 1);
    }
  }, [imageBatchQuery.dataUpdatedAt]);

  const isComplete =
    imageBatchQuery.data?.status === 'complete' ||
    imageBatchQuery.data?.status === 'error';

  const hasTimedOut = retryCount >= maxRetries;

  // Get the first image from the batch (AI generated)
  const firstImage = imageBatchQuery.data?.images?.[0];
  const imageUrl = firstImage?.image_url || null;
  const imageId = firstImage?.id || null;

  // Get uploaded image data
  const uploadedImageUrl = uploadImageMutation.data?.imageUrl || null;
  const uploadedImageId = uploadImageMutation.data?.uploadId || null;

  // Determine loading state
  // Only include cover query loading when we're actually using the cover
  const isLoading =
    generateFromSongMutation.isPending ||
    uploadImageMutation.isPending ||
    (!!batchId && !isComplete && !hasTimedOut) ||
    (shouldUseCover && existingCoverQuery.isFetching);

  // Determine error message with priority order
  let error: string | null = null;

  if (imageBatchQuery.data?.status === 'error') {
    error = 'Failed to generate';
  } else if (hasTimedOut) {
    error = 'Failed to generate';
  } else if (imageBatchQuery.error) {
    error = 'Failed to generate';
  } else if (generateFromSongMutation.error) {
    error = 'Failed to generate';
  } else if (shouldUseCover && existingCoverQuery.error) {
    error =
      existingCoverQuery.error instanceof Error
        ? existingCoverQuery.error.message
        : 'Failed to fetch cover art';
  } else if (uploadImageMutation.error) {
    error =
      uploadImageMutation.error instanceof Error
        ? uploadImageMutation.error.message
        : 'Failed to upload image';
  }

  const reset = useCallback(() => {
    setBatchId(null);
    setRetryCount(0);
  }, []);

  const resetUploadedImage = useCallback(() => {
    uploadImageMutation.reset();
  }, [uploadImageMutation]);

  return {
    // AI generation
    startGenerationFromSongWithPolling: generateFromSongMutation.mutateAsync,
    imageUrl,
    imageId,

    // Cover image
    coverImageUrl,
    coverImageId,
    existingCoverError,
    refetchExistingCover,
    reset,

    // Image upload
    uploadImage: uploadImageMutation.mutateAsync,
    uploadedImageUrl,
    uploadedImageId,
    resetUploadedImage,

    // Shared state
    isLoading,
    error,
    existingCoverQuery, // Expose for refetch() if needed
    imageBatchQuery,
    generateFromSongMutation,
    uploadImageMutation,
  };
}