viewer/src/hooks/useChatApi.ts

import { useState, useCallback, useEffect } from "react";
import { useMutation } from "@tanstack/react-query";
import { v4 as uuidv4 } from "uuid";

// Generate a conversation ID to help backend keep track of the chat history
const getConversationId = () => {
  const storedId = localStorage.getItem("chat_conversation_id");
  if (storedId) return storedId;

  const newId = uuidv4();
  localStorage.setItem("chat_conversation_id", newId);
  return newId;
};

export interface ChatMessage {
  sender: string;
  text: string;
  imageUrl?: string;
}

export const useChatApi = () => {
  const [conversationId] = useState(getConversationId);
  const [isStreaming, setIsStreaming] = useState(false);

  // Use mutation for the API call
  const chatMutation = useMutation({
    mutationFn: async (prompt: string) => {
      const controller = new AbortController();
      const signal = controller.signal;

      const response = await fetch("http://127.0.0.1:8000/prompt", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          session_id: conversationId,
          prompt: prompt,
        }),
        signal,
      });

      if (!response.ok) {
        throw new Error("Network response was not ok");
      }

      return { response, controller };
    },
  });

  // Process the streaming response
  const streamResponse = useCallback(
    async (
      response: Response,
      onChunk: (chunk: string) => void,
      onImage?: (imageUrl: string) => void
    ) => {
      if (!response.body) {
        throw new Error("Response body is null");
      }

      setIsStreaming(true);

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = "";

      try {
        while (true) {
          const { done, value } = await reader.read();

          if (done) {
            break;
          }

          // Decode the chunk
          const text = decoder.decode(value, { stream: true });
          buffer += text;

          // Process SSE format: "data: {json}\n\n"
          const lines = buffer.split("\n\n");
          buffer = lines.pop() || "";

          for (const line of lines) {
            if (line.startsWith("data: ")) {
              const jsonStr = line.slice(6);
              try {
                const data = JSON.parse(jsonStr);

                // Check if there's text content
                if (data.text) {
                  onChunk(data.text);
                }

                // Check if there's an image URL
                if (data.image_url) {
                  onImage?.(data.image_url);
                }
              } catch (e) {
                console.error("Error parsing SSE JSON:", e);
              }
            }
          }
        }
      } catch (error) {
        console.error("Error reading stream:", error);
      } finally {
        setIsStreaming(false);
      }
    },
    []
  );

  // Cleanup function to abort any pending requests
  useEffect(() => {
    return () => {
      if (chatMutation.data?.controller) {
        chatMutation.data.controller.abort();
      }
    };
  }, [chatMutation.data]);

  return {
    sendMessage: chatMutation.mutate,
    streamResponse,
    isLoading: chatMutation.isPending || isStreaming,
    error: chatMutation.error,
  };
};