stack/apps/web/src/hooks/useChat.ts

/**
 * useChat hook
 * Manages chat state, LLM interactions, and conversation persistence
 */

import { useState, useCallback, useRef } from "react";
import {
  sendChatMessage,
  streamChatMessage,
  type ChatMessage as ApiChatMessage,
} from "@/lib/api/chat";
import { createConversation, updateConversation, getIdea, type Idea } from "@/lib/api/ideas";
import { safeJsonParse, isMessageArray } from "@/lib/utils/safe-json";

export interface Message {
  id: string;
  role: "user" | "assistant" | "system";
  content: string;
  thinking?: string;
  createdAt: string;
  model?: string;
  provider?: string;
  promptTokens?: number;
  completionTokens?: number;
  totalTokens?: number;
}

export interface UseChatOptions {
  model?: string;
  temperature?: number;
  maxTokens?: number;
  systemPrompt?: string;
  projectId?: string | null;
  onError?: (error: Error) => void;
}

export interface UseChatReturn {
  messages: Message[];
  isLoading: boolean;
  isStreaming: boolean;
  error: string | null;
  conversationId: string | null;
  conversationTitle: string | null;
  sendMessage: (content: string) => Promise<void>;
  abortStream: () => void;
  loadConversation: (ideaId: string) => Promise<void>;
  startNewConversation: (projectId?: string | null) => void;
  setMessages: React.Dispatch<React.SetStateAction<Message[]>>;
  clearError: () => void;
}

const DEFAULT_MODEL = "llama3.2";
const WELCOME_MESSAGE: Message = {
  id: "welcome",
  role: "assistant",
  content: "Hello! I'm your AI assistant. How can I help you today?",
  createdAt: new Date().toISOString(),
};

/**
 * Hook for managing chat conversations
 */
export function useChat(options: UseChatOptions = {}): UseChatReturn {
  const {
    model = DEFAULT_MODEL,
    temperature,
    maxTokens,
    systemPrompt,
    projectId,
    onError,
  } = options;

  const [messages, setMessages] = useState<Message[]>([WELCOME_MESSAGE]);
  const [isLoading, setIsLoading] = useState(false);
  const [isStreaming, setIsStreaming] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [conversationId, setConversationId] = useState<string | null>(null);
  const [conversationTitle, setConversationTitle] = useState<string | null>(null);

  // Track project ID in ref to prevent stale closures
  const projectIdRef = useRef<string | null>(projectId ?? null);
  projectIdRef.current = projectId ?? null;

  // Track messages in ref to prevent stale closures during rapid sends
  const messagesRef = useRef<Message[]>(messages);
  messagesRef.current = messages;

  // AbortController ref for the active stream
  const abortControllerRef = useRef<AbortController | null>(null);

  // Track conversation state in refs to avoid stale closures in streaming callbacks
  const conversationIdRef = useRef<string | null>(conversationId);
  conversationIdRef.current = conversationId;

  const conversationTitleRef = useRef<string | null>(conversationTitle);
  conversationTitleRef.current = conversationTitle;

  /**
   * Convert our Message format to API ChatMessage format
   */
  const convertToApiMessages = useCallback((msgs: Message[]): ApiChatMessage[] => {
    return msgs
      .filter((msg) => msg.role !== "system" || msg.id !== "welcome")
      .map((msg) => ({
        role: msg.role,
        content: msg.content,
      }));
  }, []);

  /**
   * Generate a conversation title from the first user message
   */
  const generateTitle = useCallback((firstMessage: string): string => {
    const maxLength = 60;
    const trimmed = firstMessage.trim();

    if (trimmed.length <= maxLength) {
      return trimmed;
    }

    return trimmed.substring(0, maxLength - 3) + "...";
  }, []);

  /**
   * Serialize messages to JSON for storage
   */
  const serializeMessages = useCallback((msgs: Message[]): string => {
    return JSON.stringify(msgs, null, 2);
  }, []);

  /**
   * Deserialize messages from JSON with runtime type validation
   */
  const deserializeMessages = useCallback((json: string): Message[] => {
    return safeJsonParse(json, isMessageArray, [WELCOME_MESSAGE]);
  }, []);

  /**
   * Save conversation to backend.
   * Uses refs for conversation state to avoid stale closures in streaming callbacks.
   */
  const saveConversation = useCallback(
    async (msgs: Message[], title: string): Promise<string> => {
      const content = serializeMessages(msgs);
      const currentConvId = conversationIdRef.current;

      if (currentConvId) {
        await updateConversation(currentConvId, content, title);
        return currentConvId;
      } else {
        const idea = await createConversation(title, content, projectIdRef.current ?? undefined);
        setConversationId(idea.id);
        setConversationTitle(title);
        conversationIdRef.current = idea.id;
        conversationTitleRef.current = title;
        return idea.id;
      }
    },
    [serializeMessages]
  );

  /**
   * Abort an active stream
   */
  const abortStream = useCallback((): void => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      abortControllerRef.current = null;
    }
    setIsStreaming(false);
    setIsLoading(false);
  }, []);

  /**
   * Send a message to the LLM using streaming, with fallback to non-streaming
   */
  const sendMessage = useCallback(
    async (content: string): Promise<void> => {
      if (!content.trim() || isLoading || isStreaming) {
        return;
      }

      const userMessage: Message = {
        id: `user-${Date.now().toString()}-${Math.random().toString(36).slice(2, 8)}`,
        role: "user",
        content: content.trim(),
        createdAt: new Date().toISOString(),
      };

      setMessages((prev) => {
        const updated = [...prev, userMessage];
        messagesRef.current = updated;
        return updated;
      });
      setIsLoading(true);
      setError(null);

      const assistantMessageId = `assistant-${Date.now().toString()}-${Math.random().toString(36).slice(2, 8)}`;
      const placeholderMessage: Message = {
        id: assistantMessageId,
        role: "assistant",
        content: "",
        createdAt: new Date().toISOString(),
        model,
      };

      const currentMessages = messagesRef.current;
      const apiMessages = convertToApiMessages(currentMessages);

      const request = {
        model,
        messages: apiMessages,
        ...(temperature !== undefined && { temperature }),
        ...(maxTokens !== undefined && { maxTokens }),
        ...(systemPrompt !== undefined && { systemPrompt }),
      };

      const controller = new AbortController();
      abortControllerRef.current = controller;

      let streamingSucceeded = false;

      try {
        await new Promise<void>((resolve, reject) => {
          let hasReceivedData = false;

          streamChatMessage(
            request,
            (chunk: string) => {
              if (!hasReceivedData) {
                hasReceivedData = true;
                setIsLoading(false);
                setIsStreaming(true);
                setMessages((prev) => {
                  const updated = [...prev, { ...placeholderMessage }];
                  messagesRef.current = updated;
                  return updated;
                });
              }

              setMessages((prev) => {
                const updated = prev.map((msg) =>
                  msg.id === assistantMessageId ? { ...msg, content: msg.content + chunk } : msg
                );
                messagesRef.current = updated;
                return updated;
              });
            },
            () => {
              streamingSucceeded = true;
              setIsStreaming(false);
              abortControllerRef.current = null;
              resolve();
            },
            (err: Error) => {
              reject(err);
            },
            controller.signal
          );
        });
      } catch (err: unknown) {
        if (controller.signal.aborted) {
          setIsStreaming(false);
          setIsLoading(false);
          abortControllerRef.current = null;

          // Remove placeholder if no content was received
          setMessages((prev) => {
            const assistantMsg = prev.find((m) => m.id === assistantMessageId);
            if (assistantMsg?.content === "") {
              const updated = prev.filter((m) => m.id !== assistantMessageId);
              messagesRef.current = updated;
              return updated;
            }
            messagesRef.current = prev;
            return prev;
          });
          return;
        }

        // Streaming failed — fall back to non-streaming
        console.warn("Streaming failed, falling back to non-streaming", {
          error: err instanceof Error ? err : new Error(String(err)),
        });

        setMessages((prev) => {
          const withoutPlaceholder = prev.filter((m) => m.id !== assistantMessageId);
          messagesRef.current = withoutPlaceholder;
          return withoutPlaceholder;
        });
        setIsStreaming(false);

        try {
          const response = await sendChatMessage(request);

          const assistantMessage: Message = {
            id: `assistant-${Date.now().toString()}`,
            role: "assistant",
            content: response.message.content,
            createdAt: new Date().toISOString(),
            model: response.model,
            promptTokens: response.promptEvalCount ?? 0,
            completionTokens: response.evalCount ?? 0,
            totalTokens: (response.promptEvalCount ?? 0) + (response.evalCount ?? 0),
          };

          setMessages((prev) => {
            const updated = [...prev, assistantMessage];
            messagesRef.current = updated;
            return updated;
          });

          streamingSucceeded = true;
        } catch (fallbackErr: unknown) {
          const errorMsg =
            fallbackErr instanceof Error ? fallbackErr.message : "Failed to send message";
          setError("Unable to send message. Please try again.");
          onError?.(fallbackErr instanceof Error ? fallbackErr : new Error(errorMsg));
          console.error("Failed to send chat message", {
            error: fallbackErr,
            errorType: "LLM_ERROR",
            conversationId: conversationIdRef.current,
            messageLength: content.length,
            messagePreview: content.substring(0, 50),
            model,
            messageCount: messagesRef.current.length,
            timestamp: new Date().toISOString(),
          });

          const errorMessage: Message = {
            id: `error-${String(Date.now())}`,
            role: "assistant",
            content: "Something went wrong. Please try again.",
            createdAt: new Date().toISOString(),
          };
          setMessages((prev) => {
            const updated = [...prev, errorMessage];
            messagesRef.current = updated;
            return updated;
          });
          setIsLoading(false);
          return;
        }
      }

      setIsLoading(false);

      if (!streamingSucceeded) {
        return;
      }

      const finalMessages = messagesRef.current;

      const isFirstMessage =
        !conversationIdRef.current && finalMessages.filter((m) => m.role === "user").length === 1;
      const title = isFirstMessage
        ? generateTitle(content)
        : (conversationTitleRef.current ?? "Chat Conversation");

      try {
        await saveConversation(finalMessages, title);
      } catch (saveErr) {
        const saveErrorMsg =
          saveErr instanceof Error ? saveErr.message : "Unknown persistence error";
        setError("Message sent but failed to save. Please try again.");
        onError?.(saveErr instanceof Error ? saveErr : new Error(saveErrorMsg));
        console.error("Failed to save conversation", {
          error: saveErr,
          errorType: "PERSISTENCE_ERROR",
          conversationId: conversationIdRef.current,
          detail: saveErrorMsg,
        });
      }
    },
    [
      isLoading,
      isStreaming,
      model,
      temperature,
      maxTokens,
      systemPrompt,
      onError,
      convertToApiMessages,
      generateTitle,
      saveConversation,
    ]
  );

  /**
   * Load an existing conversation from the backend
   */
  const loadConversation = useCallback(
    async (ideaId: string): Promise<void> => {
      try {
        setIsLoading(true);
        setError(null);

        const idea: Idea = await getIdea(ideaId);
        const msgs = deserializeMessages(idea.content);

        setMessages(msgs);
        setConversationId(idea.id);
        setConversationTitle(idea.title ?? null);
        conversationIdRef.current = idea.id;
        conversationTitleRef.current = idea.title ?? null;
      } catch (err) {
        const errorMsg = err instanceof Error ? err.message : "Failed to load conversation";
        setError("Unable to load conversation. Please try again.");
        onError?.(err instanceof Error ? err : new Error(errorMsg));
        console.error("Failed to load conversation", {
          error: err,
          errorType: "LOAD_ERROR",
          ideaId,
          timestamp: new Date().toISOString(),
        });
      } finally {
        setIsLoading(false);
      }
    },
    [deserializeMessages, onError]
  );

  /**
   * Start a new conversation
   */
  const startNewConversation = useCallback((newProjectId?: string | null): void => {
    setMessages([WELCOME_MESSAGE]);
    setConversationId(null);
    setConversationTitle(null);
    setError(null);
    conversationIdRef.current = null;
    conversationTitleRef.current = null;
    projectIdRef.current = newProjectId ?? null;
  }, []);

  /**
   * Clear error message
   */
  const clearError = useCallback((): void => {
    setError(null);
  }, []);

  return {
    messages,
    isLoading,
    isStreaming,
    error,
    conversationId,
    conversationTitle,
    sendMessage,
    abortStream,
    loadConversation,
    startNewConversation,
    setMessages,
    clearError,
  };
}