src/api/predictions/hooks.ts

import { useMutation, UseMutationResult, useQuery, useQueryClient } from "@tanstack/react-query";

import { useRef } from "react";
import {
  LlmResponseType,
  LlmRequestType,
  ChunksResponseType,
  ChunksArrayItem,
  ChunkRequestType,
  FeedbackRequestType,
  LogRequestType,
} from "./types";
import { getLlmResponse, getChunks, postFeedback, getLogs, getAbbreviations } from "./predictionsApi";

export function useGetLlmResponse() {
  const queryClient = useQueryClient();
  const abortControllerRef = useRef<AbortController | null>(null);

  const cancelQuery = () => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      abortControllerRef.current = null;
    }
    queryClient.cancelQueries({ queryKey: ["llm"] });
  };

  const mutation: UseMutationResult<LlmResponseType | undefined, Error, LlmRequestType, unknown> = useMutation({
    mutationFn: async (data: LlmRequestType) => {
      // Отменить предыдущий запрос
      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
      }

      const abortController = new AbortController();
      abortControllerRef.current = abortController;

      try {
        const response = await getLlmResponse(data, abortController.signal);
        return response;
      } catch (error) {
        if (error instanceof Error && error.name === "AbortError") {
          console.log("Запрос был отменен");
          return undefined;
        } else {
          throw error;
        }
      } finally {
        if (abortControllerRef.current === abortController) {
          abortControllerRef.current = null;
        }
      }
    },
    onSuccess: () => {},
    onError: () => {},
  });

  return { mutation, cancelQuery };
}

export function useGetChunks() {
  const queryClient = useQueryClient();
  const abortControllerRef = useRef<AbortController | null>(null);

  const cancelQuery = () => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      abortControllerRef.current = null;
    }
    queryClient.cancelQueries({ queryKey: ["predictions"] });
  };

  const mutation: UseMutationResult<
    { chunks: ChunksResponseType; chunksArray: ChunksArrayItem[] } | undefined,
    Error,
    ChunkRequestType,
    unknown
  > = useMutation({
    mutationFn: async (data: ChunkRequestType) => {
      // Отменить предыдущий запрос
      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
      }

      const abortController = new AbortController();
      abortControllerRef.current = abortController;

      try {
        const response = await getChunks(data, abortController.signal);
        return response;
      } catch (error) {
        if (error instanceof Error && error.name === "AbortError") {
          console.log("Запрос был отменен");
          return undefined;
        } else {
          throw error;
        }
      } finally {
        if (abortControllerRef.current === abortController) {
          abortControllerRef.current = null;
        }
      }
    },
    onSuccess: () => {
      // if (chunksData) {
      //   llmMutation.mutate({ query: query, chunks: chunksData });
      // }
    },
    onError: () => {},
  });

  return { mutation, cancelQuery };
}

export const usePostFeedback = () => {
  return useMutation({
    mutationFn: async (data: FeedbackRequestType) => {
      const response = await postFeedback(data);
      return response;
    },
    onSuccess: () => {},
    onError: () => {},
  });
};

export const useGetLogs = () => {
  return useMutation({
    mutationKey: ["logs"],
    mutationFn: async (data: LogRequestType | undefined) => {
      const response = await getLogs(data);
      return response;
    },
  });
};

export const useGetAbbreviations = () => {
  return useQuery({
    queryKey: ["abbreviations"],
    queryFn: async () => {
      const response = await getAbbreviations();
      return response;
    },
  });
};