import { useCallback, useEffect, useRef } from "react"; import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; import { useChatStore } from "@/stores/chat-store"; import { getChats, getMessages, createChat, updateChat, deleteChat, sendMessage, type Message, } from "@/api/chats"; import { useAuthStore } from "@/stores/auth-store"; export function useChat() { const queryClient = useQueryClient(); const user = useAuthStore((s) => s.user); const { currentChatId, setChats, setCurrentChat, addChat, removeChat, updateChatInList, setMessages, appendMessage, setStreaming, appendStreamingContent, finalizeStreamingMessage, clearStreamingContent, } = useChatStore(); const abortRef = useRef(null); // Fetch chats const chatsQuery = useQuery({ queryKey: ["chats"], queryFn: () => getChats(false), }); useEffect(() => { if (chatsQuery.data) { setChats(chatsQuery.data); } }, [chatsQuery.data, setChats]); // Fetch messages for current chat const messagesQuery = useQuery({ queryKey: ["messages", currentChatId], queryFn: () => getMessages(currentChatId!, 200), enabled: !!currentChatId, }); useEffect(() => { if (messagesQuery.data) { setMessages(messagesQuery.data); } }, [messagesQuery.data, setMessages]); // Create chat const createMutation = useMutation({ mutationFn: (title?: string) => createChat(title), onSuccess: (chat) => { addChat(chat); queryClient.invalidateQueries({ queryKey: ["chats"] }); }, }); // Delete chat const deleteMutation = useMutation({ mutationFn: (chatId: string) => deleteChat(chatId), onSuccess: (_, chatId) => { removeChat(chatId); if (currentChatId === chatId) { setCurrentChat(null); } queryClient.invalidateQueries({ queryKey: ["chats"] }); }, }); // Archive/unarchive const archiveMutation = useMutation({ mutationFn: ({ chatId, archived }: { chatId: string; archived: boolean }) => updateChat(chatId, { is_archived: archived }), onSuccess: (chat) => { if (chat.is_archived) { removeChat(chat.id); } else { updateChatInList(chat); } queryClient.invalidateQueries({ queryKey: ["chats"] }); }, }); // Send message with SSE const handleSendMessage = useCallback( (content: string) => { if (!currentChatId) return; // Add user message optimistically const tempUserMsg: Message = { id: crypto.randomUUID(), chat_id: currentChatId, role: "user", content, metadata: null, created_at: new Date().toISOString(), }; appendMessage(tempUserMsg); setStreaming(true); // Abort previous stream if any abortRef.current?.abort(); const controller = new AbortController(); abortRef.current = controller; sendMessage(currentChatId, content, { onDelta: (delta) => appendStreamingContent(delta), onComplete: (messageId, metadata) => { const finalContent = useChatStore.getState().streamingContent; finalizeStreamingMessage(messageId, finalContent, metadata); queryClient.invalidateQueries({ queryKey: ["chats"] }); }, onError: (detail) => { clearStreamingContent(); console.error("SSE error:", detail); }, }, controller.signal); }, [ currentChatId, appendMessage, setStreaming, appendStreamingContent, finalizeStreamingMessage, clearStreamingContent, queryClient, ] ); // Cleanup on unmount useEffect(() => { return () => { abortRef.current?.abort(); }; }, []); // Cleanup on chat switch useEffect(() => { abortRef.current?.abort(); clearStreamingContent(); }, [currentChatId, clearStreamingContent]); const limitReached = chatsQuery.data ? chatsQuery.data.length >= (user?.max_chats ?? 10) && user?.role !== "admin" : false; return { chatsQuery, messagesQuery, createChat: createMutation.mutateAsync, deleteChat: deleteMutation.mutate, archiveChat: (chatId: string, archived: boolean) => archiveMutation.mutate({ chatId, archived }), sendMessage: handleSendMessage, limitReached, }; }