adiuva/src/renderer/hooks/useAIChat.ts

import { useState, useCallback, useRef } from 'react';
import { trpc } from '@/lib/trpc';

interface ChatMessage {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  error?: boolean;
}

export interface ChatContext {
  type: 'global' | 'project';
  projectId?: string;
  uiContext?: string;
}

interface UseAIChatReturn {
  messages: ChatMessage[];
  input: string;
  setInput: (v: string) => void;
  isStreaming: boolean;
  streamingContent: string;
  handleSend: (overrideMessage?: string, overrideContext?: ChatContext) => void;
  clearMessages: () => void;
}

interface UseAIChatOptions {
  onSectionTag?: (sectionId: string) => void;
}

export function useAIChat(defaultContext: ChatContext, options?: UseAIChatOptions): UseAIChatReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [input, setInput] = useState('');
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingContent, setStreamingContent] = useState('');

  const streamingContentRef = useRef('');
  const chatMutation = trpc.ai.chat.useMutation();

  const clearMessages = useCallback(() => {
    setMessages([]);
    setStreamingContent('');
    streamingContentRef.current = '';
  }, []);

  const handleSend = useCallback(
    (overrideMessage?: string, overrideContext?: ChatContext) => {
      const trimmed = (overrideMessage ?? input).trim();
      if (!trimmed || isStreaming) return;

      const userMsg: ChatMessage = {
        id: crypto.randomUUID(),
        role: 'user',
        content: trimmed,
      };

      setMessages((prev) => [...prev, userMsg]);
      if (!overrideMessage) setInput('');
      setIsStreaming(true);
      setStreamingContent('');
      streamingContentRef.current = '';

      const unsubscribe = window.electronAI.onStreamChunk(({ token, done }) => {
        if (done) {
          let finalContent = streamingContentRef.current;

          // Parse and strip [SECTION:xxx] tag from AI response
          const sectionMatch = finalContent.match(/^\[SECTION:([\w-]+)\]\s*/);
          if (sectionMatch) {
            finalContent = finalContent.slice(sectionMatch[0].length);
            options?.onSectionTag?.(sectionMatch[1]!);
          }

          setMessages((prev) => [
            ...prev,
            { id: crypto.randomUUID(), role: 'assistant', content: finalContent },
          ]);
          setStreamingContent('');
          streamingContentRef.current = '';
          setIsStreaming(false);
          unsubscribe();
          return;
        }
        streamingContentRef.current += token;
        setStreamingContent(streamingContentRef.current);
      });

      const ctx = overrideContext ?? defaultContext;

      chatMutation.mutate(
        {
          message: trimmed,
          context: {
            type: ctx.type,
            ...(ctx.type === 'project' && ctx.projectId ? { projectId: ctx.projectId } : {}),
            ...(ctx.uiContext ? { uiContext: ctx.uiContext } : {}),
          },
        },
        {
          onSuccess: (data) => {
            if (data.error) {
              unsubscribe();
              setMessages((prev) => [
                ...prev,
                { id: crypto.randomUUID(), role: 'assistant', content: data.error!, error: true },
              ]);
              setStreamingContent('');
              streamingContentRef.current = '';
              setIsStreaming(false);
            }
          },
          onError: (err) => {
            unsubscribe();
            setMessages((prev) => [
              ...prev,
              {
                id: crypto.randomUUID(),
                role: 'assistant',
                content: err.message || 'An unexpected error occurred.',
                error: true,
              },
            ]);
            setStreamingContent('');
            streamingContentRef.current = '';
            setIsStreaming(false);
          },
        },
      );
    },
    [input, isStreaming, defaultContext, chatMutation],
  );

  return {
    messages,
    input,
    setInput,
    isStreaming,
    streamingContent,
    handleSend,
    clearMessages,
  };
}