import { useState, useEffect, useRef, useContext, useCallback } from 'react';
import { useParams, useNavigate } from 'react-router-dom';
import { socketClient } from '../lib/socket';
import { chatSessionApi, projectApi, providerApi, type ChatSession, type ChatTurn, type ChatTurnBlock, ChatTurnStats, ChatSessionMode, type AiProvider, type Project } from '../lib/api';
import { WorkspaceMode } from '../lib/types';
import SessionPanel from '../components/SessionPanel';
import ProjectPanel from '../components/ProjectPanel';
import FilesPanel from '../components/FilesPanel';
import LogPanel from '../components/LogPanel';
import ChatTurnComponent from '../components/ChatTurn';
import EditorPanel from '../components/EditorPanel';
import ErrorBoundary from '../components/ErrorBoundary';
import { AppContext } from '../App';

interface LogEntry {
  id: string;
  timestamp: Date;
  level: string;
  component: string;
  message: string;
  metadata?: unknown;
}

interface StreamingState {
  currentMode: 'thinking' | 'responding' | null;
  thinkingContent: string;
  respondingContent: string;
  currentBlockIndex: number | null;
}

interface SubagentStreamState {
  agentId: string;
  name: string;
  parameters: string;
  thinking: string;
  response: string;
  toolCalls: Array<{
    callId: string;
    name: string;
    parameters?: string;
    response?: string;
  }>;
  stats?: ChatTurnStats;
}

export default function ChatSessionView() {
  const { projectId, sessionId } = useParams<{ projectId: string; sessionId: string }>();
  const navigate = useNavigate();
  const appContext = useContext(AppContext);

  const [project, setProject] = useState<Project | null>(null);
  const [session, setSession] = useState<ChatSession | null>(null);
  const [turns, setTurns] = useState<ChatTurn[]>([]);
  const [promptInput, setPromptInput] = useState('');
  const [isProcessing, setIsProcessing] = useState(false);
  const [isAborting, setIsAborting] = useState(false);
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState('');
  const [sessionLocked, setSessionLocked] = useState(true);
  const [workspaceMode, setWorkspaceMode] = useState<WorkspaceMode>(WorkspaceMode.Idle);
  const [isUpdatingMode, setIsUpdatingMode] = useState(false);
  const [isUpdatingProvider, setIsUpdatingProvider] = useState(false);
  const [isUpdatingReasoning, setIsUpdatingReasoning] = useState(false);
  const [logExpanded, setLogExpanded] = useState(false);
  const [logs, setLogs] = useState<LogEntry[]>([]);
  const [toast, setToast] = useState<string | null>(null);
  const [providers, setProviders] = useState<AiProvider[]>([]);
  const [selectedProviderId, setSelectedProviderId] = useState<string>('');
  const [selectedModelId, setSelectedModelId] = useState<string>('');
  const [sessionReasoningEffort, setSessionReasoningEffort] = useState<string>('off');
  const [sessionNumCtx, setSessionNumCtx] = useState<number>(131072);
  const [isUpdatingNumCtx, setIsUpdatingNumCtx] = useState(false);
  const [isEditingProvider, setIsEditingProvider] = useState(false);
  const [editProviderId, setEditProviderId] = useState('');
  const [editModelId, setEditModelId] = useState('');
  const [isEditingName, setIsEditingName] = useState(false);
  const [editName, setEditName] = useState('');
  const [isUpdatingName, setIsUpdatingName] = useState(false);
  const [_connectionState, setConnectionState] = useState<'disconnected' | 'connecting' | 'connected' | 'error'>('disconnected');
  const [isOtherTab, setIsOtherTab] = useState(false);
  const [reconnectAttempts, setReconnectAttempts] = useState(0);
  const [editorFilePath, setEditorFilePath] = useState<string | undefined>(undefined);

  const messagesEndRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLTextAreaElement>(null);
  const toastTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const pendingUpdatesRef = useRef<Map<string, Partial<ChatTurn>>>(new Map());
  const updateRafRef = useRef<number | null>(null);
  const currentTurnIdRef = useRef<string | null>(null);
  const streamingStateRef = useRef<Map<string, StreamingState>>(new Map());
  const escTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const escFlagRef = useRef(false);
  const subagentStateRef = useRef<Map<string, SubagentStreamState>>(new Map());
  const sessionRef = useRef<ChatSession | null>(null);
  const projectRef = useRef<Project | null>(null);

  useEffect(() => {
    loadSessionData();
  }, [projectId, sessionId]);

  useEffect(() => {
    setupSocketListeners();
    return () => cleanupSocketListeners();
  }, [sessionId]);

  useEffect(() => {
    scrollToBottom();
  }, [turns]);

  useEffect(() => {
    return () => {
      if (toastTimerRef.current) {
        clearTimeout(toastTimerRef.current);
      }
      if (updateRafRef.current) {
        cancelAnimationFrame(updateRafRef.current);
      }
    };
  }, []);

  // Keep refs in sync with state for use in cleanup closures
  useEffect(() => {
    sessionRef.current = session;
  }, [session]);

  useEffect(() => {
    projectRef.current = project;
  }, [project]);

  // Start heartbeat when session+project are loaded
  useEffect(() => {
    if (session && project) {
      socketClient.startSessionHeartbeat();
    }
    return () => {
      socketClient.stopSessionHeartbeat();
    };
  }, [session, project]);

  // Handle socket reconnection and tab lock
  const handleSocketConnect = useCallback(() => {
    setConnectionState('connected');
    appContext?.setStatusMessage('Connected');
    
    // If we had a processing turn, backend will auto-reconnect
    const hasProcessingTurn = turns.some(t => t.status === 'processing');
    if (hasProcessingTurn) {
      appContext?.setStatusMessage('Reconnecting to active session...');
    }
  }, [turns, appContext]);

  const handleTabLockDenied = useCallback((data: { message: string }) => {
    setIsOtherTab(true);
    setConnectionState('error');
    appContext?.setStatusMessage(data.message);
  }, [appContext]);

  const handleReconnectAttempt = useCallback(() => {
    setReconnectAttempts(prev => prev + 1);
    appContext?.setStatusMessage(`Reconnecting... (${reconnectAttempts + 1})`);
  }, [reconnectAttempts, appContext]);

  const handleReconnectFailed = useCallback(() => {
    setConnectionState('error');
    appContext?.setStatusMessage('Reconnection failed');
  }, [appContext]);

  const handleReconnect = useCallback(() => {
    setConnectionState('connected');
    setReconnectAttempts(0);
    appContext?.setStatusMessage('Reconnected');
  }, [appContext]);

  // Release session lock on unmount only
  useEffect(() => {
    return () => {
      const droneJson = localStorage.getItem('dtp_drone_registration');
      if (droneJson && sessionRef.current && projectRef.current) {
        try {
          const registration = JSON.parse(droneJson);
          socketClient.releaseSessionLock(
            registration,
            projectRef.current,
            sessionRef.current,
          );
        } catch (err) {
          console.error('Failed to release session lock', err);
        }
      }
    };
  }, []);

  const loadSessionData = async () => {
    try {
      if (sessionId) {
        const sessionData = await chatSessionApi.get(sessionId);
        setSession(sessionData);

        const projectRef = sessionData.project;
        const projectObjectId = typeof projectRef === 'string' ? projectRef : projectRef?._id;
        if (projectObjectId) {
          const projectData = await projectApi.get(projectObjectId);
          setProject(projectData);
        }

        const turnsData = await chatSessionApi.getTurns(sessionId);
        setTurns(turnsData);
        setSessionLocked(true);

        const allProviders = await providerApi.getAll();
        setProviders(allProviders.sort((a, b) => a.name.localeCompare(b.name)));

        const providerId = typeof sessionData.provider === 'string' 
          ? sessionData.provider 
          : sessionData.provider?._id;
        setSelectedProviderId(providerId || '');
        setSelectedModelId(sessionData.selectedModel || '');
        setSessionReasoningEffort(sessionData.reasoningEffort || 'off');

        // Initialize numCtx: persisted value wins, otherwise model's contextWindow
        if (sessionData.numCtx) {
          setSessionNumCtx(sessionData.numCtx);
        } else {
          const provider = allProviders.find(p => p._id === providerId);
          const model = provider?.models.find(m => m.id === sessionData.selectedModel);
          setSessionNumCtx(model?.contextWindow || 131072);
        }
        
        // Set connection state to connecting - will update on socket connect
        setConnectionState('connecting');
        appContext?.setStatusMessage('Connecting...');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to load session');
      setConnectionState('error');
    } finally {
      setLoading(false);
    }
  };

  const handleSessionUpdated = useCallback((updates: Partial<ChatSession>) => {
    setSession(prev => prev ? { ...prev, ...updates } : null);
  }, []);

  const setupSocketListeners = () => {
    socketClient.on('thinking', handleThinking);
    socketClient.on('response', handleResponse);
    socketClient.on('toolCall', handleToolCall);
    socketClient.on('workOrderComplete', handleWorkOrderComplete);
    socketClient.on('workspaceModeChanged', handleWorkspaceModeChanged);
    socketClient.on('sessionUpdated', handleSessionUpdated);
    socketClient.on('log:entry', handleLogEntry);
    socketClient.on('status', handleStatus);
    socketClient.on('agent:thinking', handleAgentThinking);
    socketClient.on('agent:response', handleAgentResponse);
    socketClient.on('agent:tool-call', handleAgentToolCall);
    socketClient.on('agent:tool-result', handleAgentToolResult);
    socketClient.on('agent:complete', handleAgentComplete);
    socketClient.on('connect', handleSocketConnect);
    socketClient.on('tabLockDenied', handleTabLockDenied);
    socketClient.on('reconnect_attempt', handleReconnectAttempt);
    socketClient.on('reconnect_failed', handleReconnectFailed);
    socketClient.on('reconnect', handleReconnect);
  };

  const cleanupSocketListeners = () => {
    socketClient.off('thinking', handleThinking);
    socketClient.off('response', handleResponse);
    socketClient.off('toolCall', handleToolCall);
    socketClient.off('workOrderComplete', handleWorkOrderComplete);
    socketClient.off('workspaceModeChanged', handleWorkspaceModeChanged);
    socketClient.off('sessionUpdated', handleSessionUpdated);
    socketClient.off('log:entry', handleLogEntry);
    socketClient.off('status', handleStatus);
    socketClient.off('agent:thinking', handleAgentThinking);
    socketClient.off('agent:response', handleAgentResponse);
    socketClient.off('agent:tool-call', handleAgentToolCall);
    socketClient.off('agent:tool-result', handleAgentToolResult);
    socketClient.off('agent:complete', handleAgentComplete);
    socketClient.off('connect', handleSocketConnect);
    socketClient.off('tabLockDenied', handleTabLockDenied);
    socketClient.off('reconnect_attempt', handleReconnectAttempt);
    socketClient.off('reconnect_failed', handleReconnectFailed);
    socketClient.off('reconnect', handleReconnect);
  };

  const scheduleUpdate = useCallback(() => {
    if (updateRafRef.current) return;
    
    updateRafRef.current = requestAnimationFrame(() => {
      const pendingEntries = Array.from(pendingUpdatesRef.current.entries());
      if (pendingEntries.length === 0) {
        updateRafRef.current = null;
        return;
      }

      setTurns(prevTurns => {
        const newTurns = [...prevTurns];
        let changed = false;

        for (const [turnId, turnUpdates] of pendingEntries) {
          const turnIndex = newTurns.findIndex(t => t._id === turnId);
          if (turnIndex === -1) continue;

            const oldTurn = newTurns[turnIndex]!;
            const newTurn = { ...oldTurn, ...turnUpdates };
          
          if (turnUpdates.blocks !== undefined) {
            const state = streamingStateRef.current.get(turnId);
            const updatedBlocks = [...(oldTurn!.blocks || [])];

            for (const updateBlock of turnUpdates.blocks) {
              let blockIndex: number | null = state?.currentBlockIndex ?? null;

              if (updateBlock.mode === 'tool') {
                // For tool blocks with callId, find and replace existing
                if (updateBlock.content?.callId) {
                  const existingIndex = updatedBlocks.findIndex(
                    b => b.mode === 'tool' && b.content?.callId === updateBlock.content?.callId
                  );
                  blockIndex = existingIndex !== -1 ? existingIndex : null;
                } else {
                  blockIndex = null;
                }
              } else if (
                blockIndex === null ||
                updatedBlocks[blockIndex]?.mode !== updateBlock.mode
              ) {
                const lastIndex = updatedBlocks.length - 1;
                blockIndex = updatedBlocks[lastIndex]?.mode === updateBlock.mode
                  ? lastIndex
                  : null;
              }

              if (blockIndex !== null) {
                updatedBlocks[blockIndex] = updateBlock;
              } else {
                updatedBlocks.push(updateBlock);
                blockIndex = updatedBlocks.length - 1;
              }

              if (state) {
                state.currentBlockIndex = blockIndex;
              }
            }

            newTurn.blocks = updatedBlocks;
          }
          if (turnUpdates.toolCalls !== undefined) {
            newTurn.toolCalls = [...(oldTurn!.toolCalls || []), ...turnUpdates.toolCalls];
            newTurn.stats = {
              ...oldTurn!.stats,
              toolCallCount: newTurn.toolCalls.length
            };
          }
          if (turnUpdates.status !== undefined) {
            newTurn.status = turnUpdates.status;
          }

          newTurns[turnIndex] = newTurn;
          changed = true;
        }

        return changed ? newTurns : prevTurns;
      });

      pendingUpdatesRef.current.clear();
      updateRafRef.current = null;
    });
  }, []);

  const mergePendingUpdate = useCallback((turnId: string, updates: Partial<ChatTurn>) => {
    const existing = pendingUpdatesRef.current.get(turnId);
    const merged: Partial<ChatTurn> = {
      ...existing,
      ...updates,
    };

    if (existing?.blocks || updates.blocks) {
      const existingBlocks = [...(existing?.blocks || [])];
      const updateBlocks = updates.blocks || [];

      // Deduplicate tool blocks by callId: replace existing tool blocks
      // that have the same callId as an incoming update block
      merged.blocks = [...existingBlocks];
      for (const block of updateBlocks) {
        if (block.mode === 'tool' && block.content?.callId) {
          const existingIdx = merged.blocks.findIndex(
            b => b.mode === 'tool' && b.content?.callId === block.content?.callId
          );
          if (existingIdx !== -1) {
            merged.blocks[existingIdx] = block;
          } else {
            merged.blocks.push(block);
          }
        } else {
          merged.blocks.push(block);
        }
      }
    }

    if (existing?.toolCalls || updates.toolCalls) {
      merged.toolCalls = [
        ...(existing?.toolCalls || []),
        ...(updates.toolCalls || []),
      ];
    }

    pendingUpdatesRef.current.set(turnId, merged);
  }, []);

  const handleThinking = useCallback((content: string) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    let state = streamingStateRef.current.get(turnId);
    if (!state) {
      state = { currentMode: null, thinkingContent: '', respondingContent: '', currentBlockIndex: null };
      streamingStateRef.current.set(turnId, state);
    }

    // Check for mode transition
    if (state.currentMode !== 'thinking') {
      // Flush previous mode
      if (state.currentMode === 'responding' && state.respondingContent) {
        mergePendingUpdate(turnId, {
          blocks: [{
            mode: 'responding' as const,
            createdAt: new Date().toISOString(),
            content: state.respondingContent,
          }],
        });
        state.respondingContent = '';
        state.currentBlockIndex = null;
      }
      state.currentMode = 'thinking';
    }

    // Aggregate content
    state.thinkingContent += content;
    
    // Update with aggregated content
    mergePendingUpdate(turnId, {
      blocks: [{
        mode: 'thinking' as const,
        createdAt: new Date().toISOString(),
        content: state.thinkingContent,
      }],
    });
    scheduleUpdate();
  }, [mergePendingUpdate, scheduleUpdate]);

  const handleResponse = useCallback((content: string) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    let state = streamingStateRef.current.get(turnId);
    if (!state) {
      state = { currentMode: null, thinkingContent: '', respondingContent: '', currentBlockIndex: null };
      streamingStateRef.current.set(turnId, state);
    }

    // Check for mode transition
    if (state.currentMode !== 'responding') {
      // Flush previous mode
      if (state.currentMode === 'thinking' && state.thinkingContent) {
        mergePendingUpdate(turnId, {
          blocks: [{
            mode: 'thinking' as const,
            createdAt: new Date().toISOString(),
            content: state.thinkingContent,
          }],
        });
        state.thinkingContent = '';
        state.currentBlockIndex = null;
      }
      state.currentMode = 'responding';
    }

    // Aggregate content
    state.respondingContent += content;
    
    // Update with aggregated content
    mergePendingUpdate(turnId, {
      blocks: [{
        mode: 'responding' as const,
        createdAt: new Date().toISOString(),
        content: state.respondingContent,
      }],
    });
    scheduleUpdate();
  }, [mergePendingUpdate, scheduleUpdate]);

  const handleToolCall = useCallback((callId: string, name: string, params: string, response: string) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    // Flush current streaming state
    const state = streamingStateRef.current.get(turnId);
    if (state) {
      const blocksToFlush: ChatTurnBlock[] = [];
      
      if (state.currentMode === 'thinking' && state.thinkingContent) {
        blocksToFlush.push({
          mode: 'thinking' as const,
          createdAt: new Date().toISOString(),
          content: state.thinkingContent,
        });
        state.thinkingContent = '';
      }
      if (state.currentMode === 'responding' && state.respondingContent) {
        blocksToFlush.push({
          mode: 'responding' as const,
          createdAt: new Date().toISOString(),
          content: state.respondingContent,
        });
        state.respondingContent = '';
      }
      
      if (blocksToFlush.length > 0) {
        mergePendingUpdate(turnId, {
          blocks: blocksToFlush,
        });
        scheduleUpdate();
      }
      
      state.currentMode = null;
      state.currentBlockIndex = null;
    }

    // Initialize subagent state for subagent tool calls
    if (name === 'subagent') {
      subagentStateRef.current.set(callId, {
        agentId: callId,
        name,
        parameters: params,
        thinking: '',
        response: '',
        toolCalls: [],
      });
    }

    // Add tool block
    mergePendingUpdate(turnId, {
      blocks: [{
        mode: 'tool' as const,
        createdAt: new Date().toISOString(),
        content: { callId, name, parameters: params, response },
      }],
      toolCalls: [{ callId, name, parameters: params, response }],
    });
    scheduleUpdate();
  }, [mergePendingUpdate, scheduleUpdate]);

  const showToast = useCallback((message: string) => {
    setToast(message);
    if (toastTimerRef.current) {
      clearTimeout(toastTimerRef.current);
    }
    toastTimerRef.current = setTimeout(() => {
      setToast(null);
    }, 4000);
  }, []);

  const handleWorkOrderComplete = useCallback((turnId: string, success: boolean, message?: string) => {
    // Backend has already flushed and persisted all streaming content
    // Just clean up frontend streaming state and update status
    if (streamingStateRef.current.has(turnId)) {
      if (pendingUpdatesRef.current.has(turnId) || updateRafRef.current) {
        requestAnimationFrame(() => streamingStateRef.current.delete(turnId));
      } else {
        streamingStateRef.current.delete(turnId);
      }
    }

    // Clean up subagent state for this turn
    const subagentEntries = Array.from(subagentStateRef.current.entries());
    for (const [agentId] of subagentEntries) {
      subagentStateRef.current.delete(agentId);
    }

    if (success && message === 'aborted') {
      setTurns(prevTurns =>
        prevTurns.map(turn =>
          turn._id === turnId
            ? { ...turn, status: 'aborted', errorMessage: 'The turn was aborted by you.' }
            : turn
        )
      );
      showToast('The turn was aborted by you.');
    } else {
      setTurns(prevTurns =>
        prevTurns.map(turn =>
          turn._id === turnId
            ? { ...turn, status: success ? 'finished' : 'error', errorMessage: message && !success ? message : turn.errorMessage }
            : turn
        )
      );
      if (!success) {
        setError(message || 'Work order failed');
      }
    }

    // Clean up abort state
    if (escTimerRef.current) {
      clearTimeout(escTimerRef.current);
      escTimerRef.current = null;
    }
    escFlagRef.current = false;
    setIsAborting(false);
    setIsProcessing(false);
    currentTurnIdRef.current = null;
  }, [showToast]);

  const handleWorkspaceModeChanged = useCallback((mode: string) => {
    setWorkspaceMode(mode as WorkspaceMode);
  }, []);

  const handleLogEntry = useCallback((data: { level: string; component: string; message: string; metadata?: unknown; timestamp: number }) => {
    setLogs(prev => [
      ...prev,
      {
        id: `log-${Date.now()}-${Math.random()}`,
        timestamp: new Date(data.timestamp),
        level: data.level,
        component: data.component,
        message: data.message,
        metadata: data.metadata,
      },
    ]);
  }, []);

  const handleStatus = useCallback((content: string) => {
    appContext?.setStatusMessage(content);
  }, [appContext]);

  const defaultSubagentStats: ChatTurnStats = {
    toolCallCount: 0,
    inputTokens: 0,
    thinkingTokenCount: 0,
    responseTokens: 0,
    durationMs: 0,
    durationLabel: '0ms',
  };

  const buildSubagentBlock = useCallback((subState: SubagentStreamState): ChatTurnBlock => ({
    mode: 'tool' as const,
    createdAt: new Date().toISOString(),
    content: {
      callId: subState.agentId,
      name: subState.name,
      parameters: subState.parameters,
      response: subState.response,
      subagent: {
        agentId: subState.agentId,
        thinking: subState.thinking,
        response: subState.response,
        toolCalls: subState.toolCalls,
        stats: subState.stats || defaultSubagentStats,
      },
    },
  }), []);

  const updateSubagentBlock = useCallback((turnId: string, agentId: string) => {
    const subState = subagentStateRef.current.get(agentId);
    if (!subState) return;
    mergePendingUpdate(turnId, {
      blocks: [buildSubagentBlock(subState)],
    });
    scheduleUpdate();
  }, [mergePendingUpdate, scheduleUpdate, buildSubagentBlock]);

  const handleAgentThinking = useCallback((data: { agentId: string; thinking: string }) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    let subState = subagentStateRef.current.get(data.agentId);
    if (!subState) {
      subState = { agentId: data.agentId, name: '', parameters: '', thinking: '', response: '', toolCalls: [] };
      subagentStateRef.current.set(data.agentId, subState);
    }
    subState.thinking += data.thinking;
    updateSubagentBlock(turnId, data.agentId);
  }, [updateSubagentBlock]);

  const handleAgentResponse = useCallback((data: { agentId: string; chunk: string }) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    let subState = subagentStateRef.current.get(data.agentId);
    if (!subState) {
      subState = { agentId: data.agentId, name: '', parameters: '', thinking: '', response: '', toolCalls: [] };
      subagentStateRef.current.set(data.agentId, subState);
    }
    subState.response += data.chunk;
    updateSubagentBlock(turnId, data.agentId);
  }, [updateSubagentBlock]);

  const handleAgentToolCall = useCallback((data: { agentId: string; tool: string; args: unknown }) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    let subState = subagentStateRef.current.get(data.agentId);
    if (!subState) {
      subState = { agentId: data.agentId, name: '', parameters: '', thinking: '', response: '', toolCalls: [] };
      subagentStateRef.current.set(data.agentId, subState);
    }
    const argsStr = typeof data.args === 'string' ? data.args : JSON.stringify(data.args);
    subState.toolCalls.push({ callId: `tc_${Date.now()}_${subState.toolCalls.length}`, name: data.tool, parameters: argsStr });
    updateSubagentBlock(turnId, data.agentId);
  }, [updateSubagentBlock]);

  const handleAgentToolResult = useCallback((data: { agentId: string; tool: string; result: unknown }) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    const subState = subagentStateRef.current.get(data.agentId);
    if (!subState) return;

    const lastToolCall = subState.toolCalls[subState.toolCalls.length - 1];
    if (lastToolCall && lastToolCall.name === data.tool) {
      lastToolCall.response = typeof data.result === 'string' ? data.result : JSON.stringify(data.result);
    }
    updateSubagentBlock(turnId, data.agentId);
  }, [updateSubagentBlock]);

  const handleAgentComplete = useCallback((data: { agentId: string; response?: string; subagent?: Record<string, unknown>; stats?: ChatTurnStats }) => {
    const turnId = currentTurnIdRef.current;
    if (!turnId) return;

    const subState = subagentStateRef.current.get(data.agentId);
    if (!subState) return;

    // Apply final response if provided
    if (data.response) {
      subState.response = data.response;
    }

    // Apply stats from agent:complete
    if (data.stats) {
      subState.stats = data.stats;
    }

    // If subagent data was passed directly (e.g., from agent:complete with full payload),
    // apply the thinking and toolCalls from it
    if (data.subagent) {
      const sa = data.subagent as Record<string, unknown>;
      if (typeof sa.thinking === 'string') subState.thinking = sa.thinking;
      if (typeof sa.response === 'string') subState.response = sa.response;
      if (Array.isArray(sa.toolCalls)) subState.toolCalls = sa.toolCalls as SubagentStreamState['toolCalls'];
      if (sa.stats) subState.stats = { ...defaultSubagentStats, ...(sa.stats as Record<string, unknown>) } as ChatTurnStats;
    }

    updateSubagentBlock(turnId, data.agentId);

    // Clean up subagent state after final update
    requestAnimationFrame(() => {
      subagentStateRef.current.delete(data.agentId);
    });
  }, [updateSubagentBlock]);

  const handleWorkspaceModeChange = async (mode: WorkspaceMode) => {
    if (!session || !project) return;

    if (mode === workspaceMode) return;

    try {
      const droneJson = localStorage.getItem('dtp_drone_registration');
      if (!droneJson) {
        showToast('No drone registration found');
        return;
      }
      const registration = JSON.parse(droneJson);
      const result = await socketClient.requestWorkspaceMode(
        registration,
        project,
        session,
        mode,
      );
      if (!result.success) {
        showToast(result.reason || `Cannot switch to ${mode} mode`);
        setWorkspaceMode(workspaceMode);
      }
    } catch (err) {
      showToast(`Failed to change workspace mode: ${err instanceof Error ? err.message : 'Unknown error'}`);
    }
  };

  const handleModeChange = async (e: React.ChangeEvent<HTMLSelectElement>) => {
    if (!session) return;

    const newMode = e.target.value as ChatSessionMode;
    if (newMode === session.mode) return;

    setIsUpdatingMode(true);
    try {
      const updatedSession = await chatSessionApi.update(session._id, { mode: newMode });
      setSession(updatedSession);
      showToast(`Session mode changed to ${newMode}`);
    } catch (err) {
      showToast(`Failed to change mode: ${err instanceof Error ? err.message : 'Unknown error'}`);
    } finally {
      setIsUpdatingMode(false);
    }
  };

  const getSelectedModel = useCallback(() => {
    const provider = providers.find(p => p._id === selectedProviderId);
    if (!provider) return null;
    return provider.models.find(m => m.id === selectedModelId) || null;
  }, [providers, selectedProviderId, selectedModelId]);

  const getSelectedModelCapabilities = useCallback(() => {
    return getSelectedModel()?.capabilities || null;
  }, [getSelectedModel]);

  const handleReasoningChange = async (e: React.ChangeEvent<HTMLSelectElement>) => {
    if (!session) return;

    const newValue = e.target.value;
    if (newValue === sessionReasoningEffort) return;

    setIsUpdatingReasoning(true);
    try {
      const updatedSession = await chatSessionApi.update(session._id, {
        reasoningEffort: newValue
      });
      setSession(updatedSession);
      setSessionReasoningEffort(newValue);
      showToast(`Reasoning effort set to ${newValue}`);
    } catch (err) {
      showToast(`Failed to change reasoning effort: ${err instanceof Error ? err.message : 'Unknown error'}`);
    } finally {
      setIsUpdatingReasoning(false);
    }
  };

  const handleNumCtxChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    setSessionNumCtx(parseInt(e.target.value, 10));
  };

  const handleNumCtxCommit = async (e: React.MouseEvent<HTMLInputElement> | React.TouchEvent<HTMLInputElement>) => {
    if (!session || isUpdatingNumCtx) return;

    const newValue = parseInt(e.currentTarget.value, 10);
    setIsUpdatingNumCtx(true);
    try {
      const updatedSession = await chatSessionApi.update(session._id, {
        numCtx: newValue
      });
      setSession(updatedSession);
      showToast(`Context window set to ${newValue.toLocaleString()}`);
    } catch (err) {
      showToast(`Failed to change context window: ${err instanceof Error ? err.message : 'Unknown error'}`);
    } finally {
      setIsUpdatingNumCtx(false);
    }
  };

  const getProviderName = (id: string) => {
    return providers.find(p => p._id === id)?.name ?? 'None';
  };

  const getModelName = (providerId: string, modelId: string) => {
    if (!providerId || !modelId) return 'None';
    const provider = providers.find(p => p._id === providerId);
    return provider?.models.find(m => m.id === modelId)?.name ?? modelId;
  };

  const getSortedModels = (providerId: string) => {
    const provider = providers.find(p => p._id === providerId);
    if (!provider) return [];
    return [...provider.models].sort((a, b) => a.name.localeCompare(b.name));
  };

  const truncateModelName = (name: string, maxLength = 25) => {
    if (name.length <= maxLength) return name;
    return name.slice(0, maxLength - 3) + '...';
  };

  const handleStartEditProvider = () => {
    setEditProviderId(selectedProviderId);
    setEditModelId(selectedModelId);
    setIsEditingProvider(true);
  };

  const handleSaveProvider = async () => {
    if (!session || !editProviderId || !editModelId) return;
    setIsUpdatingProvider(true);
    try {
      const updatedSession = await chatSessionApi.update(session._id, {
        provider: editProviderId,
        selectedModel: editModelId,
      });
      setSession(updatedSession);
      setSelectedProviderId(editProviderId);
      setSelectedModelId(editModelId);

      // Reset numCtx to new model's max context window
      const newProvider = providers.find(p => p._id === editProviderId);
      const newModel = newProvider?.models.find(m => m.id === editModelId);
      const newMaxCtx = newModel?.contextWindow || 131072;
      setSessionNumCtx(newMaxCtx);

      setIsEditingProvider(false);
      showToast('Provider and model updated');
    } catch (err) {
      showToast(`Failed to update: ${err instanceof Error ? err.message : 'Unknown error'}`);
    } finally {
      setIsUpdatingProvider(false);
    }
  };

  const handleCancelEditProvider = () => {
    setEditProviderId(selectedProviderId);
    setEditModelId(selectedModelId);
    setIsEditingProvider(false);
  };

  const handleStartEditName = () => {
    setEditName(session?.name || '');
    setIsEditingName(true);
  };

  const handleSaveName = async () => {
    if (!session || !editName.trim()) return;
    setIsUpdatingName(true);
    try {
      const updatedSession = await chatSessionApi.update(session._id, { name: editName.trim() });
      setSession(updatedSession);
      setIsEditingName(false);
      showToast('Session name updated');
    } catch (err) {
      showToast(`Failed to update name: ${err instanceof Error ? err.message : 'Unknown error'}`);
    } finally {
      setIsUpdatingName(false);
    }
  };

  const handleCancelEditName = () => {
    setIsEditingName(false);
  };

  const handleSubmitPrompt = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!promptInput.trim() || isProcessing || !socketClient.connected || !sessionLocked) return;
    if (workspaceMode !== WorkspaceMode.Agent) return;
    if (!session?.selectedModel) {
      showToast('Please select a model before submitting a prompt');
      return;
    }

    const tempTurnId = `temp-${Date.now()}`;
    const userTurn: ChatTurn = {
      _id: tempTurnId,
      createdAt: new Date().toISOString(),
      user: session?.user?._id || '',
      project: session?.project?._id || projectId || '',
      session: sessionId || '',
      provider: typeof session?.provider === 'object' ? session?.provider?._id || '' : session?.provider || '',
      llm: session?.selectedModel || '',
      mode: session?.mode || 'develop',
      status: 'processing',
      prompts: { user: promptInput.trim() },
      blocks: [],
      toolCalls: [],
      subagents: [],
      stats: {
        toolCallCount: 0,
        inputTokens: 0,
        thinkingTokenCount: 0,
        responseTokens: 0,
        durationMs: 0,
        durationLabel: '0ms'
      }
    };

    setTurns(prev => [...prev, userTurn]);
    currentTurnIdRef.current = tempTurnId;
    setPromptInput('');
    setIsProcessing(true);
    setError('');

    socketClient.submitPrompt(promptInput.trim(), (success, data) => {
      if (success && data.turnId) {
        setTurns(prevTurns =>
          prevTurns.map(turn =>
            turn._id === tempTurnId
              ? { ...turn, _id: data.turnId! }
              : turn
          )
        );
        currentTurnIdRef.current = data.turnId;
      } else if (!success) {
        setTurns(prevTurns =>
          prevTurns.map(turn =>
            turn._id === tempTurnId
              ? { ...turn, status: 'error', errorMessage: data.message || 'Failed to submit prompt' }
              : turn
          )
        );
        setIsProcessing(false);
        currentTurnIdRef.current = null;
      }
    });
  };

  const handleCancel = useCallback(() => {
    if (isAborting) return;
    setIsAborting(true);
    socketClient.abortWorkOrder((success, message) => {
      if (success) {
        showToast('Aborting Agentic Workflow Loop...');
      } else {
        showToast(message || 'Failed to abort');
        setIsAborting(false);
      }
    });
  }, [isAborting, showToast]);

  // Global Esc key handler for abort: first Esc shows prompt, second Esc within 3s aborts
  useEffect(() => {
    if (!isProcessing) {
      if (escTimerRef.current) {
        clearTimeout(escTimerRef.current);
        escTimerRef.current = null;
      }
      escFlagRef.current = false;
      return;
    }

    const handleKeyDown = (e: KeyboardEvent) => {
      if (e.key !== 'Escape') return;
      e.preventDefault();

      if (!escFlagRef.current) {
        escFlagRef.current = true;
        showToast('Press Esc again to abort');
        escTimerRef.current = setTimeout(() => {
          escFlagRef.current = false;
          setToast(null);
          escTimerRef.current = null;
        }, 3000);
      } else {
        if (escTimerRef.current) {
          clearTimeout(escTimerRef.current);
          escTimerRef.current = null;
        }
        escFlagRef.current = false;
        setToast(null);
        handleCancel();
      }
    };

    window.addEventListener('keydown', handleKeyDown);
    return () => {
      window.removeEventListener('keydown', handleKeyDown);
      if (escTimerRef.current) {
        clearTimeout(escTimerRef.current);
        escTimerRef.current = null;
      }
      escFlagRef.current = false;
    };
  }, [isProcessing, showToast, handleCancel]);

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  };

  const promptDisabled = isProcessing || workspaceMode !== WorkspaceMode.Agent || !sessionLocked || !session?.selectedModel || isEditingProvider;
  const promptPlaceholder = isEditingProvider
    ? 'Save or cancel provider changes to continue.'
    : workspaceMode !== WorkspaceMode.Agent
    ? 'Select Agent mode to enter a prompt.'
    : !session?.selectedModel
    ? 'Please select a model first.'
    : 'Enter your prompt... (Shift+Enter for new line)';

  if (loading) {
    return (
      <div className="flex-1 flex items-center justify-center bg-bg-primary">
        <p className="text-text-muted">Loading chat session...</p>
      </div>
    );
  }

  if (error && !session) {
    return (
      <div className="flex-1 flex items-center justify-center bg-bg-primary">
        <div className="text-center">
          <p className="text-red-500 mb-4">{error}</p>
          <button
            onClick={() => navigate(`/projects/${projectId}`)}
            className="px-4 py-2 bg-brand text-white rounded hover:bg-red-700 transition-colors"
          >
            Back to Project
          </button>
        </div>
      </div>
    );
  }

  // Render "other tab" state
  if (isOtherTab) {
    return (
      <div className="flex-1 flex items-center justify-center bg-bg-primary">
        <div className="text-center">
          <p className="text-text-secondary mb-4">
            This chat session is open in another browser tab.
          </p>
          <button
            onClick={() => navigate(`/projects/${projectId}`)}
            className="px-4 py-2 bg-brand text-white rounded hover:bg-red-700 transition-colors"
          >
            Back to Project
          </button>
        </div>
      </div>
    );
  }

  return (
    <div className="flex-1 flex bg-bg-primary overflow-hidden relative">
      {/* Toast notification */}
      {toast && (
        <div className="absolute top-4 left-1/2 -translate-x-1/2 z-50 px-4 py-2 bg-bg-elevated border border-border-default rounded shadow-lg text-sm text-text-primary">
          {toast}
        </div>
      )}

      {/* Content Area */}
      <div className="flex-1 flex flex-col relative min-w-0">
        {isEditingProvider && (
          <div className="absolute inset-0 bg-bg-primary/80 z-30" />
        )}
        
        {/* File Editor View (replaces Chat View when file is open) */}
        {editorFilePath ? (
          <ErrorBoundary>
            <EditorPanel
              workspaceMode={workspaceMode}
              filePath={editorFilePath}
              onCloseFile={() => setEditorFilePath(undefined)}
            />
          </ErrorBoundary>
        ) : (
          /* Chat View (75%) */
          <div className="flex-1 flex flex-col overflow-hidden">
            {/* Messages */}
            <div className="flex-1 overflow-y-auto p-3 space-y-2">
              {turns.map((turn) => (
                <ChatTurnComponent
                  key={turn._id}
                  turn={turn}
                />
              ))}
              <div ref={messagesEndRef} />
            </div>

          {/* Prompt Input */}
          <div className="border-t border-border-subtle p-4 bg-bg-secondary shrink-0">
            <form onSubmit={handleSubmitPrompt} className="flex gap-2">
              <textarea
                ref={inputRef}
                value={promptInput}
                onChange={(e) => setPromptInput(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === 'Enter' && !e.shiftKey) {
                    e.preventDefault();
                    handleSubmitPrompt(e);
                  }
                }}
                placeholder={promptPlaceholder}
                className="flex-1 px-3 py-2 bg-bg-tertiary border border-border-default rounded text-text-primary focus:border-brand focus:outline-none resize-none disabled:opacity-50"
                rows={3}
                disabled={promptDisabled}
              />
              {isProcessing ? (
                <button
                  type="button"
                  onClick={handleCancel}
                  disabled={isAborting}
                  className="px-6 py-2 bg-red-600 text-white rounded hover:bg-red-700 transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
                >
                  {isAborting ? 'Aborting...' : 'Cancel'}
                </button>
              ) : (
                <button
                  type="submit"
                  disabled={promptDisabled || !promptInput.trim()}
                  className="px-6 py-2 bg-brand text-white rounded hover:bg-brand/80 transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
                >
                  Send
                </button>
              )}
            </form>
          </div>
        </div>
        )}

        {/* Log Panel (25%) */}
        <LogPanel
          logs={logs}
          expanded={logExpanded}
          onToggleExpand={() => setLogExpanded(!logExpanded)}
        />
      </div>

      {/* Sidebar */}
      <div className="w-80 border-l border-border-subtle bg-bg-secondary flex flex-col">
        <SessionPanel
          session={session}
          workspaceMode={workspaceMode}
          sessionLocked={sessionLocked}
          providers={providers}
          selectedProviderId={selectedProviderId}
          selectedModelId={selectedModelId}
          sessionReasoningEffort={sessionReasoningEffort}
          sessionNumCtx={sessionNumCtx}
          isEditingProvider={isEditingProvider}
          editProviderId={editProviderId}
          editModelId={editModelId}
          isUpdatingProvider={isUpdatingProvider}
          isEditingName={isEditingName}
          editName={editName}
          isUpdatingName={isUpdatingName}
          isUpdatingMode={isUpdatingMode}
          isUpdatingReasoning={isUpdatingReasoning}
          isUpdatingNumCtx={isUpdatingNumCtx}
          onWorkspaceModeChange={handleWorkspaceModeChange}
          onModeChange={handleModeChange}
          onReasoningChange={handleReasoningChange}
          onNumCtxChange={handleNumCtxChange}
          onNumCtxCommit={handleNumCtxCommit}
          onStartEditProvider={handleStartEditProvider}
          onSaveProvider={handleSaveProvider}
          onCancelEditProvider={handleCancelEditProvider}
          onEditProviderIdChange={setEditProviderId}
          onEditModelIdChange={setEditModelId}
          onStartEditName={handleStartEditName}
          onSaveName={handleSaveName}
          onCancelEditName={handleCancelEditName}
          onEditNameChange={setEditName}
          getProviderName={getProviderName}
          getModelName={getModelName}
          getSortedModels={getSortedModels}
          truncateModelName={truncateModelName}
          getSelectedModelCapabilities={getSelectedModelCapabilities}
          getSelectedModel={getSelectedModel}
        />

        <ProjectPanel
          project={project}
          onNavigateToProject={() => navigate(`/projects/${project?.slug}`)}
        />

        <FilesPanel 
          workspaceMode={workspaceMode} 
          onFileSelect={(path) => {
            setEditorFilePath(path);
          }}
        />
      </div>
    </div>
  );
}
