import { useRef, useState, useTransition, useCallback } from "react";
import { useQueryClient } from "@tanstack/react-query";
import { chatApi } from "../api/chat";

export function useStreamingChat(conversationId: string | null) {
  const [streamingContent, setStreamingContent] = useState<string>("");
  const [isStreaming, setIsStreaming] = useState(false);
  const abortRef = useRef<AbortController | null>(null);
  const queryClient = useQueryClient();
  const [, startTransition] = useTransition();

  const startStream = useCallback(
    (userMessage: string, agentId?: string) => {
      if (!conversationId) return;

      setIsStreaming(true);
      setStreamingContent("");

      const abort = new AbortController();
      abortRef.current = abort;

      chatApi.postMessageAndStream(
        conversationId,
        { content: userMessage, agentId },
        {
          onToken: (token: string) => {
            startTransition(() => {
              setStreamingContent((prev) => prev + token);
            });
          },
          onDone: (_messageId: string, _content: string) => {
            setIsStreaming(false);
            setStreamingContent("");
            abortRef.current = null;
            // Optimistically insert the completed message into cache to avoid flash (Pitfall 2)
            queryClient.setQueryData(
              ["chat", "messages", conversationId],
              (old: unknown) => old, // Keep existing data -- invalidation will refetch
            );
            queryClient.invalidateQueries({ queryKey: ["chat", "messages", conversationId] });
            queryClient.invalidateQueries({ queryKey: ["chat", "conversations"] });
          },
          onError: (error: string) => {
            setIsStreaming(false);
            abortRef.current = null;
            console.error("[useStreamingChat] Stream error:", error);
          },
        },
        abort.signal,
      );
    },
    [conversationId, queryClient, startTransition],
  );

  const stop = useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
    const partial = streamingContent;
    setIsStreaming(false);
    setStreamingContent("");
    // Persist partial content with [stopped] suffix (Open Question 3)
    if (conversationId && partial.trim()) {
      chatApi.savePartialMessage(conversationId, {
        role: "assistant",
        content: partial.trim() + " [stopped]",
      }).then(() => {
        queryClient.invalidateQueries({ queryKey: ["chat", "messages", conversationId] });
      });
    }
  }, [conversationId, streamingContent, queryClient]);

  return { streamingContent, isStreaming, startStream, stop };
}
