/**
 * Chat Agent - LangGraph ReAct Agent
 *
 * A dedicated chat agent for handling conversational interactions with support for:
 * - Style formatting
 * - Built-in commands (prompts)
 * - Optional tool usage
 * - Conversation state management
 * - Streaming responses
 * - Search results context
 */

import {
  MessagesAnnotation,
  StateGraph,
  START,
  Annotation,
} from "@langchain/langgraph/web";
import {
  SystemMessage,
  HumanMessage,
  ToolMessage,
  AIMessage,
  RemoveMessage,
} from "@langchain/core/messages";
import { ToolNode } from "@langchain/langgraph/prebuilt";
import { concat } from "@langchain/core/utils/stream";
import {
  LlmInfos,
  modelViaLanggraph,
  TokensUsage,
} from "../langraphModelsLoader";
import { updateTokenCounter } from "../../modelsInfo";
import { hasCapability } from "../../modelRegistry";
import { StructuredOutputType } from "@langchain/core/language_models/base";
import {
  buildChatSystemPrompt,
  buildCompleteCommandPrompt,
  buildConversationContext,
  buildResultsContext,
  SUMMARIZATION_PROMPT,
} from "./chat-agent-prompts";
import { getChatTools } from "./chat-tools";
import {
  isImageGenerationRequest,
  handleImageGenerationCommand,
  hasAudioContent,
  handleAudioAnalysisRequest,
  isRequestingFreshAudioTranscription,
  handleFreshAudioTranscriptionRequest,
  hasVideoContent,
  handleVideoAnalysisRequest,
  hasPdfContent,
  handlePdfAnalysisRequest,
  isWebSearchRequest,
  handleWebSearchRequest,
  isFileExportRequest,
  handleFileExportCommand,
} from "./multimodal-commands";
import { getContextFromQueries } from "../../queryContextExtractor";
import { alwaysExtractQuery } from "../../..";

// Chat Agent State
const ChatAgentState = Annotation.Root({
  ...MessagesAnnotation.spec,
  // Model configuration
  model: Annotation<LlmInfos>,
  // Response formatting
  style: Annotation<string | undefined>,
  // Command/prompt context
  commandPrompt: Annotation<string | undefined>,
  // Chat session identifier for multi-turn image editing
  chatSessionId: Annotation<string | undefined>,
  // Tool control
  toolsEnabled: Annotation<boolean>,
  enabledTools: Annotation<Set<string> | undefined>,
  permissions: Annotation<{ contentAccess: boolean }>,
  chatTools: Annotation<any[]>,
  // Conversation state
  conversationHistory: Annotation<string[] | undefined>,
  conversationSummary: Annotation<string | undefined>,
  exchangesSinceLastSummary: Annotation<number>,
  // Results context
  resultsContext: Annotation<any[] | undefined>,
  resultsDescription: Annotation<string | undefined>,
  // Access mode
  accessMode: Annotation<"Balanced" | "Full Access">,
  // Agent mode (enables deep analysis features)
  isAgentMode: Annotation<boolean>,
  // Streaming
  streamingCallback: Annotation<((content: string) => void) | undefined>,
  // Abort signal
  abortSignal: Annotation<AbortSignal | undefined>,
  // Tool usage callback
  toolUsageCallback: Annotation<
    | ((toolInfo: { toolName: string; args?: Record<string, any> }) => void)
    | undefined
  >,
  // Agent callbacks
  addResultsCallback: Annotation<((results: any[]) => void) | undefined>,
  selectResultsCallback: Annotation<((uids: string[]) => void) | undefined>,
  expandedResultsCallback: Annotation<
    ((currentResults: any[]) => Promise<any[]>) | undefined
  >,
  toolResponseCallback: Annotation<
    ((toolInfo: { toolName: string; response: string }) => void) | undefined
  >,
  // Tool confirmation callback for sensitive operations
  toolConfirmationCallback: Annotation<
    | ((confirmationRequest: {
        toolName: string;
        toolCallId: string;
        args: Record<string, any>;
      }) => Promise<{
        approved: boolean;
        alwaysApprove?: boolean;
        declineReason?: string;
      }>)
    | undefined
  >,
  // User choice callback for inline choice forms
  userChoiceCallback: Annotation<
    | ((choiceRequest: {
        commandId: string;
        title: string;
        hintsEnabled?: boolean;
        options: Array<{
          id: string;
          label: string;
          type: "radio" | "checkbox" | "text" | "slider";
          choices?: Array<{ value: string; label: string; hint?: string }>;
          defaultValue?: string;
          placeholder?: string;
          min?: number;
          max?: number;
          step?: number;
        }>;
      }) => Promise<{
        selectedOptions: Record<string, string>;
        cancelled: boolean;
      }>)
    | undefined
  >,
  // Set of tools that have been "always approved" for this session
  alwaysApprovedTools: Annotation<Set<string> | undefined>,
  needsExpansion: Annotation<boolean>,
  // Timing
  startTime: Annotation<number>,
  // Tool results cache
  toolResultsCache: Annotation<Record<string, any>>,
  // Active skill instructions (replaces previous when new skill is loaded)
  activeSkillInstructions: Annotation<string | undefined>,
  // Response tracking
  finalAnswer: Annotation<string | undefined>,
  // Token usage tracking (cumulative across the conversation)
  tokensUsage: Annotation<TokensUsage>,
  // This turn's token usage only (for per-message display)
  turnTokensDelta: Annotation<TokensUsage | undefined>,
  // Invalid tool call retry counter
  invalidToolCallRetries: Annotation<number>,
  // Audio transcription cache (URL -> transcription)
  audioTranscriptionCache: Annotation<Map<string, string>>,
  // Image edition mode
  imageEditionMode: Annotation<boolean>,
  // Model used for the last image generation (to reuse for editing)
  imageGenerationModelId: Annotation<string | undefined>,
  // URL of the last generated image (for non-Gemini models that need it passed explicitly)
  lastGeneratedImageUrl: Annotation<string | undefined>,
  // Per-session PDF extraction override
  includePdf: Annotation<boolean | undefined>,
});

// Module-level variables
let llm: StructuredOutputType;
let turnTokensUsage: TokensUsage;
let initialTokensUsage: TokensUsage;
let sys_msg: SystemMessage;
let originalUserMessageForHistory: string;

// Conversation summarization threshold
const SUMMARIZATION_THRESHOLD = 25;

// Number of most recent turns ALWAYS kept verbatim (never summarized away), so
// the model retains exact recent context even after summarization.
const RECENT_TURNS_TO_KEEP = 6;

// Nodes

/**
 * Load model and prepare tools based on configuration
 */
const loadModel = async (state: typeof ChatAgentState.State) => {
  const startTime = state.startTime || Date.now();

  // Initialize token usage tracking
  turnTokensUsage = state.tokensUsage || { input_tokens: 0, output_tokens: 0 };
  initialTokensUsage = { ...turnTokensUsage };

  llm = modelViaLanggraph(state.model, turnTokensUsage);

  // Get tools based on permissions and tool enablement
  const chatTools = getChatTools(
    state.toolsEnabled,
    state.permissions,
    state.enabledTools,
  );

  // Build conversation context
  const conversationContext = buildConversationContext(
    state.conversationHistory,
    state.conversationSummary,
  );

  // Build results context
  const resultsContext = state.resultsContext
    ? buildResultsContext(state.resultsContext, state.resultsDescription)
    : undefined;

  // Extract text content from the last message
  // LangChain messages can have content as string or array of content parts (for multimodal)
  const lastMessageContent = state.messages?.at(-1)?.content;
  let lastMessage = "";
  if (typeof lastMessageContent === "string") {
    lastMessage = lastMessageContent;
  } else if (Array.isArray(lastMessageContent)) {
    // For multimodal messages, extract text parts and preserve image markdown
    lastMessage = lastMessageContent
      .map((part: any) => {
        if (typeof part === "string") return part;
        if (part.type === "text") return part.text;
        // For image parts, try to reconstruct markdown if URL is available
        if (part.type === "image_url" && part.image_url?.url) {
          return `![](${part.image_url.url})`;
        }
        return "";
      })
      .filter(Boolean)
      .join("\n");
  }

  // Store original message for history tracking before modification
  originalUserMessageForHistory = lastMessage;

  let isConversationContextToInclude = lastMessage || resultsContext;
  if (state.commandPrompt) {
    let completedLastMessage = buildCompleteCommandPrompt(
      state.commandPrompt,
      lastMessage,
      !!resultsContext,
    );
    state.messages.pop();
    state.messages.push(new HumanMessage(completedLastMessage));
  }

  // Extract query results if alwaysExtractQuery is enabled
  let queryContext: string | null = null;
  if (alwaysExtractQuery) {
    queryContext = await getContextFromQueries({
      prompt: lastMessage,
      context: resultsContext || "",
      model: state.model.id,
      rootUid: "",
    });
  }

  // Build system prompt - NOTE: Don't include tool descriptions
  // LangChain's bindTools() handles tool schemas automatically via the API
  const systemPrompt = await buildChatSystemPrompt({
    lastMessage: lastMessage,
    style: state.style,
    // Command prompt is included in system prompt and also added to conversationHistory
    // in finalize() so it persists across turns until summarization
    commandPrompt: state.commandPrompt,
    toolsEnabled: state.toolsEnabled,
    conversationContext: isConversationContextToInclude && conversationContext,
    resultsContext,
    accessMode: state.accessMode,
    isAgentMode: state.isAgentMode,
    activeSkillInstructions: state.activeSkillInstructions,
    enabledTools: state.enabledTools,
    hasAudioContent: hasAudioContent(lastMessage, state.resultsContext),
    hasVideoContent: hasVideoContent(lastMessage, state.resultsContext),
    hasPdfContent: hasPdfContent(lastMessage, state.resultsContext, state.includePdf),
    maxTokensLimit: state.model.advancedParams?.maxTokens,
  });

  const finalSystemPrompt = queryContext
    ? systemPrompt + "\n\n" + queryContext
    : systemPrompt;
  sys_msg = new SystemMessage({ content: finalSystemPrompt });

  return {
    messages: state.messages,
    chatTools,
    startTime,
  };
};

/**
 * Check if conversation needs summarization
 */
const shouldSummarize = (state: typeof ChatAgentState.State): boolean => {
  return (
    state.exchangesSinceLastSummary >= SUMMARIZATION_THRESHOLD &&
    state.conversationHistory &&
    state.conversationHistory.length > 0
  );
};

/**
 * Summarize conversation to reduce token usage
 */
const summarizeConversation = async (state: typeof ChatAgentState.State) => {
  if (!shouldSummarize(state)) {
    return {};
  }

  const history = state.conversationHistory || [];
  // History holds ~2 entries per turn (<User>: …, <Assistant>: …). Keep the most
  // recent turns verbatim and only summarize the OLDER ones — never discard the
  // whole history (that was the cause of context being lost mid-chat).
  const keepEntries = Math.min(RECENT_TURNS_TO_KEEP * 2, history.length);
  const olderEntries = history.slice(0, history.length - keepEntries);
  const recentEntries = history.slice(history.length - keepEntries);

  // Not enough old material to summarize yet: keep history, just reset counter.
  if (olderEntries.length === 0) {
    return {
      conversationHistory: history,
      exchangesSinceLastSummary: Math.ceil(recentEntries.length / 2),
    };
  }

  // Fold the previous summary into the new one so context ACCUMULATES across
  // successive summarizations instead of being repeatedly thrown away.
  const priorSummary = state.conversationSummary
    ? `Summary of earlier conversation (already condensed):\n${state.conversationSummary}\n\n`
    : "";
  const conversationText = priorSummary + olderEntries.join("\n\n");
  const prompt = SUMMARIZATION_PROMPT.replace("{conversation}", conversationText);

  // Fail safe: if summarization errors, keep the recent history and the prior
  // summary rather than losing everything.
  let summary = state.conversationSummary || "";
  try {
    const response = await llm.invoke([new HumanMessage({ content: prompt })]);
    summary = response.content.toString();
    // Account for the summarization call's cost. The token callback in
    // modelViaLanggraph only fires for non-thinking, non-Google models; for
    // thinking/Google models we must read usage_metadata here (as the assistant
    // node does) so this extra LLM call isn't a hidden, uncounted cost.
    if (
      (response as any).usage_metadata &&
      (state.model.thinking || state.model.provider === "Google")
    ) {
      turnTokensUsage.input_tokens +=
        (response as any).usage_metadata.input_tokens || 0;
      turnTokensUsage.output_tokens +=
        (response as any).usage_metadata.output_tokens || 0;
    }
  } catch (error) {
    console.warn(
      "[ChatAgent] Conversation summarization failed; keeping recent history verbatim:",
      error,
    );
  }

  return {
    conversationSummary: summary,
    // Recent turns remain un-summarized; count them so we only summarize again
    // once enough NEW turns accumulate.
    exchangesSinceLastSummary: Math.ceil(recentEntries.length / 2),
    // Keep recent exchanges verbatim instead of clearing all history.
    conversationHistory: recentEntries,
  };
};

/**
 * Assistant node - generates response with optional tool calls
 */
const assistant = async (state: typeof ChatAgentState.State) => {
  const msgTypes = state.messages.map(
    (m: any) => m.getType?.() || m._getType?.() || "unknown",
  );

  let messages = [sys_msg, ...state.messages];
  let gathered: any = undefined;

  // Handle multimodal commands (image generation, audio, video, web search)

  // Handle web search command
  if (isWebSearchRequest(state.commandPrompt)) {
    const result = await handleWebSearchRequest(
      originalUserMessageForHistory,
      state.model.id, // Pass model ID for web search
      state.resultsContext,
      state.messages,
    );

    if (result.tokensUsage) {
      turnTokensUsage = { ...result.tokensUsage };
    }

    return { messages: result.messages };
  }

  // Handle file export commands (PDF, DOCX, PPTX)
  if (isFileExportRequest(state.commandPrompt)) {
    const result = await handleFileExportCommand(
      state.commandPrompt!,
      originalUserMessageForHistory,
      state.resultsContext,
      state.conversationHistory,
      state.conversationSummary,
      state.userChoiceCallback,
      state.messages,
    );
    return { messages: result.messages };
  }

  // Handle image edition mode - route all prompts to image editing
  if (state.imageEditionMode) {
    // Check for exit commands (safety net - normally handled at UI level)
    const userInput = originalUserMessageForHistory.trim().toLowerCase();
    if (
      userInput === "/exit-edit" ||
      userInput === "/conversation"
    ) {
      return {
        messages: [
          ...state.messages,
          new AIMessage(
            "*Exited image edition mode. Your messages will now be processed as regular conversation.*",
          ),
        ],
        imageEditionMode: false,
      };
    }

    // Determine which model to use for editing
    let editModel = state.imageGenerationModelId;

    // Check if the stored model supports editing
    if (editModel && !hasCapability(editModel, "editImage")) {
      editModel = "gemini-3-pro-image-preview"; // Fallback to Nano Banana Pro
    }

    // If no model stored, use Nano Banana Pro
    if (!editModel) {
      editModel = "gemini-3-pro-image-preview";
    }

    // For non-Gemini models, prepend last image as markdown so imageGeneration() picks it up
    const isGeminiModel = editModel.includes("gemini");
    let editPrompt = originalUserMessageForHistory;

    if (!isGeminiModel && state.lastGeneratedImageUrl) {
      editPrompt = `![last_generated](${state.lastGeneratedImageUrl})\n\n${originalUserMessageForHistory}`;
    }

    const result = await handleImageGenerationCommand(
      editPrompt,
      undefined, // No command prompt - this is an edit
      editModel,
      state.resultsContext,
      state.chatSessionId,
      state.messages,
      undefined, // No userChoiceCallback - don't ask again
      true, // Already in edit mode
    );

    if (result.tokensUsage) {
      turnTokensUsage = { ...result.tokensUsage };
    }

    return {
      messages: result.messages,
      imageEditionMode: true, // Stay in edit mode
      imageGenerationModelId: editModel,
      lastGeneratedImageUrl:
        result.lastGeneratedImageUrl || state.lastGeneratedImageUrl,
    };
  }

  // Handle image generation
  if (isImageGenerationRequest(state.commandPrompt, state.chatSessionId, state.imageEditionMode)) {
    const result = await handleImageGenerationCommand(
      originalUserMessageForHistory,
      state.commandPrompt,
      state.model.id,
      state.resultsContext,
      state.chatSessionId,
      state.messages,
      state.userChoiceCallback,
      false,
    );

    if (result.tokensUsage) {
      turnTokensUsage = { ...result.tokensUsage };
    }

    return {
      messages: result.messages,
      imageEditionMode: result.imageEditionMode ?? state.imageEditionMode,
      imageGenerationModelId:
        result.imageGenerationModelId ?? state.imageGenerationModelId,
      lastGeneratedImageUrl:
        result.lastGeneratedImageUrl ?? state.lastGeneratedImageUrl,
    };
  }

  // Handle audio analysis requests (automatic, not command-based)
  // This modifies the messages to include transcription before passing to LLM
  let updatedAudioCache = state.audioTranscriptionCache || new Map();
  if (hasAudioContent(originalUserMessageForHistory, state.resultsContext)) {
    const audioResult = await handleAudioAnalysisRequest(
      originalUserMessageForHistory,
      state.model.id,
      state.resultsContext,
      state.messages,
      updatedAudioCache,
    );

    // If this is transcription-only (no user instructions), return directly
    if (audioResult.isTranscriptionOnly) {
      return {
        messages: audioResult.messages,
        audioTranscriptionCache:
          audioResult.audioTranscriptions || updatedAudioCache,
      };
    }

    // Update messages and cache if transcription occurred
    if (audioResult.messages !== state.messages) {
      // Use the enhanced messages with transcription context
      messages = [sys_msg, ...audioResult.messages];
      updatedAudioCache = audioResult.audioTranscriptions || updatedAudioCache;

      // Note: We don't return here - we continue to process with the LLM
      // The transcription is now part of the message context
    } else if (audioResult.audioTranscriptions) {
      // Just update cache if it changed
      updatedAudioCache = audioResult.audioTranscriptions;
    }
  }

  // Handle video analysis requests (automatic, not command-based, requires Gemini)
  // This directly returns the video analysis without passing to LLM
  if (hasVideoContent(originalUserMessageForHistory, state.resultsContext)) {
    const videoResult = await handleVideoAnalysisRequest(
      originalUserMessageForHistory,
      state.model.id,
      state.resultsContext,
      state.messages,
    );

    // If this is analysis-only (no additional user instructions after video), return directly
    if (videoResult.isAnalysisOnly) {
      return {
        messages: videoResult.messages,
      };
    }

    // If user provided additional instructions beyond the video, update messages
    // and continue to LLM (video is already processed, analysis in message)
    if (videoResult.messages !== state.messages) {
      // Video analysis is complete, return it
      return {
        messages: videoResult.messages,
      };
    }
  }

  // Handle PDF analysis requests (automatic, not command-based, requires Gemini)
  // This directly returns the PDF analysis without passing to LLM
  if (hasPdfContent(originalUserMessageForHistory, state.resultsContext, state.includePdf)) {
    const pdfResult = await handlePdfAnalysisRequest(
      originalUserMessageForHistory,
      state.model.id,
      state.resultsContext,
      state.messages,
      state.includePdf,
    );

    // If this is analysis-only (no additional user instructions after PDF), return directly
    if (pdfResult.isAnalysisOnly) {
      return {
        messages: pdfResult.messages,
      };
    }

    // If user provided additional instructions beyond the PDF, update messages
    // and continue to LLM (PDF is already processed, analysis in message)
    if (pdfResult.messages !== state.messages) {
      // PDF analysis is complete, return it
      return {
        messages: pdfResult.messages,
      };
    }
  }

  console.log("Complete systemPrompt :>> ", sys_msg.content);
  console.log("messages in current chat turn :>> ", messages);

  // Bind tools if enabled
  let llm_with_tools = llm;
  if (state.enabledTools) {
    llm_with_tools = llm.bindTools(state.chatTools);
  }

  // Streaming setup
  let streamingContent = "";
  const streamCallback = state.streamingCallback;

  // Workaround: Disable streaming for Gemini models due to tool call parsing issues
  const isGeminiModel = state.model.id.toLowerCase().includes("gemini");
  const shouldStream =
    streamCallback && (!isGeminiModel || !state.toolsEnabled);

  if (isGeminiModel && streamCallback) {
    console.log(
      "⚠️ Streaming disabled for Gemini model to avoid tool call issues",
    );
  }

  // Reasoning not working with current version of langchain/openai
  // const options = {};
  // if (state.model.id.includes("gpt-5") && state.model.thinking)
  //   options["reasoning"] = {
  //     effort: reasoningEffort,
  //     summary: "auto",
  //   };

  if (shouldStream) {
    // Stream the response - use concat to properly accumulate tool call chunks

    const stream = await llm_with_tools.stream(messages);

    for await (const chunk of stream) {
      // Stop streaming if abort signal fired
      if (state.abortSignal?.aborted) break;

      // console.log("chunk :>> ", chunk);
      // Use concat to properly merge chunks including tool_call_chunks
      gathered = gathered !== undefined ? concat(gathered, chunk) : chunk;

      // Capture token usage from streaming chunks
      // For thinking models, callback is disabled so we must read usage_metadata here.
      // For Anthropic/Google models, the handleLLMEnd callback doesn't receive usage data reliably.
      if (
        chunk.usage_metadata &&
        (state.model.thinking ||
          state.model.provider === "Anthropic" ||
          state.model.provider === "Google")
      ) {
        if (chunk.usage_metadata.input_tokens) {
          turnTokensUsage.input_tokens += chunk.usage_metadata.input_tokens;
        }
        if (chunk.usage_metadata.output_tokens) {
          turnTokensUsage.output_tokens += chunk.usage_metadata.output_tokens;
        }
      }

      // Stream content to callback as it arrives
      if (chunk.content) {
        // Handle different content types from different providers
        let textContent = "";

        if (typeof chunk.content === "string") {
          // OpenAI and most providers send plain strings
          textContent = chunk.content;
        } else if (Array.isArray(chunk.content)) {
          // Anthropic/Gemini may send arrays with text and/or function calls
          // Extract only text content, ignore function calls for streaming
          textContent = chunk.content
            .map((item: any) => {
              if (typeof item === "string") {
                return item;
              } else if (item?.type === "text" && item?.text) {
                // Handle Anthropic's {type: "text", text: "..."} format
                return item.text;
              }
              return "";
            })
            .join("");
        }

        if (textContent) {
          streamingContent += textContent;
          streamCallback(textContent);
        }
      }
    }

    // Check for tool calls and notify via callback
    if (
      state.toolUsageCallback &&
      gathered &&
      gathered.tool_calls &&
      gathered.tool_calls.length > 0
    ) {
      const firstToolCall = gathered.tool_calls[0];
      if (firstToolCall && firstToolCall.name) {
        state.toolUsageCallback({
          toolName: firstToolCall.name,
          args: firstToolCall.args,
        });
      }
    }

    // Check if LLM requested fresh audio transcription (streaming path)
    const gatheredContent = gathered?.content?.toString() || "";
    if (
      isRequestingFreshAudioTranscription(gatheredContent) &&
      hasAudioContent(originalUserMessageForHistory, state.resultsContext)
    ) {
      const freshResult = await handleFreshAudioTranscriptionRequest(
        originalUserMessageForHistory,
        state.model.id,
        state.resultsContext,
        state.messages,
        llm_with_tools,
        sys_msg,
      );

      // Update token usage from re-invocation
      if (freshResult.messages.length > state.messages.length) {
        const freshResponse =
          freshResult.messages[freshResult.messages.length - 1];
        if (freshResponse.usage_metadata) {
          turnTokensUsage.input_tokens +=
            freshResponse.usage_metadata.input_tokens || 0;
          turnTokensUsage.output_tokens +=
            freshResponse.usage_metadata.output_tokens || 0;
        }
      }

      return {
        messages: freshResult.messages,
        audioTranscriptionCache: freshResult.audioTranscriptionCache,
      };
    }

    // Return the gathered message (which has proper tool_calls)

    return {
      messages: [...state.messages, gathered],
      audioTranscriptionCache: updatedAudioCache,
    };
  } else {
    // Non-streaming response
    const response = await llm_with_tools.invoke(messages);

    // Capture token usage from response
    // For thinking models, callback is disabled so we must read usage_metadata here.
    // For Anthropic/Google models, the handleLLMEnd callback doesn't receive usage data reliably.
    if (
      response.usage_metadata &&
      (state.model.thinking ||
        state.model.provider === "Anthropic" ||
        state.model.provider === "Google")
    ) {
      if (response.usage_metadata.input_tokens) {
        turnTokensUsage.input_tokens += response.usage_metadata.input_tokens;
      }
      if (response.usage_metadata.output_tokens) {
        turnTokensUsage.output_tokens += response.usage_metadata.output_tokens;
      }
    }

    // Check if LLM requested fresh audio transcription (non-streaming path)
    const responseContent = response.content?.toString() || "";
    if (
      isRequestingFreshAudioTranscription(responseContent) &&
      hasAudioContent(originalUserMessageForHistory, state.resultsContext)
    ) {
      const freshResult = await handleFreshAudioTranscriptionRequest(
        originalUserMessageForHistory,
        state.model.id,
        state.resultsContext,
        state.messages,
        llm_with_tools,
        sys_msg,
      );

      // Update token usage from re-invocation
      if (freshResult.messages.length > state.messages.length) {
        const freshResponse =
          freshResult.messages[freshResult.messages.length - 1];
        if (freshResponse.usage_metadata) {
          turnTokensUsage.input_tokens +=
            freshResponse.usage_metadata.input_tokens || 0;
          turnTokensUsage.output_tokens +=
            freshResponse.usage_metadata.output_tokens || 0;
        }
      }

      return {
        messages: freshResult.messages,
        audioTranscriptionCache: freshResult.audioTranscriptionCache,
      };
    }

    // Check for tool calls and notify via callback
    if (
      state.toolUsageCallback &&
      "tool_calls" in response &&
      Array.isArray(response.tool_calls) &&
      response.tool_calls?.length > 0
    ) {
      // Notify about the first tool being used
      const firstToolCall = response.tool_calls[0];
      if (firstToolCall && firstToolCall.name) {
        state.toolUsageCallback({
          toolName: firstToolCall.name,
          args: firstToolCall.args,
        });
      }
    }

    return {
      messages: [...state.messages, response],
      audioTranscriptionCache: updatedAudioCache,
    };
  }
};

/**
 * Handle invalid tool calls by creating error messages for the assistant
 */
const handleInvalidToolCalls = async (state: typeof ChatAgentState.State) => {
  const lastMessage: any = state.messages.at(-1);
  const invalidToolCalls = lastMessage?.invalid_tool_calls || [];
  const retries = state.invalidToolCallRetries || 0;

  console.log("🚨 Invalid tool calls detected:", invalidToolCalls);
  console.log(`Retry attempt: ${retries + 1}/3`);
  // console.log("Last message:", lastMessage);

  // We need to remove the problematic message with invalid_tool_calls
  // LangGraph's MessagesAnnotation uses a reducer, so we can:
  // 1. Use RemoveMessage to delete the bad message by ID
  // 2. Add new clean messages

  // Get the message ID for removal
  const badMessageId = lastMessage?.id;

  // First, extract any valid text content from the bad message to preserve it
  let preservedContent: AIMessage | null = null;
  if (lastMessage?.content) {
    const textContent = extractTextContent(lastMessage.content);
    if (textContent.trim()) {
      preservedContent = new AIMessage({ content: textContent });
    }
  }

  // Build the list of messages to add
  const messagesToAdd: any[] = [];

  // Remove the bad message if it has an ID
  if (badMessageId) {
    messagesToAdd.push(new RemoveMessage({ id: badMessageId }));
  }

  // Add preserved content if any
  if (preservedContent) {
    messagesToAdd.push(preservedContent);
  }

  // Maximum 3 retries
  const MAX_RETRIES = 3;

  // If we've exceeded max retries, provide a helpful error message and finalize
  if (retries >= MAX_RETRIES) {
    console.log("❌ Max retries exceeded for invalid tool calls");
    const errorSummary = invalidToolCalls
      .map((call: any) => `- ${call.name}: ${call.error}`)
      .join("\n");

    // Create an AI message explaining the failure
    const finalErrorMessage = new AIMessage({
      content: `I apologize, but I'm having trouble calling the tools correctly. After ${MAX_RETRIES} attempts, I encountered these errors:

${errorSummary}

This might be due to the model's limitations with tool calling. Please try:
1. Rephrasing your question more simply
2. Using a more capable model
3. Or ask me without requiring tool usage`,
    });

    messagesToAdd.push(finalErrorMessage);

    return {
      messages: messagesToAdd,
      invalidToolCallRetries: 0, // Reset for next turn
    };
  }

  // Create error messages for each invalid tool call with guidance
  const errorMessages = invalidToolCalls.map(
    (invalidCall: any, index: number) => {
      // Generate a unique tool_call_id if not provided
      const toolCallId = invalidCall.id || `invalid_${Date.now()}_${index}`;

      const errorText = `Tool call error for "${invalidCall.name}":
- Error: ${invalidCall.error}
- Args received: ${
        typeof invalidCall.args === "string"
          ? invalidCall.args
          : JSON.stringify(invalidCall.args)
      }
- Tool call ID: ${toolCallId}

Please review the tool's parameter schema and try again with valid arguments. This is retry ${
        retries + 1
      }/${MAX_RETRIES}.`;

      // Create a proper ToolMessage for LangGraph
      return new ToolMessage({
        content: errorText,
        tool_call_id: toolCallId,
        name: invalidCall.name || "unknown_tool",
      });
    },
  );

  // Add the error messages to the list
  messagesToAdd.push(...errorMessages);

  // Return all messages: RemoveMessage + preserved content + error messages
  // This prevents serialization errors with providers like Gemini
  return {
    messages: messagesToAdd,
    invalidToolCallRetries: retries + 1,
  };
};

/**
 * Tools node with caching
 */
const toolsWithCaching = async (state: typeof ChatAgentState.State) => {
  const toolNode = new ToolNode(state.chatTools);

  // Pass configuration to tools via the config parameter
  const config = {
    configurable: {
      currentResultsContext: state.resultsContext || [],
      addResultsCallback: state.addResultsCallback,
      selectResultsCallback: state.selectResultsCallback,
      expandedResultsCallback: state.expandedResultsCallback,
      needsExpansion: state.needsExpansion || false,
      model: state.model, // Pass model info so tools can reuse the same LLM
      llm: llm, // Pass the initialized LLM instance directly
      toolResultsCache: state.toolResultsCache || {}, // Pass cache for tool result deduplication
      permissions: state.permissions, // Pass privacy mode so sub-agents can inherit it
      // Tool confirmation for sensitive operations
      toolConfirmationCallback: state.toolConfirmationCallback,
      alwaysApprovedTools: state.alwaysApprovedTools,
      // User choice callback for ask_user_choice tool
      userChoiceCallback: state.userChoiceCallback,
      // Enabled tools set for checking tool availability
      enabledTools: state.enabledTools,
    },
  };

  const result = await toolNode.invoke(state, config);

  // Cache tool results
  const toolMessages = result.messages.filter(
    (msg: any) => msg.getType?.() === "tool" || msg._getType?.() === "tool",
  );
  const updatedCache = { ...state.toolResultsCache };

  // Extract skill instructions if the skill tool was called
  let newActiveSkillInstructions = state.activeSkillInstructions;

  // Track if any tools that add results were called
  let needsExpansion = false;
  const resultsAddingTools = ["ask_your_graph", "add_to_context"];

  toolMessages.forEach((msg: any) => {
    if (msg.tool_call_id && msg.content) {
      updatedCache[msg.tool_call_id] = {
        content: msg.content,
        timestamp: Date.now(),
        tool_name: msg.name,
      };

      // Check if this tool adds results to context
      if (resultsAddingTools.includes(msg.name)) {
        needsExpansion = true;
      }

      // If this is the skill tool, extract and store its instructions
      if (msg.name === "live_ai_skills" && msg.content) {
        newActiveSkillInstructions = msg.content;
      }

      // Call toolResponseCallback to capture tool's response for UI display
      if (state.toolResponseCallback) {
        // For help and skills tools, extract only the display portion
        let displayResponse = msg.content;
        if (
          (msg.name === "get_help" || msg.name === "live_ai_skills") &&
          msg.content
        ) {
          const displayMatch = msg.content.match(
            /\[DISPLAY\]([\s\S]*?)\[\/DISPLAY\]/,
          );
          if (displayMatch) {
            displayResponse = displayMatch[1].trim();
          }
        }

        state.toolResponseCallback({
          toolName: msg.name,
          response: displayResponse,
        });
      }
    }
  });

  // If tools added results, trigger expansion in parent component
  // and update the resultsContext with the newly expanded results
  let updatedResultsContext = state.resultsContext;
  let shouldRebuildSystemPrompt = false;

  if (needsExpansion && state.expandedResultsCallback) {
    try {
      // Pass the original results context - the callback tracks additions via ref
      updatedResultsContext = await state.expandedResultsCallback(
        state.resultsContext || [],
      );
      console.log(
        `✅ [Chat Agent] Context expanded, now has ${updatedResultsContext.length} results`,
      );
      shouldRebuildSystemPrompt = true;
    } catch (error) {
      console.error("❌ [Chat Agent] Failed to expand context:", error);
    }
  }

  // Rebuild system prompt if context was expanded
  if (shouldRebuildSystemPrompt && updatedResultsContext) {
    // Build conversation context
    const conversationContext = buildConversationContext(
      state.conversationHistory,
      state.conversationSummary,
    );

    // Build results context with expanded results
    const resultsContextString = buildResultsContext(
      updatedResultsContext,
      state.resultsDescription,
    );

    // Rebuild system prompt
    const lastMessage = state.messages?.at(-1)?.content?.toString() || "";
    const systemPrompt = await buildChatSystemPrompt({
      lastMessage: lastMessage,
      style: state.style,
      commandPrompt: state.commandPrompt,
      toolsEnabled: state.toolsEnabled,
      conversationContext: conversationContext,
      resultsContext: resultsContextString,
      accessMode: state.accessMode,
      isAgentMode: state.isAgentMode,
      activeSkillInstructions: newActiveSkillInstructions,
      enabledTools: state.enabledTools,
      hasAudioContent: hasAudioContent(lastMessage, updatedResultsContext),
      hasVideoContent: hasVideoContent(lastMessage, updatedResultsContext),
      hasPdfContent: hasPdfContent(lastMessage, updatedResultsContext, state.includePdf),
      maxTokensLimit: state.model.advancedParams?.maxTokens,
    });

    sys_msg = new SystemMessage({ content: systemPrompt });
  }

  console.log("🔍 [toolsWithCaching] Returning result to graph");
  return {
    ...result,
    resultsContext: updatedResultsContext,
    toolResultsCache: updatedCache,
    activeSkillInstructions: newActiveSkillInstructions,
    needsExpansion: needsExpansion,
    // Reset invalid tool call counter after successful tool execution
    invalidToolCallRetries: 0,
  };
};

/**
 * Helper function to extract text content from message content
 * Handles different formats from different providers
 */
const extractTextContent = (content: any): string => {
  if (!content) return "";

  if (typeof content === "string") {
    return content;
  }

  if (Array.isArray(content)) {
    // Anthropic/Gemini format: array of objects/strings
    return content
      .map((item: any) => {
        if (typeof item === "string") {
          return item;
        } else if (item?.type === "text" && item?.text) {
          // Handle {type: "text", text: "..."} format
          return item.text;
        }
        // Ignore tool_use and other non-text items
        return "";
      })
      .join("");
  }

  // Fallback to toString for unexpected formats
  return String(content);
};

/**
 * Extract final answer and update conversation history
 */
const finalize = async (state: typeof ChatAgentState.State) => {
  const lastMessage = state.messages.at(-1);
  const finalAnswer = extractTextContent(lastMessage?.content);

  // Update conversation history
  const newHistory = [...(state.conversationHistory || [])];

  // Add user message
  const userMessage = state.messages.find(
    (msg) => msg.getType?.() === "human" || msg._getType?.() === "human",
  );
  if (userMessage) {
    newHistory.push(`\n<User>: ${userMessage.content}`);
  }

  // Add assistant response
  newHistory.push(`\n<Assistant>: ${finalAnswer}`);

  // Persist this turn's token usage to extensionStorage
  const turnDelta = {
    input_tokens:
      turnTokensUsage.input_tokens - initialTokensUsage.input_tokens,
    output_tokens:
      turnTokensUsage.output_tokens - initialTokensUsage.output_tokens,
  };
  if (turnDelta.input_tokens > 0 || turnDelta.output_tokens > 0) {
    updateTokenCounter(state.model.id, turnDelta);
  }

  return {
    finalAnswer,
    conversationHistory: newHistory,
    exchangesSinceLastSummary: (state.exchangesSinceLastSummary || 0) + 1,
    // Cumulative usage for the whole conversation (used to seed the next turn
    // and for the running total shown in the chat header).
    tokensUsage: turnTokensUsage,
    // This turn's usage only (incl. any summarization cost) — for per-message display.
    turnTokensDelta: turnDelta,
    // Preserve state that should persist across turns
    activeSkillInstructions: state.activeSkillInstructions,
    toolResultsCache: state.toolResultsCache,
    audioTranscriptionCache: state.audioTranscriptionCache,
    // Image edition mode state
    imageEditionMode: state.imageEditionMode,
    imageGenerationModelId: state.imageGenerationModelId,
    lastGeneratedImageUrl: state.lastGeneratedImageUrl,
    // Reset retry counter for next turn
    invalidToolCallRetries: 0,
  };
};

// Edges

/**
 * Decide whether to continue to tools or finalize
 */
const shouldContinue = (state: typeof ChatAgentState.State) => {
  const { messages } = state;
  const lastMessage = messages[messages.length - 1];

  const hasToolCalls =
    "tool_calls" in lastMessage &&
    Array.isArray(lastMessage.tool_calls) &&
    lastMessage.tool_calls?.length > 0;

  // Check for invalid tool calls
  const hasInvalidToolCalls =
    "invalid_tool_calls" in lastMessage &&
    Array.isArray(lastMessage.invalid_tool_calls) &&
    lastMessage.invalid_tool_calls?.length > 0;

  if (hasToolCalls) {
    return "tools";
  }

  if (hasInvalidToolCalls) {
    return "handleInvalidToolCalls";
  }

  return "finalize";
};

/**
 * Check if summarization is needed before assistant
 */
const checkSummarization = (state: typeof ChatAgentState.State) => {
  if (shouldSummarize(state)) {
    return "summarizeConversation";
  }
  return "assistant";
};

/**
 * Decide whether to retry or finalize after handling invalid tool calls
 */
const shouldRetryOrFinalize = (state: typeof ChatAgentState.State) => {
  const lastMessage: any = state.messages.at(-1);
  // Check if the last message is an error message (assistant role means max retries exceeded)
  const messageType = lastMessage?.getType?.() || lastMessage?._getType?.();
  if (messageType === "ai") {
    // AI/assistant message means we've hit max retries and created an error message
    return "finalize";
  }
  // Otherwise, retry with the assistant
  return "assistant";
};

// Build Graph

/**
 * Create the chat agent graph
 */
export const createChatGraph = () => {
  const builder = new StateGraph(ChatAgentState);

  builder
    .addNode("loadModel", loadModel)
    .addNode("summarizeConversation", summarizeConversation)
    .addNode("assistant", assistant)
    .addNode("tools", toolsWithCaching)
    .addNode("handleInvalidToolCalls", handleInvalidToolCalls)
    .addNode("finalize", finalize)

    .addEdge(START, "loadModel")
    .addConditionalEdges("loadModel", checkSummarization)
    .addEdge("summarizeConversation", "assistant")
    .addConditionalEdges("assistant", shouldContinue)
    .addEdge("tools", "assistant")
    .addConditionalEdges("handleInvalidToolCalls", shouldRetryOrFinalize)
    .addEdge("finalize", "__end__");

  return builder.compile();
};

/**
 * Export state type for external use
 */
export type ChatAgentStateType = typeof ChatAgentState.State;
