import React from 'react'
import { QueryClient, QueryClientProvider, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { queryDocument, ingestDocument, listDocuments, listChunks, deleteDocument, deleteChunk } from './api'
import type { QueryRequest, QueryResponse, IngestResponse, DocumentListResponse, ChunkInfo, DeleteResponse } from '../types'

export const queryClient = new QueryClient()

export const useQueryDocument = () => {
  return useMutation<QueryResponse, Error, QueryRequest>({
    mutationFn: queryDocument,
  })
}

export const useIngestDocument = () => {
  return useMutation<IngestResponse, Error, File>({
    mutationFn: ingestDocument,
  })
}

export const useDocuments = () => {
  return useQuery<DocumentListResponse, Error>({
    queryKey: ['documents'],
    queryFn: listDocuments,
  })
}

export const useDocumentChunks = (documentId: string | null) => {
  return useQuery<ChunkInfo[], Error>({
    queryKey: ['documents', documentId, 'chunks'],
    queryFn: () => listChunks(documentId!),
    enabled: documentId !== null,
  })
}

export const useDeleteDocument = () => {
  const queryClient = useQueryClient()
  return useMutation<DeleteResponse, Error, string>({
    mutationFn: deleteDocument,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['documents'] })
    },
  })
}

export const useDeleteChunk = () => {
  const queryClient = useQueryClient()
  return useMutation<DeleteResponse, Error, string>({
    mutationFn: deleteChunk,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['documents'] })
    },
  })
}

export const AppQueryProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
  return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
}