import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
import { datasetsApi } from '../api/endpoints'
import type {
  DatasetCreateRequest,
  DatasetDetailResponse,
  DatasetListResponse,
  DatasetTrainRequest,
} from '../api/types'

export const useDatasets = (params?: {
  status?: string
  limit?: number
  offset?: number
}) => {
  const queryClient = useQueryClient()

  const { data, isLoading, error, refetch } = useQuery<DatasetListResponse>({
    queryKey: ['datasets', params],
    queryFn: () => datasetsApi.list(params),
    staleTime: 30000,
    // Poll every 5 seconds when there's an active training task
    refetchInterval: (query) => {
      const datasets = query.state.data?.datasets ?? []
      const hasActiveTraining = datasets.some(
        d => d.training_status === 'running' || d.training_status === 'pending' || d.training_status === 'scheduled'
      )
      return hasActiveTraining ? 5000 : false
    },
  })

  const createMutation = useMutation({
    mutationFn: (req: DatasetCreateRequest) => datasetsApi.create(req),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['datasets'] })
    },
  })

  const deleteMutation = useMutation({
    mutationFn: (datasetId: string) => datasetsApi.remove(datasetId),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['datasets'] })
    },
  })

  const trainMutation = useMutation({
    mutationFn: ({ datasetId, req }: { datasetId: string; req: DatasetTrainRequest }) =>
      datasetsApi.trainFromDataset(datasetId, req),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['datasets'] })
      queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
    },
  })

  return {
    datasets: data?.datasets ?? [],
    total: data?.total ?? 0,
    isLoading,
    error,
    refetch,
    createDataset: createMutation.mutate,
    createDatasetAsync: createMutation.mutateAsync,
    isCreating: createMutation.isPending,
    deleteDataset: deleteMutation.mutate,
    isDeleting: deleteMutation.isPending,
    trainFromDataset: trainMutation.mutate,
    trainFromDatasetAsync: trainMutation.mutateAsync,
    isTraining: trainMutation.isPending,
  }
}

export const useDatasetDetail = (datasetId: string | null) => {
  const { data, isLoading, error } = useQuery<DatasetDetailResponse>({
    queryKey: ['datasets', datasetId],
    queryFn: () => datasetsApi.getDetail(datasetId!),
    enabled: !!datasetId,
    staleTime: 30000,
  })

  return {
    dataset: data ?? null,
    isLoading,
    error,
  }
}