/**
 * Hook for managing augmentation operations.
 *
 * Provides functions for fetching augmentation types, presets, and previewing augmentations.
 */

import { useQuery, useMutation } from '@tanstack/react-query'
import {
  augmentationApi,
  type AugmentationTypesResponse,
  type PresetsResponse,
  type PreviewResponse,
  type BatchRequest,
  type BatchResponse,
  type AugmentationConfig,
} from '../api/endpoints/augmentation'

interface PreviewParams {
  documentId: string
  augmentationType: string
  params: Record<string, unknown>
  page?: number
}

interface PreviewConfigParams {
  documentId: string
  config: AugmentationConfig
  page?: number
}

export const useAugmentation = () => {
  // Fetch augmentation types
  const {
    data: typesData,
    isLoading: isLoadingTypes,
    error: typesError,
  } = useQuery<AugmentationTypesResponse>({
    queryKey: ['augmentation', 'types'],
    queryFn: () => augmentationApi.getTypes(),
    staleTime: 5 * 60 * 1000, // Cache for 5 minutes
  })

  // Fetch presets
  const {
    data: presetsData,
    isLoading: isLoadingPresets,
    error: presetsError,
  } = useQuery<PresetsResponse>({
    queryKey: ['augmentation', 'presets'],
    queryFn: () => augmentationApi.getPresets(),
    staleTime: 5 * 60 * 1000,
  })

  // Preview single augmentation mutation
  const previewMutation = useMutation<PreviewResponse, Error, PreviewParams>({
    mutationFn: ({ documentId, augmentationType, params, page = 1 }) =>
      augmentationApi.preview(
        documentId,
        { augmentation_type: augmentationType, params },
        page
      ),
    onError: (error) => {
      console.error('Preview augmentation failed:', error)
    },
  })

  // Preview full config mutation
  const previewConfigMutation = useMutation<PreviewResponse, Error, PreviewConfigParams>({
    mutationFn: ({ documentId, config, page = 1 }) =>
      augmentationApi.previewConfig(documentId, config, page),
    onError: (error) => {
      console.error('Preview config failed:', error)
    },
  })

  // Create augmented dataset mutation
  const createBatchMutation = useMutation<BatchResponse, Error, BatchRequest>({
    mutationFn: (request) => augmentationApi.createBatch(request),
    onError: (error) => {
      console.error('Create augmented dataset failed:', error)
    },
  })

  return {
    // Types data
    augmentationTypes: typesData?.augmentation_types || [],
    isLoadingTypes,
    typesError,

    // Presets data
    presets: presetsData?.presets || [],
    isLoadingPresets,
    presetsError,

    // Preview single augmentation
    preview: previewMutation.mutate,
    previewAsync: previewMutation.mutateAsync,
    isPreviewing: previewMutation.isPending,
    previewData: previewMutation.data,
    previewError: previewMutation.error,

    // Preview full config
    previewConfig: previewConfigMutation.mutate,
    previewConfigAsync: previewConfigMutation.mutateAsync,
    isPreviewingConfig: previewConfigMutation.isPending,
    previewConfigData: previewConfigMutation.data,
    previewConfigError: previewConfigMutation.error,

    // Create batch
    createBatch: createBatchMutation.mutate,
    createBatchAsync: createBatchMutation.mutateAsync,
    isCreatingBatch: createBatchMutation.isPending,
    batchData: createBatchMutation.data,
    batchError: createBatchMutation.error,

    // Reset functions for clearing stale mutation state
    resetPreview: previewMutation.reset,
    resetPreviewConfig: previewConfigMutation.reset,
    resetBatch: createBatchMutation.reset,
  }
}
