frontend/src/hooks/react-query/knowledge-base/use-knowledge-base-queries.ts

import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
import { toast } from 'sonner';
import { createClient } from '@/lib/supabase/client';
import { knowledgeBaseKeys } from './keys';
import { 
  CreateKnowledgeBaseEntryRequest, 
  KnowledgeBaseEntry, 
  KnowledgeBaseListResponse, 
  UpdateKnowledgeBaseEntryRequest,
  FileUploadRequest,
  GitCloneRequest,
  ProcessingJob,
  ProcessingJobsResponse,
  UploadResponse,
  CloneResponse
} from './types';

const API_URL = process.env.NEXT_PUBLIC_BACKEND_URL || '';

const useAuthHeaders = () => {
  const getHeaders = async () => {
    const supabase = createClient();
    const { data: { session } } = await supabase.auth.getSession();
    
    if (!session?.access_token) {
      throw new Error('No access token available');
    }
    return {
      'Authorization': `Bearer ${session.access_token}`,
      'Content-Type': 'application/json',
    };  
  };
  
  return { getHeaders };
};

export function useKnowledgeBaseEntries(threadId: string, includeInactive = false) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.thread(threadId),
    queryFn: async (): Promise<KnowledgeBaseListResponse> => {
      const headers = await getHeaders();
      const url = new URL(`${API_URL}/knowledge-base/threads/${threadId}`);
      url.searchParams.set('include_inactive', includeInactive.toString());
      
      const response = await fetch(url.toString(), { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch knowledge base entries');
      }
      
      return await response.json();
    },
    enabled: !!threadId,
  });
}

export function useKnowledgeBaseEntry(entryId: string) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.entry(entryId),
    queryFn: async (): Promise<KnowledgeBaseEntry> => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/${entryId}`, { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch knowledge base entry');
      }
      
      return await response.json();
    },
    enabled: !!entryId,
  });
}

export function useKnowledgeBaseContext(threadId: string, maxTokens = 4000) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.context(threadId),
    queryFn: async () => {
      const headers = await getHeaders();
      const url = new URL(`${API_URL}/knowledge-base/threads/${threadId}/context`);
      url.searchParams.set('max_tokens', maxTokens.toString());
      
      const response = await fetch(url.toString(), { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch knowledge base context');
      }
      
      return await response.json();
    },
    enabled: !!threadId,
  });
}

export function useCreateKnowledgeBaseEntry() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async ({ threadId, data }: { threadId: string; data: CreateKnowledgeBaseEntryRequest }) => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/threads/${threadId}`, {
        method: 'POST',
        headers: {
          ...headers,
          'Content-Type': 'application/json',
        },
        body: JSON.stringify(data),
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to create knowledge base entry');
      }
      
      return await response.json();
    },
    onSuccess: (_, { threadId }) => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.thread(threadId) });
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.context(threadId) });
      toast.success('Knowledge base entry created successfully');
    },
    onError: (error) => {
      toast.error(`Failed to create knowledge base entry: ${error.message}`);
    },
  });
}

export function useUpdateKnowledgeBaseEntry() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async ({ entryId, data }: { entryId: string; data: UpdateKnowledgeBaseEntryRequest }) => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/${entryId}`, {
        method: 'PUT',
        headers: {
          ...headers,
          'Content-Type': 'application/json',
        },
        body: JSON.stringify(data),
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to update knowledge base entry');
      }
      
      return await response.json();
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.all });
      toast.success('Knowledge base entry updated successfully');
    },
    onError: (error) => {
      toast.error(`Failed to update knowledge base entry: ${error.message}`);
    },
  });
}

export function useDeleteKnowledgeBaseEntry() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async (entryId: string) => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/${entryId}`, {
        method: 'DELETE',
        headers,
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to delete knowledge base entry');
      }
      
      return await response.json();
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.all });
      toast.success('Knowledge base entry deleted successfully');
    },
    onError: (error) => {
      toast.error(`Failed to delete knowledge base entry: ${error.message}`);
    },
  });
}

export function useAgentKnowledgeBaseEntries(agentId: string, includeInactive = false) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.agent(agentId),
    queryFn: async (): Promise<KnowledgeBaseListResponse> => {
      const headers = await getHeaders();
      const url = new URL(`${API_URL}/knowledge-base/agents/${agentId}`);
      url.searchParams.set('include_inactive', includeInactive.toString());
      
      const response = await fetch(url.toString(), { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch agent knowledge base entries');
      }
      
      return await response.json();
    },
    enabled: !!agentId,
  });
}

export function useCreateAgentKnowledgeBaseEntry() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async ({ agentId, data }: { agentId: string; data: CreateKnowledgeBaseEntryRequest }) => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/agents/${agentId}`, {
        method: 'POST',
        headers: {
          ...headers,
          'Content-Type': 'application/json',
        },
        body: JSON.stringify(data),
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to create agent knowledge base entry');
      }
      
      return await response.json();
    },
    onSuccess: (_, { agentId }) => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.agent(agentId) });
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.agentContext(agentId) });
      toast.success('Agent knowledge entry created successfully');
    },
    onError: (error) => {
      toast.error(`Failed to create agent knowledge entry: ${error.message}`);
    },
  });
}

export function useAgentKnowledgeBaseContext(agentId: string, maxTokens = 4000) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.agentContext(agentId),
    queryFn: async () => {
      const headers = await getHeaders();
      const url = new URL(`${API_URL}/knowledge-base/agents/${agentId}/context`);
      url.searchParams.set('max_tokens', maxTokens.toString());
      
      const response = await fetch(url.toString(), { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch agent knowledge base context');
      }
      
      return await response.json();
    },
    enabled: !!agentId,
  });
}

export function useCombinedKnowledgeBaseContext(threadId: string, agentId?: string, maxTokens = 4000) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.combinedContext(threadId, agentId),
    queryFn: async () => {
      const headers = await getHeaders();
      const url = new URL(`${API_URL}/knowledge-base/threads/${threadId}/combined-context`);
      url.searchParams.set('max_tokens', maxTokens.toString());
      if (agentId) {
        url.searchParams.set('agent_id', agentId);
      }
      
      const response = await fetch(url.toString(), { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch combined knowledge base context');
      }
      
      return await response.json();
    },
    enabled: !!threadId,
  });
}

// New hooks for file upload and git clone operations
export function useUploadAgentFiles() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async ({ agentId, file }: FileUploadRequest): Promise<UploadResponse> => {
      const supabase = createClient();
      const { data: { session } } = await supabase.auth.getSession();
      
      if (!session?.access_token) {
        throw new Error('No access token available');
      }

      const formData = new FormData();
      formData.append('file', file);

      const response = await fetch(`${API_URL}/knowledge-base/agents/${agentId}/upload-file`, {
        method: 'POST',
        headers: {
          'Authorization': `Bearer ${session.access_token}`,
        },
        body: formData,
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to upload file');
      }
      
      return await response.json();
    },
    onSuccess: (data, { agentId }) => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.agent(agentId) });
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.processingJobs(agentId) });
      toast.success('File uploaded successfully. Processing in background.');
    },
    onError: (error) => {
      toast.error(`Failed to upload file: ${error.message}`);
    },
  });
}

export function useCloneGitRepository() {
  const queryClient = useQueryClient();
  const { getHeaders } = useAuthHeaders();
  
  return useMutation({
    mutationFn: async ({ agentId, git_url, branch = 'main' }: GitCloneRequest): Promise<CloneResponse> => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/agents/${agentId}/clone-git-repo`, {
        method: 'POST',
        headers,
        body: JSON.stringify({ git_url, branch }),
      });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to clone repository');
      }
      
      return await response.json();
    },
    onSuccess: (data, { agentId }) => {
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.agent(agentId) });
      queryClient.invalidateQueries({ queryKey: knowledgeBaseKeys.processingJobs(agentId) });
      toast.success('Repository cloning started. Processing in background.');
    },
    onError: (error) => {
      toast.error(`Failed to clone repository: ${error.message}`);
    },
  });
}

export function useAgentProcessingJobs(agentId: string) {
  const { getHeaders } = useAuthHeaders();
  
  return useQuery({
    queryKey: knowledgeBaseKeys.processingJobs(agentId),
    queryFn: async (): Promise<ProcessingJobsResponse> => {
      const headers = await getHeaders();
      const response = await fetch(`${API_URL}/knowledge-base/agents/${agentId}/processing-jobs`, { headers });
      
      if (!response.ok) {
        const error = await response.text();
        throw new Error(error || 'Failed to fetch processing jobs');
      }
      
      return await response.json();
    },
    enabled: !!agentId,
    refetchInterval: 5000,
  });
}