import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
import { trainingApi } from '../api/endpoints'
import type { TrainingModelsResponse } from '../api/types'

export const useTraining = () => {
  const queryClient = useQueryClient()

  const { data: modelsData, isLoading: isLoadingModels } =
    useQuery<TrainingModelsResponse>({
      queryKey: ['training', 'models'],
      queryFn: () => trainingApi.getModels(),
      staleTime: 30000,
    })

  const startTrainingMutation = useMutation({
    mutationFn: (config: {
      name: string
      description?: string
      document_ids: string[]
      epochs?: number
      batch_size?: number
      model_base?: string
    }) => trainingApi.startTraining(config),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
    },
  })

  const cancelTaskMutation = useMutation({
    mutationFn: (taskId: string) => trainingApi.cancelTask(taskId),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
    },
  })

  const downloadModelMutation = useMutation({
    mutationFn: (taskId: string) => trainingApi.downloadModel(taskId),
    onSuccess: (blob, taskId) => {
      const url = window.URL.createObjectURL(blob)
      const a = document.createElement('a')
      a.href = url
      a.download = `model-${taskId}.pt`
      document.body.appendChild(a)
      a.click()
      window.URL.revokeObjectURL(url)
      document.body.removeChild(a)
    },
  })

  return {
    models: modelsData?.models || [],
    total: modelsData?.total || 0,
    isLoadingModels,
    startTraining: startTrainingMutation.mutate,
    startTrainingAsync: startTrainingMutation.mutateAsync,
    isStartingTraining: startTrainingMutation.isPending,
    cancelTask: cancelTaskMutation.mutate,
    isCancelling: cancelTaskMutation.isPending,
    downloadModel: downloadModelMutation.mutate,
    isDownloading: downloadModelMutation.isPending,
  }
}

export const useTrainingDocuments = (params?: {
  has_annotations?: boolean
  min_annotation_count?: number
  exclude_used_in_training?: boolean
  limit?: number
  offset?: number
}) => {
  const { data, isLoading, error } = useQuery({
    queryKey: ['training', 'documents', params],
    queryFn: () => trainingApi.getDocumentsForTraining(params),
    staleTime: 30000,
  })

  return {
    documents: data?.documents || [],
    total: data?.total || 0,
    isLoading,
    error,
  }
}
