import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import {
  getConnections,
  saveConnection,
  deleteConnection,
  testConnection,
  connectDb,
  disconnectDb,
} from "@/lib/tauri";
import { useAppStore } from "@/stores/app-store";
import type { ConnectionConfig } from "@/types";

export function useConnections() {
  return useQuery({
    queryKey: ["connections"],
    queryFn: getConnections,
  });
}

export function useSaveConnection() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (config: ConnectionConfig) => saveConnection(config),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["connections"] });
    },
  });
}

export function useDeleteConnection() {
  const queryClient = useQueryClient();
  const { removeConnectedId, setActiveConnectionId, activeConnectionId } =
    useAppStore();

  return useMutation({
    mutationFn: (id: string) => deleteConnection(id),
    onSuccess: (_data, id) => {
      queryClient.invalidateQueries({ queryKey: ["connections"] });
      removeConnectedId(id);
      if (activeConnectionId === id) {
        setActiveConnectionId(null);
      }
    },
  });
}

export function useTestConnection() {
  return useMutation({
    mutationFn: (config: ConnectionConfig) => testConnection(config),
  });
}

export function useConnect() {
  const {
    addConnectedId,
    setActiveConnectionId,
    setPgVersion,
    setDbFlavor,
    setCurrentDatabase,
    addTab,
    tabs,
  } = useAppStore();

  return useMutation({
    mutationFn: async (config: ConnectionConfig) => {
      const result = await connectDb(config);
      return { id: config.id, ...result, database: config.database };
    },
    onSuccess: ({ id, version, flavor, database }) => {
      addConnectedId(id);
      setActiveConnectionId(id);
      setPgVersion(version);
      setDbFlavor(id, flavor);
      setCurrentDatabase(database);

      const hasChatForConnection = tabs.some(
        (t) => t.type === "chat" && t.connectionId === id
      );
      if (!hasChatForConnection) {
        addTab({
          id: crypto.randomUUID(),
          type: "chat",
          title: "Chat",
          connectionId: id,
          database,
        });
      }
    },
  });
}

export function useDisconnect() {
  const {
    removeConnectedId,
    setActiveConnectionId,
    activeConnectionId,
    setPgVersion,
  } = useAppStore();

  return useMutation({
    mutationFn: (id: string) => disconnectDb(id),
    onSuccess: (_data, id) => {
      removeConnectedId(id);
      if (activeConnectionId === id) {
        setActiveConnectionId(null);
        setPgVersion(null);
      }
    },
  });
}

export function useReconnect() {
  const queryClient = useQueryClient();
  const { setPgVersion, setDbFlavor, setCurrentDatabase } = useAppStore();

  return useMutation({
    mutationFn: async (config: ConnectionConfig) => {
      await disconnectDb(config.id);
      const result = await connectDb(config);
      return { id: config.id, ...result, database: config.database };
    },
    onSuccess: ({ id, version, flavor, database }) => {
      setPgVersion(version);
      setDbFlavor(id, flavor);
      setCurrentDatabase(database);
      queryClient.invalidateQueries({ queryKey: ["databases"] });
      queryClient.invalidateQueries({ queryKey: ["schemas"] });
      queryClient.invalidateQueries({ queryKey: ["completion-schema"] });
    },
  });
}
