/**
 * Claude SDK Integration
 *
 * This module provides SDK-based integration with Claude using the @anthropic-ai/claude-agent-sdk.
 * It mirrors the interface of claude-cli.js but uses the SDK internally for better performance
 * and maintainability.
 *
 * Key features:
 * - Direct SDK integration without child processes
 * - Session management with abort capability
 * - Options mapping between CLI and SDK formats
 * - WebSocket message streaming
 */

import { query, renameSession, forkSession } from '@anthropic-ai/claude-agent-sdk';
import { promises as fs } from 'fs';
import path from 'path';
import { CLAUDE_MODELS } from '../shared/modelConstants.js';
import { PRODUCT_SYSTEM_DESC } from '../shared/brand.js';
import { getUserPaths } from './services/user-directories.js';
import { usageDb } from './database/db.js';
import { calculateCost, normalizeModelName } from './services/pricing.js';
import { evaluate as evaluateToolGuard } from './services/tool-guard/index.js';
import builtinTools, { createRequestId } from './services/builtin-tools/index.js';

// Session tracking: Map of session IDs to active query instances
const activeSessions = new Map();

/**
 * MutableWriter wraps a WebSocketWriter and buffers messages while the client
 * is disconnected. When the client reconnects, switchTo() replays the buffer
 * to the new writer so no streaming chunks are lost.
 */
class MutableWriter {
  constructor(ws) {
    this.current = ws;
    this.buffer = [];
    this.MAX_BUFFER = 500;
  }

  send(data) {
    if (this.current && (!this.current.ws || this.current.ws.readyState === 1)) {
      this.current.send(data);
    } else if (this.buffer.length < this.MAX_BUFFER) {
      this.buffer.push(data);
    }
  }

  setSessionId(sessionId) {
    this.current?.setSessionId(sessionId);
  }

  getSessionId() {
    return this.current?.getSessionId();
  }

  /** Attach a new WebSocketWriter and replay any buffered messages to it. */
  switchTo(newWs) {
    this.current = newWs;
    const buf = this.buffer;
    this.buffer = [];
    for (const msg of buf) {
      newWs.send(msg);
    }
  }
}

// In-memory registry of pending tool approvals keyed by requestId.
// This does not persist approvals or share across processes; it exists so the
// SDK can pause tool execution while the UI decides what to do.
const pendingToolApprovals = new Map();

// Default approval timeout kept under the SDK's 60s control timeout.
// This does not change SDK limits; it only defines how long we wait for the UI,
// introduced to avoid hanging the run when no decision arrives.
const TOOL_APPROVAL_TIMEOUT_MS = parseInt(process.env.CLAUDE_TOOL_APPROVAL_TIMEOUT_MS, 10) || 55000;

// Wait for a UI approval decision, honoring SDK cancellation.
// This does not auto-approve or auto-deny; it only resolves with UI input,
// and it cleans up the pending map to avoid leaks, introduced to prevent
// replying after the SDK cancels the control request.
function waitForToolApproval(requestId, options = {}) {
  const { timeoutMs = TOOL_APPROVAL_TIMEOUT_MS, signal, onCancel } = options;

  return new Promise(resolve => {
    let settled = false;

    const finalize = (decision) => {
      if (settled) return;
      settled = true;
      cleanup();
      resolve(decision);
    };

    const cleanup = () => {
      pendingToolApprovals.delete(requestId);
      clearTimeout(timeout);
      if (signal && abortHandler) {
        signal.removeEventListener('abort', abortHandler);
      }
    };

    // Timeout is local to this process; it does not override SDK timing.
    // It exists to prevent the UI prompt from lingering indefinitely.
    const timeout = setTimeout(() => {
      onCancel?.('timeout');
      finalize(null);
    }, timeoutMs);

    const abortHandler = () => {
      // If the SDK cancels the control request, stop waiting to avoid
      // replying after the process is no longer ready for writes.
      onCancel?.('cancelled');
      finalize({ cancelled: true });
    };

    if (signal) {
      if (signal.aborted) {
        onCancel?.('cancelled');
        finalize({ cancelled: true });
        return;
      }
      signal.addEventListener('abort', abortHandler, { once: true });
    }

    pendingToolApprovals.set(requestId, (decision) => {
      finalize(decision);
    });
  });
}

// Resolve a pending approval. This does not validate the decision payload;
// validation and tool matching remain in canUseTool, which keeps this as a
// lightweight WebSocket -> SDK relay.
function resolveToolApproval(requestId, decision) {
  const resolver = pendingToolApprovals.get(requestId);
  if (resolver) {
    resolver(decision);
  }
}

// Match stored permission entries against a tool + input combo.
// This only supports exact tool names and the Bash(command:*) shorthand
// used by the UI; it intentionally does not implement full glob semantics,
// introduced to stay consistent with the UI's "Allow rule" format.
function matchesToolPermission(entry, toolName, input) {
  if (!entry || !toolName) {
    return false;
  }

  if (entry === toolName) {
    return true;
  }

  const bashMatch = entry.match(/^Bash\((.+):\*\)$/);
  if (toolName === 'Bash' && bashMatch) {
    const allowedPrefix = bashMatch[1];
    let command = '';

    if (typeof input === 'string') {
      command = input.trim();
    } else if (input && typeof input === 'object' && typeof input.command === 'string') {
      command = input.command.trim();
    }

    if (!command) {
      return false;
    }

    return command.startsWith(allowedPrefix);
  }

  return false;
}

/**
 * Maps CLI options to SDK-compatible options format
 * @param {Object} options - CLI options
 * @returns {Object} SDK-compatible options
 */
function mapCliOptionsToSDK(options = {}) {
  const { sessionId, cwd, toolsSettings, permissionMode, images } = options;

  const sdkOptions = {};

  // Map working directory
  if (cwd) {
    sdkOptions.cwd = cwd;
  }

  // Map permission mode
  if (permissionMode && permissionMode !== 'default') {
    sdkOptions.permissionMode = permissionMode;
  }

  // Map tool settings
  const settings = toolsSettings || {
    allowedTools: [],
    disallowedTools: [],
    skipPermissions: false
  };

  // Handle tool permissions
  if (settings.skipPermissions && permissionMode !== 'plan') {
    // When skipping permissions, use bypassPermissions mode
    sdkOptions.permissionMode = 'bypassPermissions';
  }

  // Map allowed tools (always set to avoid implicit "allow all" defaults).
  // This does not grant permissions by itself; it just configures the SDK,
  // introduced because leaving it undefined made the SDK treat it as "all tools allowed."
  let allowedTools = [...(settings.allowedTools || [])];

  // Add plan mode default tools
  if (permissionMode === 'plan') {
    const planModeTools = ['Read', 'Task', 'exit_plan_mode', 'TodoRead', 'TodoWrite', 'WebFetch', 'WebSearch'];
    for (const tool of planModeTools) {
      if (!allowedTools.includes(tool)) {
        allowedTools.push(tool);
      }
    }
  }

  sdkOptions.allowedTools = allowedTools;

  // Map disallowed tools (always set so the SDK doesn't treat "undefined" as permissive).
  // This does not override allowlists; it only feeds the canUseTool gate.
  sdkOptions.disallowedTools = settings.disallowedTools || [];

  // Map model (default to sonnet)
  // Valid models: sonnet, opus, haiku
  sdkOptions.model = options.model || CLAUDE_MODELS.DEFAULT;
  console.log(`Using model: ${sdkOptions.model}`);

  // Map system prompt configuration
  const baseAppend = PRODUCT_SYSTEM_DESC + builtinTools.getSystemPromptAppend();
  const extraAppend = options.appendSystemPrompt ? `\n\n${options.appendSystemPrompt}` : '';
  sdkOptions.systemPrompt = {
    type: 'preset',
    preset: 'claude_code',  // Required to use CLAUDE.md
    append: baseAppend + extraAppend,
  };

  // Map setting sources for CLAUDE.md loading
  // This loads CLAUDE.md from project, user (~/.config/claude/CLAUDE.md), and local directories
  sdkOptions.settingSources = ['project', 'user', 'local'];

  // Map resume session
  if (sessionId) {
    sdkOptions.resume = sessionId;
  }

  return sdkOptions;
}

/**
 * Adds a session to the active sessions map
 * @param {string} sessionId - Session identifier
 * @param {Object} queryInstance - SDK query instance
 * @param {Array<string>} tempImagePaths - Temp image file paths for cleanup
 * @param {string} tempDir - Temp directory for cleanup
 */
function addSession(sessionId, queryInstance, tempImagePaths = [], tempDir = null, abortController = null, mutableWriter = null) {
  activeSessions.set(sessionId, {
    instance: queryInstance,
    startTime: Date.now(),
    status: 'active',
    tempImagePaths,
    tempDir,
    abortController,
    mutableWriter
  });
}

/**
 * Removes a session from the active sessions map
 * @param {string} sessionId - Session identifier
 */
function removeSession(sessionId) {
  activeSessions.delete(sessionId);
}

/**
 * Gets a session from the active sessions map
 * @param {string} sessionId - Session identifier
 * @returns {Object|undefined} Session data or undefined
 */
function getSession(sessionId) {
  return activeSessions.get(sessionId);
}

/**
 * Attaches a new WebSocketWriter to an active session, replaying any messages
 * that were buffered while the previous connection was closed.
 * @param {string} sessionId - Session identifier
 * @param {Object} newWs - New WebSocketWriter instance
 * @returns {boolean} True if the session was found and updated
 */
function updateSessionWriter(sessionId, newWs) {
  const session = activeSessions.get(sessionId);
  if (session?.mutableWriter) {
    session.mutableWriter.switchTo(newWs);
    return true;
  }
  return false;
}

/**
 * Gets all active session IDs
 * @returns {Array<string>} Array of active session IDs
 */
function getAllSessions() {
  return Array.from(activeSessions.keys());
}

// Periodic cleanup of stale sessions (every 5 minutes, remove sessions older than 2 hours)
const SESSION_MAX_AGE_MS = 2 * 60 * 60 * 1000;
setInterval(() => {
  const now = Date.now();
  for (const [sessionId, session] of activeSessions) {
    if (now - session.startTime > SESSION_MAX_AGE_MS) {
      console.log(`[WARN] Cleaning up stale active session: ${sessionId} (age: ${Math.round((now - session.startTime) / 60000)}min)`);
      activeSessions.delete(sessionId);
    }
  }
}, 5 * 60 * 1000);

/**
 * Transforms SDK messages to WebSocket format expected by frontend
 * @param {Object} sdkMessage - SDK message object
 * @returns {Object} Transformed message ready for WebSocket
 */
function transformMessage(sdkMessage) {
  // SDK messages are already in a format compatible with the frontend
  // The CLI sends them wrapped in {type: 'claude-response', data: message}
  // We'll do the same here to maintain compatibility
  return sdkMessage;
}

/**
 * Extracts token usage from SDK result messages
 * @param {Object} resultMessage - SDK result message
 * @returns {Object|null} Token budget object or null
 */
function extractTokenBudget(resultMessage) {
  if (resultMessage.type !== 'result' || !resultMessage.modelUsage) {
    return null;
  }

  // Get the first model's usage data
  const modelKey = Object.keys(resultMessage.modelUsage)[0];
  const modelData = resultMessage.modelUsage[modelKey];

  if (!modelData) {
    return null;
  }

  // Use cumulative tokens if available (tracks total for the session)
  // Otherwise fall back to per-request tokens
  const inputTokens = modelData.cumulativeInputTokens || modelData.inputTokens || 0;
  const outputTokens = modelData.cumulativeOutputTokens || modelData.outputTokens || 0;
  const cacheReadTokens = modelData.cumulativeCacheReadInputTokens || modelData.cacheReadInputTokens || 0;
  const cacheCreationTokens = modelData.cumulativeCacheCreationInputTokens || modelData.cacheCreationInputTokens || 0;

  // Total used = input + output + cache tokens
  const totalUsed = inputTokens + outputTokens + cacheReadTokens + cacheCreationTokens;

  // Use configured context window budget from environment (default 160000)
  // This is the user's budget limit, not the model's context window
  const contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 160000;

  console.log(`Token calculation: input=${inputTokens}, output=${outputTokens}, cache=${cacheReadTokens + cacheCreationTokens}, total=${totalUsed}/${contextWindow}`);

  return {
    used: totalUsed,
    total: contextWindow
  };
}

/**
 * Handles image processing for SDK queries
 * Saves base64 images to temporary files and returns modified prompt with file paths
 * @param {string} command - Original user prompt
 * @param {Array} images - Array of image objects with base64 data
 * @param {string} cwd - Working directory for temp file creation
 * @returns {Promise<Object>} {modifiedCommand, tempImagePaths, tempDir}
 */
async function handleImages(command, images, cwd) {
  const tempImagePaths = [];
  let tempDir = null;

  if (!images || images.length === 0) {
    return { modifiedCommand: command, tempImagePaths, tempDir };
  }

  try {
    // Create temp directory in the project directory
    const workingDir = cwd || process.cwd();
    tempDir = path.join(workingDir, '.tmp', 'images', Date.now().toString());
    await fs.mkdir(tempDir, { recursive: true });

    // Save each image to a temp file
    for (const [index, image] of images.entries()) {
      // Extract base64 data and mime type
      const matches = image.data.match(/^data:([^;]+);base64,(.+)$/);
      if (!matches) {
        console.error(`[claude-sdk] handleImages: image[${index}] invalid data format, prefix="${image.data?.slice(0, 30)}"`);
        continue;
      }

      const [, mimeType, base64Data] = matches;
      const extension = mimeType.split('/')[1] || 'png';
      const filename = `image_${index}.${extension}`;
      const filepath = path.join(tempDir, filename);

      // Write base64 data to file
      await fs.writeFile(filepath, Buffer.from(base64Data, 'base64'));
      tempImagePaths.push(filepath);
    }

    // Include the full image paths in the prompt
    let modifiedCommand = command;
    if (tempImagePaths.length > 0) {
      const imageNote = `[Images provided at the following paths:]\n${tempImagePaths.map((p, i) => `${i + 1}. ${p}`).join('\n')}`;
      modifiedCommand = (command && command.trim())
        ? command + '\n\n' + imageNote
        : imageNote;
    }

    console.log(`Processed ${tempImagePaths.length} images to temp directory: ${tempDir}`);
    return { modifiedCommand, tempImagePaths, tempDir };
  } catch (error) {
    console.error('Error processing images for SDK:', error);
    return { modifiedCommand: command, tempImagePaths, tempDir };
  }
}

/**
 * Cleans up temporary image files
 * @param {Array<string>} tempImagePaths - Array of temp file paths to delete
 * @param {string} tempDir - Temp directory to remove
 */
async function cleanupTempFiles(tempImagePaths, tempDir) {
  if (!tempImagePaths || tempImagePaths.length === 0) {
    return;
  }

  try {
    // Delete individual temp files
    for (const imagePath of tempImagePaths) {
      await fs.unlink(imagePath).catch(err =>
        console.error(`Failed to delete temp image ${imagePath}:`, err)
      );
    }

    // Delete temp directory
    if (tempDir) {
      await fs.rm(tempDir, { recursive: true, force: true }).catch(err =>
        console.error(`Failed to delete temp directory ${tempDir}:`, err)
      );
    }

    console.log(`Cleaned up ${tempImagePaths.length} temp image files`);
  } catch (error) {
    console.error('Error during temp file cleanup:', error);
  }
}

/**
 * Loads MCP server configurations from user-specific .claude.json
 * @param {string} cwd - Current working directory for project-specific configs
 * @param {string} userUuid - User UUID for user-specific config (required)
 * @returns {Object|null} MCP servers object or null if none found
 */
async function loadMcpConfig(cwd, userUuid) {
  if (!userUuid) {
    console.log('No userUuid provided for loadMcpConfig, proceeding without MCP servers');
    return null;
  }
  try {
    const configDir = getUserPaths(userUuid).configDir;
    const claudeConfigPath = path.join(configDir, '.claude.json');

    // Check if config file exists
    try {
      await fs.access(claudeConfigPath);
    } catch (error) {
      // File doesn't exist, return null
      console.log('No ~/.claude.json found, proceeding without MCP servers');
      return null;
    }

    // Read and parse config file
    let claudeConfig;
    try {
      const configContent = await fs.readFile(claudeConfigPath, 'utf8');
      claudeConfig = JSON.parse(configContent);
    } catch (error) {
      console.error('Failed to parse ~/.claude.json:', error.message);
      return null;
    }

    // Extract MCP servers (merge global and project-specific)
    let mcpServers = {};

    // Add global MCP servers
    if (claudeConfig.mcpServers && typeof claudeConfig.mcpServers === 'object') {
      mcpServers = { ...claudeConfig.mcpServers };
      console.log(`Loaded ${Object.keys(mcpServers).length} global MCP servers`);
    }

    // Add/override with project-specific MCP servers
    if (claudeConfig.claudeProjects && cwd) {
      const projectConfig = claudeConfig.claudeProjects[cwd];
      if (projectConfig && projectConfig.mcpServers && typeof projectConfig.mcpServers === 'object') {
        mcpServers = { ...mcpServers, ...projectConfig.mcpServers };
        console.log(`Loaded ${Object.keys(projectConfig.mcpServers).length} project-specific MCP servers`);
      }
    }

    // Return null if no servers found
    if (Object.keys(mcpServers).length === 0) {
      console.log('No MCP servers configured');
      return null;
    }

    console.log(`Total MCP servers loaded: ${Object.keys(mcpServers).length}`);
    return mcpServers;
  } catch (error) {
    console.error('Error loading MCP config:', error.message);
    return null;
  }
}

/**
 * Executes a Claude query using the SDK
 * @param {string} command - User prompt/command
 * @param {Object} options - Query options
 * @param {Object} ws - WebSocket connection
 * @returns {Promise<void>}
 */
async function queryClaudeSDK(command, options = {}, ws) {
  const { sessionId, userUuid } = options;

  // Wrap the WebSocketWriter in a MutableWriter so messages are buffered
  // if the client disconnects mid-stream and replayed on reconnect.
  const mutableWriter = new MutableWriter(ws);
  let capturedSessionId = sessionId;
  let sessionCreatedSent = false;
  let tempImagePaths = [];
  let tempDir = null;

  // Set CLAUDE_CONFIG_DIR for user isolation
  if (userUuid) {
    const userPaths = getUserPaths(userUuid);
    process.env.CLAUDE_CONFIG_DIR = userPaths.claudeDir;
    console.log(`Set CLAUDE_CONFIG_DIR to user data directory (contains .claude folder): ${userPaths.claudeDir}`);
  }

  try {
    // Map CLI options to SDK format
    const sdkOptions = mapCliOptionsToSDK(options);

    // Create AbortController for session cancellation
    const abortController = new AbortController();
    sdkOptions.abortController = abortController;

    // Load MCP configuration
    const mcpServers = await loadMcpConfig(options.cwd, userUuid);
    if (mcpServers) {
      sdkOptions.mcpServers = mcpServers;
    }

    // Handle images - save to temp files and modify prompt
    const imageResult = await handleImages(command, options.images, options.cwd);
    const finalCommand = imageResult.modifiedCommand;
    tempImagePaths = imageResult.tempImagePaths;
    tempDir = imageResult.tempDir;

    // ===== 安全守卫（PreToolUse hook，最高优先级）=====
    // 使用 PreToolUse hook 而非 canUseTool 回调，因为 hook 会在每次工具执行前触发，
    // 包括已被 allowedTools 自动放行的工具和 bypassPermissions 模式下的工具。
    // canUseTool 仅在 CLI 发送 can_use_tool 控制消息时触发（即工具未被预批准时）。
    if (userUuid) {
      sdkOptions.hooks = {
        ...(sdkOptions.hooks || {}),
        PreToolUse: [
          ...((sdkOptions.hooks || {}).PreToolUse || []),
          {
            hooks: [async (hookInput) => {
              // ===== 内置工具拦截（在 ToolGuard 之前，匹配后直接 return） =====
              const builtinResult = await builtinTools.handlePreToolUse(hookInput, {
                userUuid, mutableWriter, cwd: sdkOptions.cwd,
              });
              if (builtinResult) return builtinResult;

              // ===== 安全守卫（Tool Guard）=====
              try {
                const guardResult = await evaluateToolGuard(hookInput.tool_name, hookInput.tool_input, {
                  userUuid,
                  cwd: sdkOptions.cwd,
                });
                if (!guardResult.allowed) {
                  console.log(`[ToolGuard] DENIED ${hookInput.tool_name} for user ${userUuid}: ${guardResult.reason}`);
                  return {
                    hookSpecificOutput: {
                      hookEventName: 'PreToolUse',
                      permissionDecision: 'deny',
                      permissionDecisionReason: `[系统安全策略] ${guardResult.reason}`,
                    }
                  };
                }
                // 路径重写：将 .claude 路径替换为用户专属目录后告知 SDK 使用新参数
                if (guardResult.modifiedInput) {
                  return {
                    hookSpecificOutput: {
                      hookEventName: 'PreToolUse',
                      permissionDecision: 'allow',
                      updatedInput: guardResult.modifiedInput,
                    }
                  };
                }
              } catch (err) {
                console.error(`[ToolGuard] Error evaluating ${hookInput.tool_name}:`, err.message);
                // 守卫出错时不阻塞正常流程，降级为原有权限控制
              }
              if (process.env.TOOL_GUARD_VERBOSE === 'true') {
                console.log(`[ToolGuard] HOOK-ALLOW ${hookInput.tool_name}`);
              }
              return {};
            }]
          }
        ]
      };
    }

    // Gate tool usage with explicit UI approval when not auto-approved.
    // This does not render UI or persist permissions; it only bridges to the UI
    // via WebSocket and waits for the response, introduced so tool calls pause
    // instead of auto-running when the allowlist is empty.
    sdkOptions.canUseTool = async (toolName, input, context) => {
      if (sdkOptions.permissionMode === 'bypassPermissions') {
        return { behavior: 'allow', updatedInput: input };
      }

      const isDisallowed = (sdkOptions.disallowedTools || []).some(entry =>
        matchesToolPermission(entry, toolName, input)
      );
      if (isDisallowed) {
        return { behavior: 'deny', message: 'Tool disallowed by settings' };
      }

      const isAllowed = (sdkOptions.allowedTools || []).some(entry =>
        matchesToolPermission(entry, toolName, input)
      );
      if (isAllowed) {
        return { behavior: 'allow', updatedInput: input };
      }

      const requestId = createRequestId();
      mutableWriter.send({
        type: 'claude-permission-request',
        requestId,
        toolName,
        input,
        sessionId: capturedSessionId || sessionId || null
      });

      // Wait for the UI; if the SDK cancels, notify the UI so it can dismiss the banner.
      // This does not retry or resurface the prompt; it just reflects the cancellation.
      // AskUserQuestion needs more time for the user to type a response.
      const approvalTimeoutMs = toolName === 'AskUserQuestion' ? 3 * 60 * 1000 : TOOL_APPROVAL_TIMEOUT_MS;
      const decision = await waitForToolApproval(requestId, {
        timeoutMs: approvalTimeoutMs,
        signal: context?.signal,
        onCancel: (reason) => {
          mutableWriter.send({
            type: 'claude-permission-cancelled',
            requestId,
            reason,
            sessionId: capturedSessionId || sessionId || null
          });
        }
      });
      if (!decision) {
        return { behavior: 'deny', message: 'Permission request timed out' };
      }

      if (decision.cancelled) {
        return { behavior: 'deny', message: 'Permission request cancelled' };
      }

      if (decision.allow) {
        // rememberEntry only updates this run's in-memory allowlist to prevent
        // repeated prompts in the same session; persistence is handled by the UI.
        if (decision.rememberEntry && typeof decision.rememberEntry === 'string') {
          if (!sdkOptions.allowedTools.includes(decision.rememberEntry)) {
            sdkOptions.allowedTools.push(decision.rememberEntry);
          }
          if (Array.isArray(sdkOptions.disallowedTools)) {
            sdkOptions.disallowedTools = sdkOptions.disallowedTools.filter(entry => entry !== decision.rememberEntry);
          }
        }
        return { behavior: 'allow', updatedInput: decision.updatedInput ?? input };
      }

      return { behavior: 'deny', message: decision.message ?? 'User denied tool use' };
    };

    // Create SDK query instance
    const queryInstance = query({
      prompt: finalCommand,
      options: sdkOptions
    });

    // Track the query instance for abort capability
    if (capturedSessionId) {
      addSession(capturedSessionId, queryInstance, tempImagePaths, tempDir, abortController, mutableWriter);
      mutableWriter.setSessionId(capturedSessionId); // resume 会话时立即同步 sessionId，确保 PreToolUse hook 可读到
    }

    // Process streaming messages
    console.log('Starting async generator loop for session:', capturedSessionId || 'NEW');
    for await (const message of queryInstance) {
      // Capture session ID from first message
      if (message.session_id && !capturedSessionId) {

        capturedSessionId = message.session_id;
        addSession(capturedSessionId, queryInstance, tempImagePaths, tempDir, abortController, mutableWriter);

        // Set session ID on writer
        if (ws.setSessionId && typeof ws.setSessionId === 'function') {
          mutableWriter.setSessionId(capturedSessionId);
        }

        // Send session-created event only once for new sessions
        if (!sessionId && !sessionCreatedSent) {
          sessionCreatedSent = true;
          mutableWriter.send({
            type: 'session-created',
            sessionId: capturedSessionId
          });
        } else {
          console.log('Not sending session-created. sessionId:', sessionId, 'sessionCreatedSent:', sessionCreatedSent);
        }
      } else {
        console.log('No session_id in message or already captured. message.session_id:', message.session_id, 'capturedSessionId:', capturedSessionId);
      }

      // Handle api_retry system messages specially
      if (message.type === 'system' && message.subtype === 'api_retry') {
        mutableWriter.send({
          type: 'claude-api-retry',
          attempt: message.attempt,
          maxRetries: message.max_retries,
          retryDelayMs: message.retry_delay_ms,
          errorStatus: message.error_status
        });
        continue;
      }

      // Transform and send message to WebSocket
      const transformedMessage = transformMessage(message);

      // 为内置工具拦截的 tool_result 打 _builtinHandled 标记，前端据此不显示错误样式
      if (transformedMessage.message?.role === 'user' && Array.isArray(transformedMessage.message?.content)) {
        for (const part of transformedMessage.message.content) {
          if (part.type === 'tool_result' && part.is_error && part.tool_use_id) {
            if (builtinTools.consumeBuiltinHandled(part.tool_use_id)) {
              part._builtinHandled = true;
            }
          }
        }
      }

      mutableWriter.send({
        type: 'claude-response',
        data: transformedMessage
      });

      // Extract and send token budget updates from result messages
      if (message.type === 'result') {
        const tokenBudget = extractTokenBudget(message);
        if (tokenBudget) {
          console.log('Token budget from modelUsage:', tokenBudget);
          mutableWriter.send({
            type: 'token-budget',
            data: tokenBudget
          });
        }

        // Record usage to database for admin statistics
        if (userUuid && message.modelUsage) {
          try {
            const modelKey = Object.keys(message.modelUsage)[0];
            const modelData = message.modelUsage[modelKey];

            if (modelData) {
              const inputTokens = modelData.inputTokens || 0;
              const outputTokens = modelData.outputTokens || 0;
              const cacheReadTokens = modelData.cacheReadInputTokens || 0;
              const cacheCreationTokens = modelData.cacheCreationInputTokens || 0;

              // Extract detailed cache creation tokens if available
              const cacheCreation = modelData.cacheCreation || {};
              const cacheCreation5mTokens = cacheCreation.ephemeral5mInputTokens || 0;
              const cacheCreation1hTokens = cacheCreation.ephemeral1hInputTokens || 0;
              const hasPreciseCacheData = cacheCreation5mTokens > 0 || cacheCreation1hTokens > 0;

              const normalizedModel = normalizeModelName(modelKey);
              const cost = calculateCost({
                model: normalizedModel,
                inputTokens,
                outputTokens,
                cacheReadTokens,
                // Use precise data if available, otherwise fallback to legacy field
                cacheCreation5mTokens: hasPreciseCacheData ? cacheCreation5mTokens : undefined,
                cacheCreation1hTokens: hasPreciseCacheData ? cacheCreation1hTokens : undefined,
                cacheCreationTokens: hasPreciseCacheData ? undefined : cacheCreationTokens
              });

              // Atomically insert usage record + update daily summary in a single transaction
              const today = new Date().toISOString().split('T')[0];
              usageDb.recordUsageTransaction(
                {
                  user_uuid: userUuid,
                  session_id: capturedSessionId,
                  model: normalizedModel,
                  raw_model: modelKey,
                  input_tokens: inputTokens,
                  output_tokens: outputTokens,
                  cache_read_tokens: cacheReadTokens,
                  cache_creation_tokens: cacheCreationTokens,
                  cost_usd: cost,
                  source: 'sdk'
                },
                {
                  user_uuid: userUuid,
                  date: today,
                  model: normalizedModel,
                  total_input_tokens: inputTokens,
                  total_output_tokens: outputTokens,
                  total_cost_usd: cost,
                  session_count: 0, // Session count updated separately
                  request_count: 1
                }
              );

              console.log(`Recorded usage for user ${userUuid}: ${normalizedModel}, cost: $${cost.toFixed(6)}`);
            }
          } catch (usageError) {
            console.error('Error recording usage:', usageError);
            // Don't throw - usage recording failure shouldn't break the query
          }
        }
      }
    }

    // Clean up session on completion
    if (capturedSessionId) {
      removeSession(capturedSessionId);
    }

    // Clean up temporary image files
    await cleanupTempFiles(tempImagePaths, tempDir);

    // Only send completion event if not aborted — when the session is aborted
    // via abortClaudeSDKSession(), that function already sends 'session-aborted'
    // to the frontend, so we must not send a duplicate 'claude-complete'.
    if (!abortController.signal.aborted) {
      console.log('Streaming complete, sending claude-complete event');
      mutableWriter.send({
        type: 'claude-complete',
        sessionId: capturedSessionId,
        exitCode: 0,
        isNewSession: !sessionId && !!command
      });
      console.log('claude-complete event sent');
    } else {
      console.log('Session was aborted, skipping claude-complete event');
    }

  } catch (error) {
    // If the session was aborted, this is expected — do not treat as an error.
    if (abortController.signal.aborted) {
      console.log('Session aborted, ignoring post-abort error:', error.message);
      if (capturedSessionId) {
        removeSession(capturedSessionId);
      }
      await cleanupTempFiles(tempImagePaths, tempDir);
      return;
    }

    console.error('SDK query error:', error);

    // Clean up session on error
    if (capturedSessionId) {
      removeSession(capturedSessionId);
    }

    // Clean up temporary image files on error
    await cleanupTempFiles(tempImagePaths, tempDir);

    // Send error to WebSocket
    mutableWriter.send({
      type: 'claude-error',
      error: error.message
    });

    throw error;
  }
}

/**
 * Aborts an active SDK session
 * @param {string} sessionId - Session identifier
 * @returns {boolean} True if session was aborted, false if not found
 */
async function abortClaudeSDKSession(sessionId) {
  const session = getSession(sessionId);

  if (!session) {
    console.log(`Session ${sessionId} not found`);
    return false;
  }

  try {
    console.log(`Aborting SDK session: ${sessionId}`);

    // Signal abort via AbortController (stops SDK internal operations)
    if (session.abortController) {
      session.abortController.abort();
    }

    // Forcefully close the query and terminate the underlying process.
    // close() is the correct method for aborting a running query — it kills the
    // CLI subprocess and cleans up all resources (MCP transports, pending requests).
    // interrupt() only works in streaming-input mode and is not reliable here.
    try {
      session.instance.close();
    } catch (closeError) {
      console.error(`Error closing session ${sessionId}:`, closeError);
    }

    // Update session status
    session.status = 'aborted';

    // Clean up temporary image files
    await cleanupTempFiles(session.tempImagePaths, session.tempDir);

    // Clean up session
    removeSession(sessionId);

    return true;
  } catch (error) {
    console.error(`Error aborting session ${sessionId}:`, error);
    return false;
  }
}

/**
 * Checks if an SDK session is currently active
 * @param {string} sessionId - Session identifier
 * @returns {boolean} True if session is active
 */
function isClaudeSDKSessionActive(sessionId) {
  const session = getSession(sessionId);
  return session && session.status === 'active';
}

/**
 * Gets all active SDK session IDs
 * @returns {Array<string>} Array of active session IDs
 */
function getActiveClaudeSDKSessions() {
  return getAllSessions();
}

/**
 * Rename a session for a specific user
 * @param {string} sessionId - Session UUID
 * @param {string} title - New session title
 * @param {string} userUuid - User UUID for directory isolation
 */
async function renameSessionForUser(sessionId, title, userUuid) {
  const userPaths = getUserPaths(userUuid);
  process.env.CLAUDE_CONFIG_DIR = userPaths.claudeDir;
  await renameSession(sessionId, title);
}

/**
 * Fork a session at a given message, creating a new branch conversation.
 * @param {string} sessionId - Session to fork
 * @param {string} userUuid - User UUID for directory isolation
 * @param {object} opts - Fork options (upToMessageId, title)
 * @returns {Promise<string>} New session UUID
 */
async function forkSessionForUser(sessionId, userUuid, opts = {}) {
  const userPaths = getUserPaths(userUuid);
  process.env.CLAUDE_CONFIG_DIR = userPaths.claudeDir;
  const result = await forkSession(sessionId, opts);
  return result.sessionId;
}

// Export public API
export {
  queryClaudeSDK,
  abortClaudeSDKSession,
  isClaudeSDKSessionActive,
  getActiveClaudeSDKSessions,
  resolveToolApproval,
  renameSessionForUser,
  forkSessionForUser,
  updateSessionWriter
};
