Atomizer/atomizer-dashboard/frontend/src/hooks/useChat.ts

import { useState, useCallback, useRef, useEffect } from 'react';
import { Message } from '../components/chat/ChatMessage';
import { ToolCall } from '../components/chat/ToolCallCard';

export type ChatMode = 'user' | 'power';

export interface CanvasState {
  nodes: any[];
  edges: any[];
  studyName?: string;
  studyPath?: string;
}

interface UseChatOptions {
  studyId?: string | null;
  mode?: ChatMode;
  useWebSocket?: boolean;
  canvasState?: CanvasState | null;
  onError?: (error: string) => void;
}

interface ChatState {
  messages: Message[];
  isThinking: boolean;
  error: string | null;
  suggestions: string[];
  sessionId: string | null;
  mode: ChatMode;
  isConnected: boolean;
}

export function useChat({
  studyId,
  mode = 'user',
  useWebSocket = true,
  canvasState: initialCanvasState,
  onError,
}: UseChatOptions = {}) {
  const [state, setState] = useState<ChatState>({
    messages: [],
    isThinking: false,
    error: null,
    suggestions: [],
    sessionId: null,
    mode,
    isConnected: false,
  });

  // Track canvas state for sending with messages
  const canvasStateRef = useRef<CanvasState | null>(initialCanvasState || null);

  const abortControllerRef = useRef<AbortController | null>(null);
  const conversationHistoryRef = useRef<Array<{ role: string; content: string }>>([]);
  const wsRef = useRef<WebSocket | null>(null);
  const currentMessageRef = useRef<string>('');
  const currentToolCallsRef = useRef<ToolCall[]>([]);

  // Load suggestions when study changes
  useEffect(() => {
    loadSuggestions();
  }, [studyId]);

  // Create session and connect WebSocket
  useEffect(() => {
    if (!useWebSocket) return;

    const createSession = async () => {
      try {
        const response = await fetch('/api/claude/sessions', {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({
            mode: state.mode,
            study_id: studyId || undefined,
          }),
        });

        if (!response.ok) {
          throw new Error('Failed to create session');
        }

        const data = await response.json();
        setState((prev) => ({ ...prev, sessionId: data.session_id }));

        // Connect WebSocket
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/api/claude/sessions/${data.session_id}/ws`;
        const ws = new WebSocket(wsUrl);

        ws.onopen = () => {
          setState((prev) => ({ ...prev, isConnected: true }));
        };

        ws.onclose = () => {
          setState((prev) => ({ ...prev, isConnected: false }));
        };

        ws.onerror = () => {
          setState((prev) => ({ ...prev, isConnected: false }));
          onError?.('WebSocket connection error');
        };

        ws.onmessage = (event) => {
          try {
            const data = JSON.parse(event.data);
            handleWebSocketMessage(data);
          } catch (e) {
            console.error('Failed to parse WebSocket message:', e);
          }
        };

        wsRef.current = ws;
      } catch (e) {
        const message = e instanceof Error ? e.message : 'Failed to create session';
        onError?.(message);
      }
    };

    createSession();

    return () => {
      wsRef.current?.close();
      wsRef.current = null;
    };
  }, [useWebSocket, state.mode, studyId]);

  // Handle WebSocket messages
  const handleWebSocketMessage = useCallback((data: any) => {
    switch (data.type) {
      case 'text':
        currentMessageRef.current += data.content || '';
        setState((prev) => ({
          ...prev,
          messages: prev.messages.map((msg, idx) =>
            idx === prev.messages.length - 1 && msg.role === 'assistant'
              ? { ...msg, content: currentMessageRef.current }
              : msg
          ),
        }));
        break;

      case 'tool_call':
        const toolCall: ToolCall = {
          name: data.tool?.name || 'unknown',
          arguments: data.tool?.arguments,
        };
        currentToolCallsRef.current.push(toolCall);
        setState((prev) => ({
          ...prev,
          messages: prev.messages.map((msg, idx) =>
            idx === prev.messages.length - 1 && msg.role === 'assistant'
              ? { ...msg, toolCalls: [...currentToolCallsRef.current] }
              : msg
          ),
        }));
        break;

      case 'tool_result':
        // Update the last tool call with its result
        if (currentToolCallsRef.current.length > 0) {
          const lastIndex = currentToolCallsRef.current.length - 1;
          currentToolCallsRef.current[lastIndex] = {
            ...currentToolCallsRef.current[lastIndex],
            result: data.result,
          };
          setState((prev) => ({
            ...prev,
            messages: prev.messages.map((msg, idx) =>
              idx === prev.messages.length - 1 && msg.role === 'assistant'
                ? { ...msg, toolCalls: [...currentToolCallsRef.current] }
                : msg
            ),
          }));
        }
        break;

      case 'done':
        setState((prev) => ({
          ...prev,
          isThinking: false,
          messages: prev.messages.map((msg, idx) =>
            idx === prev.messages.length - 1 && msg.role === 'assistant'
              ? { ...msg, isStreaming: false }
              : msg
          ),
        }));
        // Update conversation history
        conversationHistoryRef.current.push({
          role: 'assistant',
          content: currentMessageRef.current,
        });
        currentMessageRef.current = '';
        currentToolCallsRef.current = [];
        break;

      case 'error':
        setState((prev) => ({
          ...prev,
          isThinking: false,
          error: data.message || 'Unknown error',
        }));
        onError?.(data.message || 'Unknown error');
        currentMessageRef.current = '';
        currentToolCallsRef.current = [];
        break;

      case 'context_updated':
        // Study context was updated - could show notification
        break;

      case 'canvas_updated':
        // Canvas state was updated - could show notification
        break;

      case 'pong':
        // Heartbeat response - ignore
        break;
    }
  }, [onError]);

  // Switch mode (requires new session)
  const switchMode = useCallback(async (newMode: ChatMode) => {
    if (newMode === state.mode) return;

    // Close existing WebSocket
    wsRef.current?.close();

    // Update mode - useEffect will create new session
    setState((prev) => ({
      ...prev,
      mode: newMode,
      sessionId: null,
      isConnected: false,
    }));
  }, [state.mode]);

  const loadSuggestions = async () => {
    try {
      const url = studyId
        ? `/api/claude/suggestions?study_id=${encodeURIComponent(studyId)}`
        : '/api/claude/suggestions';
      const response = await fetch(url);
      if (response.ok) {
        const data = await response.json();
        setState((prev) => ({ ...prev, suggestions: data.suggestions || [] }));
      }
    } catch {
      // Silently fail - suggestions are optional
    }
  };

  const generateMessageId = () => {
    return `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  };

  const sendMessage = useCallback(
    async (content: string) => {
      if (!content.trim() || state.isThinking) return;

      // Add user message
      const userMessage: Message = {
        id: generateMessageId(),
        role: 'user',
        content: content.trim(),
        timestamp: new Date(),
      };

      setState((prev) => ({
        ...prev,
        messages: [...prev.messages, userMessage],
        isThinking: true,
        error: null,
      }));

      // Update conversation history
      conversationHistoryRef.current.push({
        role: 'user',
        content: content.trim(),
      });

      // If using WebSocket and connected, send via WebSocket
      if (useWebSocket && wsRef.current?.readyState === WebSocket.OPEN) {
        // Create assistant message placeholder
        const assistantMessage: Message = {
          id: generateMessageId(),
          role: 'assistant',
          content: '',
          timestamp: new Date(),
          isStreaming: true,
          toolCalls: [],
        };

        setState((prev) => ({
          ...prev,
          messages: [...prev.messages, assistantMessage],
        }));

        // Reset current message/tool tracking
        currentMessageRef.current = '';
        currentToolCallsRef.current = [];

        // Send message via WebSocket with canvas state
        wsRef.current.send(
          JSON.stringify({
            type: 'message',
            content: content.trim(),
            canvas_state: canvasStateRef.current || undefined,
          })
        );
        return;
      }

      // Fall back to HTTP streaming
      // Cancel any in-flight request
      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
      }
      abortControllerRef.current = new AbortController();

      try {
        // Use streaming endpoint
        const response = await fetch('/api/claude/chat/stream', {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            message: content.trim(),
            study_id: studyId || undefined,
            conversation_history: conversationHistoryRef.current.slice(0, -1), // Exclude current message
          }),
          signal: abortControllerRef.current.signal,
        });

        if (!response.ok) {
          const errorData = await response.json().catch(() => ({}));
          throw new Error(errorData.detail || `Request failed: ${response.status}`);
        }

        // Create assistant message placeholder
        const assistantMessage: Message = {
          id: generateMessageId(),
          role: 'assistant',
          content: '',
          timestamp: new Date(),
          isStreaming: true,
        };

        setState((prev) => ({
          ...prev,
          messages: [...prev.messages, assistantMessage],
        }));

        // Read stream
        const reader = response.body?.getReader();
        const decoder = new TextDecoder();
        let fullContent = '';

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

            const chunk = decoder.decode(value, { stream: true });
            const lines = chunk.split('\n');

            for (const line of lines) {
              if (line.startsWith('data: ')) {
                const data = line.slice(6);
                if (data === '[DONE]') continue;

                try {
                  const parsed = JSON.parse(data);
                  if (parsed.token) {
                    fullContent += parsed.token;
                    // Update message content
                    setState((prev) => ({
                      ...prev,
                      messages: prev.messages.map((msg) =>
                        msg.id === assistantMessage.id
                          ? { ...msg, content: fullContent }
                          : msg
                      ),
                    }));
                  } else if (parsed.error) {
                    throw new Error(parsed.error);
                  }
                } catch {
                  // Skip invalid JSON
                }
              }
            }
          }
        }

        // Mark streaming complete
        setState((prev) => ({
          ...prev,
          isThinking: false,
          messages: prev.messages.map((msg) =>
            msg.id === assistantMessage.id
              ? { ...msg, isStreaming: false, content: fullContent }
              : msg
          ),
        }));

        // Update conversation history with assistant response
        conversationHistoryRef.current.push({
          role: 'assistant',
          content: fullContent,
        });
      } catch (error: any) {
        if (error.name === 'AbortError') {
          // Request was cancelled
          return;
        }

        const errorMessage = error.message || 'Failed to send message';
        setState((prev) => ({
          ...prev,
          isThinking: false,
          error: errorMessage,
        }));

        onError?.(errorMessage);
      }
    },
    [studyId, state.isThinking, onError]
  );

  const clearMessages = useCallback(() => {
    setState((prev) => ({
      ...prev,
      messages: [],
      error: null,
    }));
    conversationHistoryRef.current = [];
  }, []);

  const cancelRequest = useCallback(() => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      setState((prev) => ({
        ...prev,
        isThinking: false,
      }));
    }
  }, []);

  // Update canvas state (call this when canvas changes)
  const updateCanvasState = useCallback((newCanvasState: CanvasState | null) => {
    canvasStateRef.current = newCanvasState;

    // Also send to backend to update context
    if (useWebSocket && wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(
        JSON.stringify({
          type: 'set_canvas',
          canvas_state: newCanvasState,
        })
      );
    }
  }, [useWebSocket]);

  return {
    messages: state.messages,
    isThinking: state.isThinking,
    error: state.error,
    suggestions: state.suggestions,
    sessionId: state.sessionId,
    mode: state.mode,
    isConnected: state.isConnected,
    sendMessage,
    clearMessages,
    cancelRequest,
    switchMode,
    updateCanvasState,
  };
}