/* eslint-disable @typescript-eslint/no-explicit-any */
import {
  AIMessage,
  AIMessageChunk,
  ToolMessage,
  BaseMessage,
  HumanMessage,
  SystemMessage,
} from '@langchain/core/messages';
import type { MessageContentImageUrl } from '@langchain/core/messages';
import type { ToolCall } from '@langchain/core/messages/tool';
import type {
  ExtendedMessageContent,
  MessageContentComplex,
  ReasoningContentText,
  ToolCallContent,
  ToolCallPart,
  TPayload,
  TMessage,
} from '@/types';
import { Providers, ContentTypes, Constants } from '@/common';
import { processToolOutput } from '@/utils/toonFormat';

interface MediaMessageParams {
  message: {
    role: string;
    content: string;
    name?: string;
    [key: string]: any;
  };
  mediaParts: MessageContentComplex[];
  endpoint?: Providers;
}

/**
 * Formats a message with media content (images, documents, videos, audios) to API payload format.
 *
 * @param params - The parameters for formatting.
 * @returns - The formatted message.
 */
export const formatMediaMessage = ({
  message,
  endpoint,
  mediaParts,
}: MediaMessageParams): {
  role: string;
  content: MessageContentComplex[];
  name?: string;
  [key: string]: any;
} => {
  // Create a new object to avoid mutating the input
  const result: {
    role: string;
    content: MessageContentComplex[];
    name?: string;
    [key: string]: any;
  } = {
    ...message,
    content: [] as MessageContentComplex[],
  };

  if (endpoint === Providers.ANTHROPIC) {
    result.content = [
      ...mediaParts,
      { type: ContentTypes.TEXT, text: message.content },
    ] as MessageContentComplex[];
    return result;
  }

  result.content = [
    { type: ContentTypes.TEXT, text: message.content },
    ...mediaParts,
  ] as MessageContentComplex[];

  return result;
};

interface MessageInput {
  role?: string;
  _name?: string;
  sender?: string;
  text?: string;
  content?: string | MessageContentComplex[];
  image_urls?: MessageContentImageUrl[];
  documents?: MessageContentComplex[];
  videos?: MessageContentComplex[];
  audios?: MessageContentComplex[];
  lc_id?: string[];
  [key: string]: any;
}

interface FormatMessageParams {
  message: MessageInput;
  userName?: string;
  assistantName?: string;
  endpoint?: Providers;
  langChain?: boolean;
}

interface FormattedMessage {
  role: string;
  content: string | MessageContentComplex[];
  name?: string;
  [key: string]: any;
}

/**
 * Formats a message to OpenAI payload format based on the provided options.
 *
 * @param params - The parameters for formatting.
 * @returns - The formatted message.
 */
export const formatMessage = ({
  message,
  userName,
  endpoint,
  assistantName,
  langChain = false,
}: FormatMessageParams):
  | FormattedMessage
  | HumanMessage
  | AIMessage
  | SystemMessage => {
  // eslint-disable-next-line prefer-const
  let { role: _role, _name, sender, text, content: _content, lc_id } = message;
  if (lc_id && lc_id[2] && !langChain) {
    const roleMapping: Record<string, string> = {
      SystemMessage: 'system',
      HumanMessage: 'user',
      AIMessage: 'assistant',
    };
    _role = roleMapping[lc_id[2]] || _role;
  }
  const role =
    _role ??
    (sender != null && sender && sender.toLowerCase() === 'user'
      ? 'user'
      : 'assistant');
  const content = _content ?? text ?? '';
  const formattedMessage: FormattedMessage = {
    role,
    content,
  };

  // Set name fields first
  if (_name != null && _name) {
    formattedMessage.name = _name;
  }

  if (userName != null && userName && formattedMessage.role === 'user') {
    formattedMessage.name = userName;
  }

  if (
    assistantName != null &&
    assistantName &&
    formattedMessage.role === 'assistant'
  ) {
    formattedMessage.name = assistantName;
  }

  if (formattedMessage.name != null && formattedMessage.name) {
    // Conform to API regex: ^[a-zA-Z0-9_-]{1,64}$
    // https://community.openai.com/t/the-format-of-the-name-field-in-the-documentation-is-incorrect/175684/2
    formattedMessage.name = formattedMessage.name.replace(
      /[^a-zA-Z0-9_-]/g,
      '_'
    );

    if (formattedMessage.name.length > 64) {
      formattedMessage.name = formattedMessage.name.substring(0, 64);
    }
  }

  const { image_urls, documents, videos, audios } = message;
  const mediaParts: MessageContentComplex[] = [];

  if (Array.isArray(documents) && documents.length > 0) {
    mediaParts.push(...documents);
  }

  if (Array.isArray(videos) && videos.length > 0) {
    mediaParts.push(...videos);
  }

  if (Array.isArray(audios) && audios.length > 0) {
    mediaParts.push(...audios);
  }

  if (Array.isArray(image_urls) && image_urls.length > 0) {
    mediaParts.push(...image_urls);
  }

  if (mediaParts.length > 0 && role === 'user') {
    const mediaMessage = formatMediaMessage({
      message: {
        ...formattedMessage,
        content:
          typeof formattedMessage.content === 'string'
            ? formattedMessage.content
            : '',
      },
      mediaParts,
      endpoint,
    });

    if (!langChain) {
      return mediaMessage;
    }

    return new HumanMessage(mediaMessage);
  }

  if (!langChain) {
    return formattedMessage;
  }

  if (role === 'user') {
    return new HumanMessage(formattedMessage);
  } else if (role === 'assistant') {
    return new AIMessage(formattedMessage);
  } else {
    return new SystemMessage(formattedMessage);
  }
};

/**
 * Formats an array of messages for LangChain.
 *
 * @param messages - The array of messages to format.
 * @param formatOptions - The options for formatting each message.
 * @returns - The array of formatted LangChain messages.
 */
export const formatLangChainMessages = (
  messages: Array<MessageInput>,
  formatOptions: Omit<FormatMessageParams, 'message' | 'langChain'>
): Array<HumanMessage | AIMessage | SystemMessage> => {
  return messages.map((msg) => {
    const formatted = formatMessage({
      ...formatOptions,
      message: msg,
      langChain: true,
    });
    return formatted as HumanMessage | AIMessage | SystemMessage;
  });
};

interface LangChainMessage {
  lc_kwargs?: {
    additional_kwargs?: Record<string, any>;
    [key: string]: any;
  };
  kwargs?: {
    additional_kwargs?: Record<string, any>;
    [key: string]: any;
  };
  [key: string]: any;
}

/**
 * Formats a LangChain message object by merging properties from `lc_kwargs` or `kwargs` and `additional_kwargs`.
 *
 * @param message - The message object to format.
 * @returns - The formatted LangChain message.
 */
export const formatFromLangChain = (
  message: LangChainMessage
): Record<string, any> => {
  const kwargs = message.lc_kwargs ?? message.kwargs ?? {};
  const { additional_kwargs = {}, ...message_kwargs } = kwargs;
  return {
    ...message_kwargs,
    ...additional_kwargs,
  };
};

/**
 * Helper function to format an assistant message
 * @param message The message to format
 * @returns Array of formatted messages
 */
function formatAssistantMessage(
  message: Partial<TMessage>
): Array<AIMessage | ToolMessage> {
  const formattedMessages: Array<AIMessage | ToolMessage> = [];
  let currentContent: MessageContentComplex[] = [];
  let lastAIMessage: AIMessage | null = null;
  let hasReasoning = false;

  if (Array.isArray(message.content)) {
    for (const part of message.content as Array<
      MessageContentComplex | undefined | null
    >) {
      if (part == null) {
        continue;
      }
      if (part.type === ContentTypes.TEXT && part.tool_call_ids) {
        /*
        If there's pending content, it needs to be aggregated as a single string to prepare for tool calls.
        For Anthropic models, the "tool_calls" field on a message is only respected if content is a string.
        */
        if (currentContent.length > 0) {
          let content = currentContent.reduce((acc, curr) => {
            if (curr.type === ContentTypes.TEXT) {
              return `${acc}${String(curr[ContentTypes.TEXT] ?? '')}\n`;
            }
            return acc;
          }, '');
          content =
            `${content}\n${part[ContentTypes.TEXT] ?? part.text ?? ''}`.trim();
          lastAIMessage = new AIMessage({ content });
          formattedMessages.push(lastAIMessage);
          currentContent = [];
          continue;
        }
        // Create a new AIMessage with this text and prepare for tool calls
        lastAIMessage = new AIMessage({
          content: part.text != null ? part.text : '',
        });
        formattedMessages.push(lastAIMessage);
      } else if (part.type === ContentTypes.TOOL_CALL) {
        // Skip malformed tool call entries without tool_call property
        if (part.tool_call == null) {
          continue;
        }

        // Note: `tool_calls` list is defined when constructed by `AIMessage` class, and outputs should be excluded from it
        const {
          output,
          args: _args,
          ..._tool_call
        } = part.tool_call as ToolCallPart;

        // Skip invalid tool calls that have no name AND no output
        if (
          _tool_call.name == null ||
          (_tool_call.name === '' && (output == null || output === ''))
        ) {
          continue;
        }

        if (!lastAIMessage) {
          // "Heal" the payload by creating an AIMessage to precede the tool call
          lastAIMessage = new AIMessage({ content: '' });
          formattedMessages.push(lastAIMessage);
        }

        const tool_call: ToolCallPart = _tool_call;
        // TODO: investigate; args as dictionary may need to be providers-or-tool-specific
        let args: any = _args;
        try {
          if (typeof _args === 'string') {
            args = JSON.parse(_args);
          }
        } catch {
          if (typeof _args === 'string') {
            args = { input: _args };
          }
        }

        tool_call.args = args;
        if (!lastAIMessage.tool_calls) {
          lastAIMessage.tool_calls = [];
        }
        lastAIMessage.tool_calls.push(tool_call as ToolCall);

        // Apply TOON compression to historical tool outputs for context efficiency
        // processToolOutput handles: JSON→TOON conversion, already-TOON detection (skip), truncation
        // Skip TOON for content_tool — its output is line-numbered source code that must stay verbatim
        // for accurate edit (str_replace) matching.
        const processedOutput =
          output != null
            ? processToolOutput(output, {
                enableToon: tool_call.name !== 'content_tool',
              }).content
            : '';

        formattedMessages.push(
          new ToolMessage({
            tool_call_id: tool_call.id ?? '',
            name: tool_call.name,
            content: processedOutput,
          })
        );
      } else if (
        part.type === ContentTypes.THINK ||
        part.type === ContentTypes.THINKING ||
        part.type === ContentTypes.REASONING_CONTENT ||
        part.type === 'redacted_thinking'
      ) {
        hasReasoning = true;
        continue;
      } else if (
        part.type === ContentTypes.ERROR ||
        part.type === ContentTypes.AGENT_UPDATE
      ) {
        continue;
      } else {
        if (
          part.type === ContentTypes.TEXT &&
          !String(part.text ?? '').trim()
        ) {
          continue;
        }
        currentContent.push(part);
      }
    }
  }

  if (hasReasoning && currentContent.length > 0) {
    const content = currentContent
      .reduce((acc, curr) => {
        if (curr.type === ContentTypes.TEXT) {
          return `${acc}${String(curr[ContentTypes.TEXT] ?? '')}\n`;
        }
        return acc;
      }, '')
      .trim();

    if (content) {
      formattedMessages.push(new AIMessage({ content }));
    }
  } else if (currentContent.length > 0) {
    formattedMessages.push(new AIMessage({ content: currentContent }));
  }

  return formattedMessages;
}

/**
 * Labels all agent content for parallel patterns (fan-out/fan-in)
 * Groups consecutive content by agent and wraps with clear labels
 */
function labelAllAgentContent(
  contentParts: MessageContentComplex[],
  agentIdMap: Record<number, string>,
  agentNames?: Record<string, string>
): MessageContentComplex[] {
  const result: MessageContentComplex[] = [];
  let currentAgentId: string | undefined;
  let agentContentBuffer: MessageContentComplex[] = [];

  const flushAgentBuffer = (): void => {
    if (agentContentBuffer.length === 0) {
      return;
    }

    if (currentAgentId != null && currentAgentId !== '') {
      const agentName = (agentNames?.[currentAgentId] ?? '') || currentAgentId;
      const formattedParts: string[] = [];

      formattedParts.push(`--- ${agentName} ---`);

      for (const part of agentContentBuffer) {
        if (part.type === ContentTypes.THINK) {
          const thinkContent = (part as ReasoningContentText).think || '';
          if (thinkContent) {
            formattedParts.push(
              `${agentName}: ${JSON.stringify({
                type: 'think',
                think: thinkContent,
              })}`
            );
          }
        } else if (part.type === ContentTypes.TEXT) {
          const textContent: string = part.text ?? '';
          if (textContent) {
            formattedParts.push(`${agentName}: ${textContent}`);
          }
        } else if (part.type === ContentTypes.TOOL_CALL) {
          formattedParts.push(
            `${agentName}: ${JSON.stringify({
              type: 'tool_call',
              tool_call: (part as ToolCallContent).tool_call,
            })}`
          );
        }
      }

      formattedParts.push(`--- End of ${agentName} ---`);

      // Create a single text content part with all agent content
      result.push({
        type: ContentTypes.TEXT,
        text: formattedParts.join('\n\n'),
      } as MessageContentComplex);
    } else {
      // No agent ID, pass through as-is
      result.push(...agentContentBuffer);
    }

    agentContentBuffer = [];
  };

  for (let i = 0; i < contentParts.length; i++) {
    const part = contentParts[i];
    const agentId = agentIdMap[i];

    // If agent changed, flush previous buffer
    if (agentId !== currentAgentId && currentAgentId !== undefined) {
      flushAgentBuffer();
    }

    currentAgentId = agentId;
    agentContentBuffer.push(part);
  }

  // Flush any remaining content
  flushAgentBuffer();

  return result;
}

/**
 * Groups content parts by agent and formats them with agent labels
 * This preprocesses multi-agent content to prevent identity confusion
 *
 * @param contentParts - The content parts from a run
 * @param agentIdMap - Map of content part index to agent ID
 * @param agentNames - Optional map of agent ID to display name
 * @param options - Configuration options
 * @param options.labelNonTransferContent - If true, labels all agent transitions (for parallel patterns)
 * @returns Modified content parts with agent labels where appropriate
 */
export const labelContentByAgent = (
  contentParts: MessageContentComplex[],
  agentIdMap?: Record<number, string>,
  agentNames?: Record<string, string>,
  options?: { labelNonTransferContent?: boolean }
): MessageContentComplex[] => {
  if (!agentIdMap || Object.keys(agentIdMap).length === 0) {
    return contentParts;
  }

  // If labelNonTransferContent is true, use a different strategy for parallel patterns
  if (options?.labelNonTransferContent === true) {
    return labelAllAgentContent(contentParts, agentIdMap, agentNames);
  }

  const result: MessageContentComplex[] = [];
  let currentAgentId: string | undefined;
  let agentContentBuffer: MessageContentComplex[] = [];
  let transferToolCallIndex: number | undefined;
  let transferToolCallId: string | undefined;

  const flushAgentBuffer = (): void => {
    if (agentContentBuffer.length === 0) {
      return;
    }

    // If this is content from a transferred agent, format it specially
    if (
      currentAgentId != null &&
      currentAgentId !== '' &&
      transferToolCallIndex !== undefined
    ) {
      const agentName = (agentNames?.[currentAgentId] ?? '') || currentAgentId;
      const formattedParts: string[] = [];

      formattedParts.push(`--- Transfer to ${agentName} ---`);

      for (const part of agentContentBuffer) {
        if (part.type === ContentTypes.THINK) {
          formattedParts.push(
            `${agentName}: ${JSON.stringify({
              type: 'think',
              think: (part as ReasoningContentText).think,
            })}`
          );
        } else if ('text' in part && part.type === ContentTypes.TEXT) {
          const textContent: string = part.text ?? '';
          if (textContent) {
            formattedParts.push(
              `${agentName}: ${JSON.stringify({
                type: 'text',
                text: textContent,
              })}`
            );
          }
        } else if (part.type === ContentTypes.TOOL_CALL) {
          formattedParts.push(
            `${agentName}: ${JSON.stringify({
              type: 'tool_call',
              tool_call: (part as ToolCallContent).tool_call,
            })}`
          );
        }
      }

      formattedParts.push(`--- End of ${agentName} response ---`);

      // Find the tool call that triggered this transfer and update its output
      if (transferToolCallIndex < result.length) {
        const transferToolCall = result[transferToolCallIndex];
        if (
          transferToolCall.type === ContentTypes.TOOL_CALL &&
          transferToolCall.tool_call?.id === transferToolCallId
        ) {
          transferToolCall.tool_call.output = formattedParts.join('\n\n');
        }
      }
    } else {
      // Not from a transfer, add as-is
      result.push(...agentContentBuffer);
    }

    agentContentBuffer = [];
    transferToolCallIndex = undefined;
    transferToolCallId = undefined;
  };

  for (let i = 0; i < contentParts.length; i++) {
    const part = contentParts[i];
    const agentId = agentIdMap[i];

    // Check if this is a transfer tool call
    const isTransferTool =
      (part.type === ContentTypes.TOOL_CALL &&
        (part as ToolCallContent).tool_call?.name?.startsWith(
          'lc_transfer_to_'
        )) ??
      false;

    // If agent changed, flush previous buffer
    if (agentId !== currentAgentId && currentAgentId !== undefined) {
      flushAgentBuffer();
    }

    currentAgentId = agentId;

    if (isTransferTool) {
      // Flush any existing buffer first
      flushAgentBuffer();
      // Add the transfer tool call to result
      result.push(part);
      // Mark that the next agent's content should be captured
      transferToolCallIndex = result.length - 1;
      transferToolCallId = (part as ToolCallContent).tool_call?.id;
      currentAgentId = undefined; // Reset to capture the next agent
    } else {
      agentContentBuffer.push(part);
    }
  }

  flushAgentBuffer();

  return result;
};

/** Extracts tool names from a tool_search output JSON string. */
function extractToolNamesFromSearchOutput(output: string): string[] {
  try {
    const parsed: unknown = JSON.parse(output);
    if (
      typeof parsed === 'object' &&
      parsed !== null &&
      Array.isArray((parsed as Record<string, unknown>).tools)
    ) {
      return (
        (parsed as Record<string, unknown>).tools as Array<{ name?: string }>
      )
        .map((t) => t.name)
        .filter((name): name is string => typeof name === 'string');
    }
  } catch {
    /** Output may have warnings prepended, try to find JSON within it */
    const jsonMatch = output.match(/\{[\s\S]*\}/);
    if (jsonMatch) {
      try {
        const parsed: unknown = JSON.parse(jsonMatch[0]);
        if (
          typeof parsed === 'object' &&
          parsed !== null &&
          Array.isArray((parsed as Record<string, unknown>).tools)
        ) {
          return (
            (parsed as Record<string, unknown>).tools as Array<{
              name?: string;
            }>
          )
            .map((t) => t.name)
            .filter((name): name is string => typeof name === 'string');
        }
      } catch {
        /* ignore */
      }
    }
  }
  return [];
}

/** Extracts the skillName from a skill tool_call's args (string or object). */
function extractSkillName(args: unknown): string | undefined {
  let parsed: Record<string, unknown> | undefined;
  if (typeof args === 'string') {
    try {
      parsed = JSON.parse(args) as Record<string, unknown>;
    } catch {
      /* malformed args — skip */
    }
  } else {
    parsed = args as Record<string, unknown> | undefined;
  }
  const name = parsed?.skillName;
  return typeof name === 'string' && name !== '' ? name : undefined;
}

/**
 * Formats an array of messages for LangChain, handling tool calls and creating ToolMessage instances.
 *
 * @param payload - The array of messages to format.
 * @param indexTokenCountMap - Optional map of message indices to token counts.
 * @param tools - Optional set of tool names that are allowed in the request.
 * @param skills - Optional map of skill name to body for reconstructing skill HumanMessages.
 *   When SkillTool is invoked, the body is injected as a HumanMessage into LangGraph state
 *   but NOT persisted to conversation history. On follow-up runs the skill body is lost.
 *   Pass `skills` to reconstruct the HumanMessage at the right position after each skill
 *   ToolMessage. Upstream PRs #96 / #97.
 * @returns - Object containing formatted messages and updated indexTokenCountMap if provided.
 */
export const formatAgentMessages = (
  payload: TPayload,
  indexTokenCountMap?: Record<number, number | undefined>,
  tools?: Set<string>,
  skills?: Map<string, string>
): {
  messages: Array<HumanMessage | AIMessage | SystemMessage | ToolMessage>;
  indexTokenCountMap?: Record<number, number>;
} => {
  const messages: Array<
    HumanMessage | AIMessage | SystemMessage | ToolMessage
  > = [];
  // If indexTokenCountMap is provided, create a new map to track the updated indices
  const updatedIndexTokenCountMap: Record<number, number> = {};
  // Keep track of the mapping from original payload indices to result indices
  const indexMapping: Record<number, number[] | undefined> = {};

  /**
   * Create a mutable copy of the tools set that can be expanded dynamically.
   * When we encounter tool_search results, we add discovered tools to this set,
   * making their subsequent tool calls valid.
   */
  const discoveredTools = tools ? new Set(tools) : undefined;

  // Process messages with tool conversion if tools set is provided
  for (let i = 0; i < payload.length; i++) {
    const message = payload[i];
    // Q: Store the current length of messages to track where this payload message starts in the result?
    // const startIndex = messages.length;
    if (typeof message.content === 'string') {
      message.content = [
        { type: ContentTypes.TEXT, [ContentTypes.TEXT]: message.content },
      ];
    }
    if (message.role !== 'assistant') {
      messages.push(
        formatMessage({
          message: message as MessageInput,
          langChain: true,
        }) as HumanMessage | AIMessage | SystemMessage
      );

      // Update the index mapping for this message
      indexMapping[i] = [messages.length - 1];
      continue;
    }

    // For assistant messages, track the starting index before processing
    const startMessageIndex = messages.length;

    /**
     * If tools set is provided, process tool_calls:
     * - Keep valid tool_calls (tools in the set or dynamically discovered)
     * - Convert invalid tool_calls to string representation for context preservation
     * - Dynamically expand the set when tool_search results are encountered
     */
    let processedMessage = message;
    let pendingSkillNames: Set<string> | undefined;
    if (discoveredTools) {
      const content = message.content;
      if (content && Array.isArray(content)) {
        const filteredContent: typeof content = [];
        const invalidToolCallIds = new Set<string>();
        const invalidToolStrings: string[] = [];

        for (const part of content) {
          if (part.type !== ContentTypes.TOOL_CALL) {
            filteredContent.push(part);
            continue;
          }

          /** Skip malformed tool_call entries */
          if (
            part.tool_call == null ||
            part.tool_call.name == null ||
            part.tool_call.name === ''
          ) {
            if (
              typeof part.tool_call?.id === 'string' &&
              part.tool_call.id !== ''
            ) {
              invalidToolCallIds.add(part.tool_call.id);
            }
            continue;
          }

          const toolName = part.tool_call.name;

          /**
           * If this is a tool_search result with output, extract discovered tool names
           * and add them to the discoveredTools set for subsequent validation.
           */
          if (
            toolName === Constants.TOOL_SEARCH &&
            typeof part.tool_call.output === 'string' &&
            part.tool_call.output !== ''
          ) {
            const extracted = extractToolNamesFromSearchOutput(
              part.tool_call.output
            );
            for (const name of extracted) {
              discoveredTools.add(name);
            }
          }

          if (discoveredTools.has(toolName)) {
            /** Valid tool - keep it */
            filteredContent.push(part);
            // Skill body reconstruction (upstream PRs #96/#97): when a
            // SkillTool tool_call is present in this assistant turn, mark
            // the skill name so we can reconstruct the body HumanMessage
            // after the assistant turn has been pushed.
            if (toolName === Constants.SKILL_TOOL && skills?.size) {
              const skillName = extractSkillName(part.tool_call.args);
              if (skillName) {
                (pendingSkillNames ??= new Set()).add(skillName);
              }
            }
          } else {
            /** Invalid tool - convert to string for context preservation */
            if (
              typeof part.tool_call.id === 'string' &&
              part.tool_call.id !== ''
            ) {
              invalidToolCallIds.add(part.tool_call.id);
            }
            const output = part.tool_call.output ?? '';
            invalidToolStrings.push(`Tool: ${toolName}, ${output}`);
          }
        }

        /** Remove tool_call_ids references to invalid tools from text parts */
        if (invalidToolCallIds.size > 0) {
          for (const part of filteredContent) {
            if (
              part.type === ContentTypes.TEXT &&
              Array.isArray(part.tool_call_ids)
            ) {
              part.tool_call_ids = part.tool_call_ids.filter(
                (id: string) => !invalidToolCallIds.has(id)
              );
              if (part.tool_call_ids.length === 0) {
                delete part.tool_call_ids;
              }
            }
          }
        }

        /** Append invalid tool strings to the content for context preservation */
        if (invalidToolStrings.length > 0) {
          /** Find the last text part or create one */
          let lastTextPartIndex = -1;
          for (let j = filteredContent.length - 1; j >= 0; j--) {
            if (filteredContent[j].type === ContentTypes.TEXT) {
              lastTextPartIndex = j;
              break;
            }
          }

          const invalidToolText = invalidToolStrings.join('\n');
          if (lastTextPartIndex >= 0) {
            const lastTextPart = filteredContent[lastTextPartIndex] as {
              type: string;
              [ContentTypes.TEXT]?: string;
              text?: string;
            };
            const existingText =
              lastTextPart[ContentTypes.TEXT] ?? lastTextPart.text ?? '';
            filteredContent[lastTextPartIndex] = {
              ...lastTextPart,
              [ContentTypes.TEXT]: existingText
                ? `${existingText}\n${invalidToolText}`
                : invalidToolText,
            };
          } else {
            /** No text part exists, create one */
            filteredContent.push({
              type: ContentTypes.TEXT,
              [ContentTypes.TEXT]: invalidToolText,
            });
          }
        }

        /** Use filtered content if we made any changes */
        if (
          filteredContent.length !== content.length ||
          invalidToolStrings.length > 0
        ) {
          processedMessage = { ...message, content: filteredContent };
        }
      }
    }

    // When tools filtering is off, still detect skill tool_calls for body reconstruction.
    if (!discoveredTools && skills?.size) {
      const content = processedMessage.content;
      if (Array.isArray(content)) {
        for (const part of content) {
          if (
            part.type !== ContentTypes.TOOL_CALL ||
            part.tool_call?.name !== Constants.SKILL_TOOL
          ) {
            continue;
          }
          const skillName = extractSkillName(part.tool_call.args);
          if (skillName) {
            (pendingSkillNames ??= new Set()).add(skillName);
          }
        }
      }
    }

    // Process the assistant message using the helper function
    const formattedMessages = formatAssistantMessage(processedMessage);
    messages.push(...formattedMessages);

    // Capture index range BEFORE skill body injection so injected
    // HumanMessages are excluded from the assistant's token distribution
    // (upstream PR #97 fix).
    const endMessageIndex = messages.length;

    if (pendingSkillNames?.size) {
      for (const skillName of pendingSkillNames) {
        const body = skills?.get(skillName);
        if (body) {
          messages.push(
            new HumanMessage({
              content: body,
              additional_kwargs: {
                role: 'user',
                isMeta: true,
                source: 'skill',
                skillName,
              },
            })
          );
        }
      }
    }

    const resultIndices = [];
    for (let j = startMessageIndex; j < endMessageIndex; j++) {
      resultIndices.push(j);
    }
    indexMapping[i] = resultIndices;
  }

  if (indexTokenCountMap) {
    for (
      let originalIndex = 0;
      originalIndex < payload.length;
      originalIndex++
    ) {
      const resultIndices = indexMapping[originalIndex] || [];
      const tokenCount = indexTokenCountMap[originalIndex];

      if (tokenCount === undefined) {
        continue;
      }

      const msgCount = resultIndices.length;
      if (msgCount === 1) {
        updatedIndexTokenCountMap[resultIndices[0]] = tokenCount;
        continue;
      }

      if (msgCount < 2) {
        continue;
      }

      let totalLength = 0;
      const lastIdx = msgCount - 1;
      const lengths = new Array<number>(msgCount);
      for (let k = 0; k < msgCount; k++) {
        const msg = messages[resultIndices[k]];
        const { content } = msg;
        let len = 0;
        if (typeof content === 'string') {
          len = content.length;
        } else if (Array.isArray(content)) {
          for (const part of content as Array<
            Record<string, unknown> | string | undefined
          >) {
            if (typeof part === 'string') {
              len += part.length;
            } else if (part != null && typeof part === 'object') {
              const val = part.text ?? part.content;
              if (typeof val === 'string') {
                len += val.length;
              }
            }
          }
        }
        const toolCalls = (msg as AIMessage).tool_calls;
        if (Array.isArray(toolCalls)) {
          for (const tc of toolCalls as Array<Record<string, unknown>>) {
            if (typeof tc.name === 'string') {
              len += tc.name.length;
            }
            const { args } = tc;
            if (typeof args === 'string') {
              len += args.length;
            } else if (args != null) {
              len += JSON.stringify(args).length;
            }
          }
        }
        lengths[k] = len;
        totalLength += len;
      }

      if (totalLength === 0) {
        const countPerMessage = Math.floor(tokenCount / msgCount);
        for (let k = 0; k < lastIdx; k++) {
          updatedIndexTokenCountMap[resultIndices[k]] = countPerMessage;
        }
        updatedIndexTokenCountMap[resultIndices[lastIdx]] =
          tokenCount - countPerMessage * lastIdx;
      } else {
        let distributed = 0;
        for (let k = 0; k < lastIdx; k++) {
          const share = Math.floor((lengths[k] / totalLength) * tokenCount);
          updatedIndexTokenCountMap[resultIndices[k]] = share;
          distributed += share;
        }
        updatedIndexTokenCountMap[resultIndices[lastIdx]] =
          tokenCount - distributed;
      }
    }
  }

  return {
    messages,
    indexTokenCountMap: indexTokenCountMap
      ? updatedIndexTokenCountMap
      : undefined,
  };
};

/**
 * Adds a value at key 0 for system messages and shifts all key indices by one in an indexTokenCountMap.
 * This is useful when adding a system message at the beginning of a conversation.
 *
 * @param indexTokenCountMap - The original map of message indices to token counts
 * @param instructionsTokenCount - The token count for the system message to add at index 0
 * @returns A new map with the system message at index 0 and all other indices shifted by 1
 */
export function shiftIndexTokenCountMap(
  indexTokenCountMap: Record<number, number>,
  instructionsTokenCount: number
): Record<number, number> {
  // Create a new map to avoid modifying the original
  const shiftedMap: Record<number, number> = {};
  shiftedMap[0] = instructionsTokenCount;

  // Shift all existing indices by 1
  for (const [indexStr, tokenCount] of Object.entries(indexTokenCountMap)) {
    const index = Number(indexStr);
    shiftedMap[index + 1] = tokenCount;
  }

  return shiftedMap;
}

/** Block types that contain binary image data and must be preserved structurally. */
const IMAGE_BLOCK_TYPES = new Set(['image_url', 'image']);

/** Checks whether a BaseMessage is a tool-role message. */
const isToolMessage = (m: BaseMessage): boolean =>
  m instanceof ToolMessage || ('role' in m && (m as any).role === 'tool');

/** Flushes accumulated text chunks into `parts` as a single text block. */
function flushTextChunks(
  textChunks: string[],
  parts: MessageContentComplex[]
): void {
  if (textChunks.length === 0) {
    return;
  }
  parts.push({
    type: ContentTypes.TEXT,
    text: textChunks.join('\n'),
  } as MessageContentComplex);
  textChunks.length = 0;
}

/**
 * Appends a single message's content to the running `textChunks` / `parts`
 * accumulators.  Image blocks are shallow-copied into `parts` as-is so that
 * binary data (base64 images) never becomes text tokens.  All other block
 * types are serialized to text — unrecognized types are JSON-serialized
 * rather than silently dropped.
 *
 * When `content` is an array containing tool_use blocks, `tool_calls` is NOT
 * additionally serialized (avoiding double output).  `tool_calls` is used as
 * a fallback when `content` is a plain string or an array with no tool_use.
 */
function appendMessageContent(
  msg: BaseMessage,
  role: string,
  textChunks: string[],
  parts: MessageContentComplex[]
): void {
  const { content } = msg;

  if (typeof content === 'string') {
    if (content) {
      textChunks.push(`${role}: ${content}`);
    }
    appendToolCalls(msg, role, textChunks);
    return;
  }

  if (!Array.isArray(content)) {
    appendToolCalls(msg, role, textChunks);
    return;
  }

  let hasToolUseBlock = false;

  for (const block of content as ExtendedMessageContent[]) {
    if (IMAGE_BLOCK_TYPES.has(block.type ?? '')) {
      flushTextChunks(textChunks, parts);
      parts.push({ ...block } as MessageContentComplex);
      continue;
    }

    if (block.type === 'tool_use') {
      hasToolUseBlock = true;
      textChunks.push(
        `${role}: [tool_use] ${String(block.name ?? '')} ${JSON.stringify(block.input ?? {})}`
      );
      continue;
    }

    const text = block.text ?? block.input;
    if (typeof text === 'string' && text) {
      textChunks.push(`${role}: ${text}`);
      continue;
    }

    // Fallback: serialize unrecognized block types to preserve context
    if (block.type != null && block.type !== '') {
      textChunks.push(`${role}: [${block.type}] ${JSON.stringify(block)}`);
    }
  }

  // If content array had no tool_use blocks, fall back to tool_calls metadata
  // (handles edge case: empty content array with tool_calls populated)
  if (!hasToolUseBlock) {
    appendToolCalls(msg, role, textChunks);
  }
}

function appendToolCalls(
  msg: BaseMessage,
  role: string,
  textChunks: string[]
): void {
  if (role !== 'AI') {
    return;
  }
  const aiMsg = msg as AIMessage;
  if (!aiMsg.tool_calls || aiMsg.tool_calls.length === 0) {
    return;
  }
  for (const tc of aiMsg.tool_calls) {
    textChunks.push(`AI: [tool_call] ${tc.name}(${JSON.stringify(tc.args)})`);
  }
}

/**
 * Ensures compatibility when switching from a non-thinking agent to a thinking-enabled agent.
 * Converts AI messages with tool calls (that lack thinking/reasoning blocks) into buffer strings,
 * avoiding the thinking block signature requirement.
 *
 * Recognizes the following as valid thinking/reasoning blocks:
 * - ContentTypes.THINKING (Anthropic)
 * - ContentTypes.REASONING_CONTENT (Bedrock)
 * - ContentTypes.REASONING (VertexAI / Google)
 * - 'redacted_thinking'
 *
 * @param messages - Array of messages to process
 * @param provider - The provider being used (unused but kept for future compatibility)
 * @returns The messages array with tool sequences converted to buffer strings if necessary
 */
export function ensureThinkingBlockInMessages(
  messages: BaseMessage[],
  _provider: Providers
): BaseMessage[] {
  if (messages.length === 0) {
    return messages;
  }

  // If the last message is already a HumanMessage, there is no trailing tool
  // sequence to convert — return early to preserve prompt caching and avoid
  // redundant token overhead from re-processing the entire history.
  const lastMsg = messages[messages.length - 1];
  const lastIsHuman =
    lastMsg instanceof HumanMessage ||
    ('role' in lastMsg && (lastMsg as any).role === 'user');
  if (lastIsHuman) {
    return messages;
  }

  const result: BaseMessage[] = [];
  let i = 0;

  while (i < messages.length) {
    const msg = messages[i];
    /** Detect AI messages by instanceof OR by role, in case cache-control cloning
     produced a plain object that lost the LangChain prototype. */
    const isAI =
      msg instanceof AIMessage ||
      msg instanceof AIMessageChunk ||
      ('role' in msg && (msg as any).role === 'assistant');

    if (!isAI) {
      result.push(msg);
      i++;
      continue;
    }

    const aiMsg = msg as AIMessage | AIMessageChunk;
    const hasToolCalls = aiMsg.tool_calls && aiMsg.tool_calls.length > 0;
    const contentIsArray = Array.isArray(aiMsg.content);

    // Check if the message has tool calls or tool_use content
    let hasToolUse = hasToolCalls ?? false;
    let hasThinkingBlock = false;

    if (contentIsArray && aiMsg.content.length > 0) {
      for (const c of aiMsg.content as ExtendedMessageContent[]) {
        if (typeof c !== 'object') {
          continue;
        }
        if (c.type === 'tool_use') {
          hasToolUse = true;
        } else if (
          c.type === ContentTypes.THINKING ||
          c.type === ContentTypes.REASONING_CONTENT ||
          c.type === ContentTypes.REASONING ||
          c.type === 'redacted_thinking'
        ) {
          hasThinkingBlock = true;
        }
        if (hasToolUse && hasThinkingBlock) {
          break;
        }
      }
    }

    // Bedrock also stores reasoning in additional_kwargs (may not be in content array)
    if (
      !hasThinkingBlock &&
      aiMsg.additional_kwargs.reasoning_content != null
    ) {
      hasThinkingBlock = true;
    }

    // If message has tool use but no thinking block, check whether this is a
    // continuation of a thinking-enabled agent's chain before converting.
    // Bedrock reasoning models can produce multiple AI→Tool rounds after an
    // initial reasoning response: the first AI message has reasoning_content,
    // but follow-ups have content: "" with only tool_calls. These are the
    // same agent's turn and must NOT be converted to HumanMessages.
    if (hasToolUse && !hasThinkingBlock) {
      // Walk backwards — if an earlier AI message in the same chain (before
      // the nearest HumanMessage) has a thinking/reasoning block, this is a
      // continuation of a thinking-enabled turn, not a non-thinking handoff.
      if (chainHasThinkingBlock(messages, i)) {
        result.push(msg);
        i++;
        continue;
      }

      // Build structured content in a single pass over the AI + following
      // ToolMessages — preserves image blocks as-is to avoid serializing
      // binary data as text (which caused 174× token amplification).
      const parts: MessageContentComplex[] = [];
      const textChunks: string[] = ['[Previous agent context]'];

      appendMessageContent(msg, 'AI', textChunks, parts);

      let j = i + 1;
      while (j < messages.length && isToolMessage(messages[j])) {
        appendMessageContent(messages[j], 'Tool', textChunks, parts);
        j++;
      }

      flushTextChunks(textChunks, parts);
      result.push(new HumanMessage({ content: parts }));
      i = j;
    } else {
      // Keep the message as is
      result.push(msg);
      i++;
    }
  }

  return result;
}

/**
 * Walks backwards from `currentIndex` through the message array to check
 * whether an earlier AI message in the same "chain" (no HumanMessage boundary)
 * contains a thinking/reasoning block.
 *
 * A "chain" is a contiguous sequence of AI + Tool messages with no intervening
 * HumanMessage. Bedrock reasoning models produce reasoning on the first AI
 * response, then issue follow-up tool calls with `content: ""` and no
 * reasoning block. These follow-ups are part of the same thinking-enabled
 * turn and should not be converted.
 */
function chainHasThinkingBlock(
  messages: BaseMessage[],
  currentIndex: number
): boolean {
  for (let k = currentIndex - 1; k >= 0; k--) {
    const prev = messages[k];

    // HumanMessage = turn boundary — stop searching
    if (
      prev instanceof HumanMessage ||
      ('role' in prev && (prev as any).role === 'user')
    ) {
      return false;
    }

    // Check AI messages for thinking/reasoning blocks
    const isPrevAI =
      prev instanceof AIMessage ||
      prev instanceof AIMessageChunk ||
      ('role' in prev && (prev as any).role === 'assistant');

    if (isPrevAI) {
      const prevAiMsg = prev as AIMessage | AIMessageChunk;

      if (Array.isArray(prevAiMsg.content) && prevAiMsg.content.length > 0) {
        const content = prevAiMsg.content as ExtendedMessageContent[];
        if (
          content.some(
            (c) =>
              typeof c === 'object' &&
              (c.type === ContentTypes.THINKING ||
                c.type === ContentTypes.REASONING_CONTENT ||
                c.type === ContentTypes.REASONING ||
                c.type === 'redacted_thinking')
          )
        ) {
          return true;
        }
      }

      // Bedrock also stores reasoning in additional_kwargs
      if (prevAiMsg.additional_kwargs.reasoning_content != null) {
        return true;
      }
    }

    // ToolMessages are part of the chain — keep walking back
  }

  return false;
}