import { dataService, QueryKeys, Tools } from 'librechat-data-provider';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import type { UseMutationResult } from '@tanstack/react-query';
import type * as t from 'librechat-data-provider';

export const useToolCallMutation = <T extends t.ToolId>(
  toolId: T,
  options?: t.ToolCallMutationOptions<T>,
): UseMutationResult<t.ToolCallResponse, Error, t.ToolParams<T>> => {
  const queryClient = useQueryClient();
  return useMutation(
    (toolParams: t.ToolParams<T>) => {
      return dataService.callTool({
        toolId,
        toolParams,
      });
    },
    {
      onMutate: (variables) => options?.onMutate?.(variables),
      onError: (error, variables, context) => options?.onError?.(error, variables, context),
      onSuccess: (response, variables, context) => {
        queryClient.setQueryData<t.ToolCallResults>(
          [QueryKeys.toolCalls, variables.conversationId],
          (prev) => [
            ...(prev ?? []),
            {
              user: '',
              toolId: Tools.execute_code,
              partIndex: variables.partIndex,
              messageId: variables.messageId,
              blockIndex: variables.blockIndex,
              conversationId: variables.conversationId,
              result: response.result,
              attachments: response.attachments,
            },
          ],
        );
        return options?.onSuccess?.(response, variables, context);
      },
    },
  );
};

interface CreateToolData {
  name: string;
  description: string;
  metadata?: Record<string, unknown>;
  // MCP-specific fields
  url?: string;
  icon?: string;
  tools?: string[];
  trust?: boolean;
  customHeaders?: Array<{
    id: string;
    name: string;
    value: string;
  }>;
  requestTimeout?: number;
  connectionTimeout?: number;
}

export const useAddToolMutation = (
  options?: t.MutationOptions<Record<string, unknown>, CreateToolData>,
): UseMutationResult<Record<string, unknown>, Error, CreateToolData> => {
  const queryClient = useQueryClient();

  return useMutation(
    (toolData: CreateToolData) => {
      return dataService.createTool(toolData);
    },
    {
      onMutate: (variables) => options?.onMutate?.(variables),
      onError: (error, variables, context) => options?.onError?.(error, variables, context),
      onSuccess: (data, variables, context) => {
        // Invalidate tools list to trigger refetch
        queryClient.invalidateQueries([QueryKeys.tools]);
        queryClient.invalidateQueries([QueryKeys.mcpTools]);
        return options?.onSuccess?.(data, variables, context);
      },
    },
  );
};
