claude-web-ui/frontend/src/hooks/useClaudeSession.js

import { useState, useRef, useCallback, useEffect } from 'react';

// Build WebSocket URL from current location
function getWsUrl() {
  if (import.meta.env.VITE_WS_URL) return import.meta.env.VITE_WS_URL;
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${protocol}//${window.location.host}/ws`;
}

const MAX_FILE_SIZE = 10 * 1024 * 1024; // 10MB

// Reconnect configuration
const RECONNECT_DELAY_BASE = 1000; // 1 second
const RECONNECT_DELAY_MAX = 30000; // 30 seconds
const RECONNECT_MAX_ATTEMPTS = 10;

export function useClaudeSession() {
  const [connected, setConnected] = useState(false);
  const [sessionActive, setSessionActive] = useState(false);
  const [messages, setMessages] = useState([]);
  const [currentProject, setCurrentProject] = useState(null);
  const [currentHost, setCurrentHost] = useState(null);
  const [isProcessing, setIsProcessing] = useState(false);
  const [error, setError] = useState(null);
  const [sessionId, setSessionId] = useState(null);

  // Session stats
  const [sessionStats, setSessionStats] = useState({
    totalCost: 0,
    inputTokens: 0,
    outputTokens: 0,
    cacheReadTokens: 0,
    cacheCreationTokens: 0,
    numTurns: 0,
    contextWindow: 200000, // Default, updated from modelUsage
    isCompacting: false,
  });

  // Permission mode and control protocol state
  // Load from localStorage, default to 'default'
  const [permissionMode, setPermissionMode] = useState(() => {
    try {
      return localStorage.getItem('claude-permission-mode') || 'default';
    } catch {
      return 'default';
    }
  });
  const [controlInitialized, setControlInitialized] = useState(false);
  const [availableModels, setAvailableModels] = useState([]);
  const [accountInfo, setAccountInfo] = useState(null);

  // Pending permission request (tool approval dialog)
  const [pendingPermission, setPendingPermission] = useState(null);

  // Note: respondedPlanApprovals removed - plan approval now handled via control protocol permission flow

  const wsRef = useRef(null);
  const currentAssistantMessage = useRef(null);
  const reconnectAttempts = useRef(0);
  const reconnectTimeout = useRef(null);
  const intentionalClose = useRef(false);

  const connect = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) return;

    const ws = new WebSocket(getWsUrl());
    wsRef.current = ws;

    ws.onopen = () => {
      console.log('WebSocket connected');
      setConnected(true);
      setError(null);
      reconnectAttempts.current = 0; // Reset on successful connection
    };

    ws.onclose = (event) => {
      console.log('WebSocket disconnected', { code: event.code, reason: event.reason, wasClean: event.wasClean });
      setConnected(false);
      setSessionActive(false);
      setIsProcessing(false);

      // Attempt reconnect unless it was intentional or auth failure
      if (!intentionalClose.current && event.code !== 1008) {
        if (reconnectAttempts.current < RECONNECT_MAX_ATTEMPTS) {
          const delay = Math.min(
            RECONNECT_DELAY_BASE * Math.pow(2, reconnectAttempts.current),
            RECONNECT_DELAY_MAX
          );
          reconnectAttempts.current++;
          console.log(`Reconnecting in ${delay}ms (attempt ${reconnectAttempts.current}/${RECONNECT_MAX_ATTEMPTS})`);

          reconnectTimeout.current = setTimeout(() => {
            console.log('Attempting reconnect...');
            connect();
          }, delay);
        } else {
          console.log('Max reconnect attempts reached');
          setError('Connection lost. Please refresh the page.');
        }
      }
    };

    ws.onerror = (err) => {
      console.error('WebSocket error:', err);
      setError('Connection error');
    };

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

  const handleMessage = useCallback((data) => {
    console.log('Received:', data.type, data);

    switch (data.type) {
      case 'session_started':
        setSessionActive(true);
        setCurrentProject(data.project);
        setSessionId(data.sessionId || `session-${Date.now()}`);
        setMessages(prev => [...prev, {
          type: 'system',
          content: `Session started in ${data.project}`,
          timestamp: data.timestamp
        }]);
        break;

      case 'session_ended':
        setSessionActive(false);
        setIsProcessing(false);
        setMessages(prev => [...prev, {
          type: 'system',
          content: `Session ended (code: ${data.code})`,
          timestamp: data.timestamp
        }]);
        break;

      case 'claude_event':
        handleClaudeEvent(data.event);
        break;

      case 'raw_output':
        console.log('Raw:', data.content);
        break;

      case 'stderr':
        console.log('Stderr:', data.content);
        break;

      case 'error':
        setError(data.message);
        setIsProcessing(false);
        break;

      case 'system':
        // Handle system messages from backend
        if (data.message) {
          setMessages(prev => [...prev, {
            type: 'system',
            content: data.message,
            timestamp: data.timestamp || Date.now()
          }]);
        }
        break;

      case 'generation_stopped':
        // User requested to stop - process killed and restarting
        setIsProcessing(false);
        if (data.message) {
          setMessages(prev => [...prev, {
            type: 'system',
            content: data.message,
            timestamp: data.timestamp || Date.now()
          }]);
        }
        break;

      case 'control_initialized':
        setControlInitialized(true);
        if (data.models) setAvailableModels(data.models);
        if (data.account) setAccountInfo(data.account);
        console.log('Control protocol initialized:', data);

        // Apply saved permission mode after control protocol is ready
        try {
          const savedMode = localStorage.getItem('claude-permission-mode');
          if (savedMode && savedMode !== 'default' && wsRef.current?.readyState === WebSocket.OPEN) {
            console.log('Applying saved permission mode:', savedMode);
            wsRef.current.send(JSON.stringify({
              type: 'set_permission_mode',
              mode: savedMode
            }));
          }
        } catch {}
        break;

      case 'permission_mode_changed':
        setPermissionMode(data.mode);
        try { localStorage.setItem('claude-permission-mode', data.mode); } catch {}
        setMessages(prev => [...prev, {
          type: 'system',
          content: `Mode changed to: ${data.mode}`,
          timestamp: data.timestamp
        }]);
        break;

      case 'permission_mode':
        setPermissionMode(data.mode);
        try { localStorage.setItem('claude-permission-mode', data.mode); } catch {}
        break;

      case 'control_error':
        setError(`Control error: ${data.error}`);
        break;

      case 'permission_request':
        console.log('Permission request:', data);
        setPendingPermission({
          requestId: data.requestId,
          toolName: data.toolName,
          toolInput: data.toolInput,
          permissionSuggestions: data.permissionSuggestions,
          blockedPath: data.blockedPath
        });
        break;
    }
  }, []);

  const handleClaudeEvent = useCallback((event) => {
    // Debug: log all event types to understand the structure
    console.log('Claude Event:', event.type, event);

    // Handle different Claude event types
    if (event.type === 'assistant') {
      setIsProcessing(true);

      // Check if this is a content block
      if (event.message?.content) {
        const newMessages = [];

        for (const block of event.message.content) {
          if (block.type === 'text' && block.text) {
            // Only add text if we don't have a streaming message with similar content
            // (to avoid duplicates from stream_event + final assistant event)
            newMessages.push({
              type: 'assistant',
              content: block.text,
              timestamp: Date.now(),
              final: true  // Mark as final message
            });
          } else if (block.type === 'tool_use') {
            // Tool use is embedded in assistant message content
            newMessages.push({
              type: 'tool_use',
              tool: block.name,
              input: block.input,
              toolUseId: block.id,
              timestamp: Date.now()
            });
          }
        }

        if (newMessages.length > 0) {
          setMessages(prev => {
            // Collect existing toolUseIds for deduplication
            const existingToolUseIds = new Set(
              prev.filter(m => m.type === 'tool_use' && m.toolUseId).map(m => m.toolUseId)
            );

            // Check if last message is a streaming message - if so, replace it with final
            const last = prev[prev.length - 1];
            if (last?.type === 'assistant' && last.streaming) {
              // Replace streaming message with final content, keep other new messages (tool_use)
              const textMessages = newMessages.filter(m => m.type === 'assistant');
              const otherMessages = newMessages.filter(m => m.type !== 'assistant');

              // Deduplicate tool_use messages by toolUseId
              const uniqueOtherMessages = otherMessages.filter(
                m => m.type !== 'tool_use' || !existingToolUseIds.has(m.toolUseId)
              );

              if (textMessages.length > 0) {
                return [
                  ...prev.slice(0, -1),
                  { ...textMessages[0], streaming: false },
                  ...uniqueOtherMessages
                ];
              }
            }

            // Deduplicate tool_use messages by toolUseId
            const uniqueNewMessages = newMessages.filter(
              m => m.type !== 'tool_use' || !existingToolUseIds.has(m.toolUseId)
            );

            return [...prev, ...uniqueNewMessages];
          });
        }
      }
    } else if (event.type === 'user') {
      // Tool results come as 'user' events with message.content containing tool_result blocks
      if (event.message?.content) {
        for (const block of event.message.content) {
          if (block.type === 'tool_result') {
            console.log('Tool result found:', block);
            setMessages(prev => [...prev, {
              type: 'tool_result',
              content: block.content,
              toolUseId: block.tool_use_id,
              isError: block.is_error || false,
              timestamp: Date.now()
            }]);
          }
        }
      }
    } else if (event.type === 'content_block_delta') {
      // Streaming delta (direct)
      if (event.delta?.text) {
        setMessages(prev => {
          const last = prev[prev.length - 1];
          if (last?.type === 'assistant' && last.streaming) {
            return [
              ...prev.slice(0, -1),
              { ...last, content: last.content + event.delta.text }
            ];
          }
          return [...prev, {
            type: 'assistant',
            content: event.delta.text,
            streaming: true,
            timestamp: Date.now()
          }];
        });
      }
    } else if (event.type === 'stream_event' && event.event?.type === 'content_block_delta') {
      // Streaming delta (wrapped in stream_event)
      const deltaText = event.event?.delta?.text;
      if (deltaText) {
        setMessages(prev => {
          const last = prev[prev.length - 1];
          if (last?.type === 'assistant' && last.streaming) {
            return [
              ...prev.slice(0, -1),
              { ...last, content: last.content + deltaText }
            ];
          }
          return [...prev, {
            type: 'assistant',
            content: deltaText,
            streaming: true,
            timestamp: Date.now()
          }];
        });
      }
    } else if (event.type === 'result') {
      // Final result - extract stats and stop processing
      setIsProcessing(false);

      // Update session stats from result event
      // Claude sends detailed modelUsage with per-model stats including contextWindow
      // usage contains: input_tokens (new), cache_read_input_tokens, cache_creation_input_tokens, output_tokens
      if (event.usage || event.modelUsage || event.total_cost_usd !== undefined) {
        // Get the primary model's usage (usually claude-opus or claude-sonnet)
        const modelUsage = event.modelUsage || {};
        const primaryModel = Object.keys(modelUsage).find(k => k.includes('opus') || k.includes('sonnet'))
          || Object.keys(modelUsage)[0];
        const primaryUsage = primaryModel ? modelUsage[primaryModel] : null;

        // Calculate effective input tokens (what's actually in context)
        // This is: new input + cache read (cache creation doesn't count towards context limit)
        const effectiveInput = (event.usage?.input_tokens || 0) +
                               (event.usage?.cache_read_input_tokens || 0);

        setSessionStats(prev => ({
          totalCost: event.total_cost_usd ?? prev.totalCost,
          // Use modelUsage for accurate per-turn counts, fallback to usage
          inputTokens: primaryUsage?.inputTokens ?? effectiveInput ?? prev.inputTokens,
          outputTokens: primaryUsage?.outputTokens ?? event.usage?.output_tokens ?? prev.outputTokens,
          cacheReadTokens: primaryUsage?.cacheReadInputTokens ?? event.usage?.cache_read_input_tokens ?? prev.cacheReadTokens,
          cacheCreationTokens: primaryUsage?.cacheCreationInputTokens ?? event.usage?.cache_creation_input_tokens ?? prev.cacheCreationTokens,
          numTurns: event.num_turns ?? prev.numTurns,
          contextWindow: primaryUsage?.contextWindow ?? prev.contextWindow,
          isCompacting: false,
        }));
      }
    } else if (event.type === 'system' && event.subtype === 'result') {
      setIsProcessing(false);
    } else if (event.type === 'system' && event.message?.includes?.('compact')) {
      // Detect compacting
      setSessionStats(prev => ({ ...prev, isCompacting: true }));
    }
  }, []);

  const loadHistory = useCallback(async (project, host = null) => {
    try {
      const encodedProject = encodeURIComponent(project);
      const hostParam = host ? `?host=${host}` : '';
      const response = await fetch(`/api/history/${encodedProject}${hostParam}`, {
        credentials: 'include',
      });
      if (response.ok) {
        const data = await response.json();
        if (data.messages && data.messages.length > 0) {
          console.log(`Loaded ${data.messages.length} messages from history (source: ${data.source || 'local'})`);
          setMessages(data.messages);
          return data.sessionId;
        }
      }
    } catch (err) {
      console.error('Failed to load history:', err);
    }
    return null;
  }, []);

  const startSession = useCallback(async (project = '/projects', resume = true, host = null) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    // Track which host we're connecting to
    setCurrentHost(host);

    // Load history before starting session if resuming
    if (resume) {
      await loadHistory(project, host);
    }

    wsRef.current.send(JSON.stringify({
      type: 'start_session',
      project,
      resume,
      host
    }));
  }, [loadHistory]);

  // Upload files to server
  const uploadFiles = useCallback(async (files) => {
    if (!files || files.length === 0) return [];

    const formData = new FormData();
    for (const fileData of files) {
      formData.append('files', fileData.file);
    }

    try {
      const response = await fetch(`/api/upload/${sessionId}`, {
        method: 'POST',
        body: formData,
        credentials: 'include',
      });

      if (!response.ok) {
        const error = await response.json();
        throw new Error(error.error || 'Upload failed');
      }

      const data = await response.json();
      return data.files;
    } catch (err) {
      console.error('Upload error:', err);
      setError(`Upload failed: ${err.message}`);
      return [];
    }
  }, [sessionId]);

  const sendMessage = useCallback(async (message, attachedFiles = []) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    if (!sessionActive) {
      setError('No active session');
      return;
    }

    let uploadedFiles = [];
    let attachmentInfo = null;

    // Upload files if any
    if (attachedFiles.length > 0) {
      setIsProcessing(true);
      uploadedFiles = await uploadFiles(attachedFiles);

      if (uploadedFiles.length > 0) {
        // Create attachment info for the message
        attachmentInfo = uploadedFiles.map(f => ({
          path: f.path,
          name: f.originalName,
          type: f.mimeType,
          isImage: f.isImage
        }));
      }
    }

    // Build the prompt with attachment info
    let finalMessage = message;
    if (uploadedFiles.length > 0) {
      const attachmentPrefix = uploadedFiles
        .map(f => `${f.path} (${f.mimeType})`)
        .join(', ');
      finalMessage = `[Attachments: ${attachmentPrefix}]\n\n${message}`;
    }

    // Add user message to display (with attachment badge info)
    setMessages(prev => [...prev, {
      type: 'user',
      content: message,
      attachments: attachmentInfo,
      timestamp: Date.now()
    }]);

    setIsProcessing(true);

    wsRef.current.send(JSON.stringify({
      type: 'user_message',
      message: finalMessage
    }));
  }, [sessionActive, uploadFiles]);

  const stopSession = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: 'stop_session' }));
    }
    setCurrentHost(null);
  }, []);

  // Stop current generation (interrupt Claude)
  const stopGeneration = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: 'stop_generation' }));
      setIsProcessing(false);
    }
  }, []);

  const clearMessages = useCallback(() => {
    setMessages([]);
  }, []);

  // Respond to a permission request (allow/deny tool execution)
  const respondToPermission = useCallback((requestId, allow, message = null) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    wsRef.current.send(JSON.stringify({
      type: 'permission_response',
      requestId,
      allow,
      message
    }));

    // Clear the pending permission
    setPendingPermission(null);
  }, []);

  // Note: respondToPlanApproval removed - plan approval now handled via permission_response with isPlanApproval flag

  // Change permission mode during session
  const changePermissionMode = useCallback((mode) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    if (!sessionActive) {
      setError('No active session');
      return;
    }

    if (!controlInitialized) {
      setError('Control protocol not initialized');
      return;
    }

    wsRef.current.send(JSON.stringify({
      type: 'set_permission_mode',
      mode
    }));
  }, [sessionActive, controlInitialized]);

  // Auto-connect on mount
  useEffect(() => {
    intentionalClose.current = false;
    connect();
    return () => {
      intentionalClose.current = true;
      if (reconnectTimeout.current) {
        clearTimeout(reconnectTimeout.current);
      }
      wsRef.current?.close();
    };
  }, [connect]);

  return {
    connected,
    sessionActive,
    sessionId,
    messages,
    currentProject,
    currentHost,
    isProcessing,
    error,
    sessionStats,
    permissionMode,
    controlInitialized,
    availableModels,
    accountInfo,
    pendingPermission,
    connect,
    startSession,
    sendMessage,
    stopSession,
    stopGeneration,
    clearMessages,
    changePermissionMode,
    respondToPermission,
    setError,
    setMessages
  };
}