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

/**
 * Hook for Claude Code CLI integration
 *
 * Connects to backend that spawns actual Claude Code CLI processes.
 * This gives full power: file editing, command execution, etc.
 *
 * Unlike useChat (which uses MCP tools), this hook:
 * - Spawns actual Claude Code CLI in the backend
 * - Has full file system access
 * - Can edit files directly (not just return instructions)
 * - Uses Opus 4.5 model
 * - Has all Claude Code capabilities
 */

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

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

interface UseClaudeCodeOptions {
  studyId?: string | null;
  canvasState?: CanvasState | null;
  onError?: (error: string) => void;
  onCanvasRefresh?: (studyId: string) => void;
}

interface ClaudeCodeState {
  messages: Message[];
  isThinking: boolean;
  error: string | null;
  sessionId: string | null;
  isConnected: boolean;
  workingDir: string | null;
}

export function useClaudeCode({
  studyId,
  canvasState: initialCanvasState,
  onError,
  onCanvasRefresh,
}: UseClaudeCodeOptions = {}) {
  const [state, setState] = useState<ClaudeCodeState>({
    messages: [],
    isThinking: false,
    error: null,
    sessionId: null,
    isConnected: false,
    workingDir: null,
  });

  // Track canvas state for sending with messages
  const canvasStateRef = useRef<CanvasState | null>(initialCanvasState || null);
  const wsRef = useRef<WebSocket | null>(null);
  const currentMessageRef = useRef<string>('');
  const reconnectAttempts = useRef(0);
  const maxReconnectAttempts = 3;

  // Keep canvas state in sync with prop changes
  useEffect(() => {
    if (initialCanvasState) {
      canvasStateRef.current = initialCanvasState;
    }
  }, [initialCanvasState]);

  // Get canvas store for auto-refresh
  const { loadFromConfig } = useCanvasStore();

  // Connect to Claude Code WebSocket
  useEffect(() => {
    const connect = () => {
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      // In development, connect directly to backend (bypass Vite proxy for WebSockets)
      // Use port 8001 to match start-dashboard.bat
      const backendHost = import.meta.env.DEV ? 'localhost:8001' : window.location.host;

      // Use study-specific endpoint if studyId provided
      const wsUrl = studyId
        ? `${protocol}//${backendHost}/api/claude-code/ws/${encodeURIComponent(studyId)}`
        : `${protocol}//${backendHost}/api/claude-code/ws`;

      console.log('[ClaudeCode] Connecting to:', wsUrl);
      const ws = new WebSocket(wsUrl);

      ws.onopen = () => {
        console.log('[ClaudeCode] Connected');
        setState((prev) => ({ ...prev, isConnected: true, error: null }));
        reconnectAttempts.current = 0;

        // If no studyId in URL, send init message
        if (!studyId) {
          ws.send(JSON.stringify({ type: 'init', study_id: null }));
        }
      };

      ws.onclose = () => {
        console.log('[ClaudeCode] Disconnected');
        setState((prev) => ({ ...prev, isConnected: false }));

        // Attempt reconnection
        if (reconnectAttempts.current < maxReconnectAttempts) {
          reconnectAttempts.current++;
          console.log(`[ClaudeCode] Reconnecting... attempt ${reconnectAttempts.current}`);
          setTimeout(connect, 2000 * reconnectAttempts.current);
        }
      };

      ws.onerror = (event) => {
        console.error('[ClaudeCode] WebSocket error:', event);
        setState((prev) => ({ ...prev, isConnected: false }));
        onError?.('Claude Code connection error');
      };

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

      wsRef.current = ws;
    };

    connect();

    return () => {
      reconnectAttempts.current = maxReconnectAttempts; // Prevent reconnection on unmount
      wsRef.current?.close();
      wsRef.current = null;
    };
  }, [studyId]);

  // Handle WebSocket messages
  const handleWebSocketMessage = useCallback(
    (data: any) => {
      switch (data.type) {
        case 'initialized':
          console.log('[ClaudeCode] Session initialized:', data.session_id);
          setState((prev) => ({
            ...prev,
            sessionId: data.session_id,
            workingDir: data.working_dir || null,
          }));
          break;

        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 'done':
          setState((prev) => ({
            ...prev,
            isThinking: false,
            messages: prev.messages.map((msg, idx) =>
              idx === prev.messages.length - 1 && msg.role === 'assistant'
                ? { ...msg, isStreaming: false }
                : msg
            ),
          }));
          currentMessageRef.current = '';
          break;

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

        case 'refresh_canvas':
          // Claude made file changes - trigger canvas refresh
          console.log('[ClaudeCode] Canvas refresh requested:', data.reason);
          if (data.study_id) {
            onCanvasRefresh?.(data.study_id);
            reloadCanvasFromStudy(data.study_id);
          }
          break;

        case 'canvas_updated':
          console.log('[ClaudeCode] Canvas state updated');
          break;

        case 'pong':
          // Heartbeat response
          break;

        default:
          console.log('[ClaudeCode] Unknown message type:', data.type);
      }
    },
    [onError, onCanvasRefresh]
  );

  // Reload canvas from study config
  const reloadCanvasFromStudy = useCallback(
    async (studyIdToReload: string) => {
      try {
        console.log('[ClaudeCode] Reloading canvas for study:', studyIdToReload);

        // Fetch fresh config from backend
        const response = await fetch(`/api/optimization/studies/${encodeURIComponent(studyIdToReload)}/config`);
        if (!response.ok) {
          throw new Error(`Failed to fetch config: ${response.status}`);
        }

        const data = await response.json();
        const config = data.config; // API returns { config: ..., path: ..., study_id: ... }

        // Reload canvas with new config
        loadFromConfig(config);

        // Add system message about refresh
        const refreshMessage: Message = {
          id: `msg_${Date.now()}_refresh`,
          role: 'system',
          content: `Canvas refreshed with latest changes from ${studyIdToReload}`,
          timestamp: new Date(),
        };

        setState((prev) => ({
          ...prev,
          messages: [...prev.messages, refreshMessage],
        }));
      } catch (error) {
        console.error('[ClaudeCode] Failed to reload canvas:', error);
      }
    },
    [loadFromConfig]
  );

  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;

      if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
        onError?.('Not connected to Claude Code');
        return;
      }

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

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

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

      // Reset current message tracking
      currentMessageRef.current = '';

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

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

  // 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 (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(
        JSON.stringify({
          type: 'set_canvas',
          canvas_state: newCanvasState,
        })
      );
    }
  }, []);

  // Send ping to keep connection alive
  useEffect(() => {
    const pingInterval = setInterval(() => {
      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(JSON.stringify({ type: 'ping' }));
      }
    }, 30000); // Every 30 seconds

    return () => clearInterval(pingInterval);
  }, []);

  return {
    messages: state.messages,
    isThinking: state.isThinking,
    error: state.error,
    sessionId: state.sessionId,
    isConnected: state.isConnected,
    workingDir: state.workingDir,
    sendMessage,
    clearMessages,
    updateCanvasState,
    reloadCanvasFromStudy,
  };
}