import { useMutation, useQuery } from "@tanstack/react-query";
import {
  activateDataset,
  addFileToDataset,
  createDataset,
  createDatasetsDraft,
  deleteDataset,
  deleteFileFromDataset,
  getDataset,
  getDatasets,
  getProcessing,
} from "./documentsApi";
import { GetDatasetRequestParams } from "./types";
import { queryClient } from "@/main";

export const useGetDatasets = () => {
  return useQuery({
    queryKey: ["datasets"],
    queryFn: async () => {
      const response = await getDatasets();
      return response;
    },
  });
};

export const useCreateDataset = () => {
  return useMutation({
    mutationFn: async (data: FormData) => {
      const response = await createDataset(data);
      return response;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
    },
  });
};

export const useActivateDataset = () => {
  return useMutation({
    mutationFn: async (id: number) => {
      const response = await activateDataset(id);
      return response;
    },
    onSuccess: (data) => {
      queryClient.invalidateQueries({ queryKey: ["processing"] });
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
      queryClient.invalidateQueries({ queryKey: ["dataset", data.id] });
    },
  });
};

export const useGetDataset = (id: number, params: GetDatasetRequestParams) => {
  return useQuery({
    queryKey: ["dataset", id, params],
    queryFn: async () => {
      const response = await getDataset(id, params);
      return response;
    },
  });
};

export const useCreateDatasetsDraft = () => {
  return useMutation({
    mutationFn: async (id: number) => {
      const response = await createDatasetsDraft(id);
      return response;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
      queryClient.invalidateQueries({ queryKey: ["dataset"] });
    },
  });
};

export const useDeleteDataset = () => {
  return useMutation({
    mutationFn: async (id: number) => {
      const response = await deleteDataset(id);
      return response;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
    },
  });
};

export const useAddFileToDataset = () => {
  return useMutation({
    mutationFn: async ({ id, data }: { id: number; data: FormData }) => {
      const response = await addFileToDataset(id, data);
      return response;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
      queryClient.invalidateQueries({ queryKey: ["dataset"] });
    },
  });
};

export const useDeleteFileFromDataset = () => {
  return useMutation({
    mutationFn: async ({ datasetId, fileId }: { datasetId: number; fileId: number }) => {
      const response = await deleteFileFromDataset(datasetId, fileId);
      return response;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["datasets"] });
      queryClient.invalidateQueries({ queryKey: ["dataset"] });
    },
  });
};

export const useGetProcessing = () => {
  return useQuery({
    queryKey: ["processing"],
    queryFn: async () => {
      const response = await getProcessing();
      return response;
    },
  });
};