personal-ai-assistant/frontend/src/hooks/use-chat.ts

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<AbortController | null>(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,
  };
}