import {
  AIMessage,
  BaseMessage,
  ToolMessage,
  UsageMetadata,
} from '@langchain/core/messages';
import type {
  ThinkingContentText,
  MessageContentComplex,
  ReasoningContentText,
} from '@/types/stream';
import type { TokenCounter } from '@/types/run';
import { ContentTypes, Constants, Providers } from '@/common';
import type { ContextPruningConfig } from '@/types/graph';
import {
  calculateMaxToolResultChars,
  truncateToolResultContent,
  truncateToolInput,
} from '@/utils/truncation';
import { applyContextPruning } from './contextPruning';

export type PruneMessagesFactoryParams = {
  provider?: Providers;
  maxTokens: number;
  startIndex: number;
  tokenCounter: TokenCounter;
  indexTokenCountMap: Record<string, number | undefined>;
  thinkingEnabled?: boolean;
  /**
   * Optional position-based content degradation config. When `enabled: true`,
   * tool result messages outside the protected zone (system + last N assistant
   * turns + image content) are soft-trimmed (head + tail) once their
   * position-age exceeds `softTrimRatio`, then hard-cleared with a placeholder
   * once age exceeds `hardClearRatio`. Defaults disabled to preserve existing
   * behavior — opt in per deployment.
   */
  contextPruningConfig?: ContextPruningConfig;
};
export type PruneMessagesParams = {
  messages: BaseMessage[];
  usageMetadata?: Partial<UsageMetadata>;
  startType?: ReturnType<BaseMessage['getType']>;
};

function isIndexInContext(
  arrayA: unknown[],
  arrayB: unknown[],
  targetIndex: number
): boolean {
  const startingIndexInA = arrayA.length - arrayB.length;
  return targetIndex >= startingIndexInA;
}

function addThinkingBlock(
  message: AIMessage,
  thinkingBlock: ThinkingContentText | ReasoningContentText
): AIMessage {
  const content: MessageContentComplex[] = Array.isArray(message.content)
    ? (message.content as MessageContentComplex[])
    : [
        {
          type: ContentTypes.TEXT,
          text: message.content,
        },
      ];
  /** Edge case, the message already has the thinking block */
  if (content[0].type === thinkingBlock.type) {
    return message;
  }
  content.unshift(thinkingBlock);
  return new AIMessage({
    ...message,
    content,
  });
}

/**
 * Calculates the total tokens from a single usage object
 *
 * @param usage The usage metadata object containing token information
 * @returns An object containing the total input and output tokens
 */
export function calculateTotalTokens(
  usage: Partial<UsageMetadata>
): UsageMetadata {
  const baseInputTokens = Number(usage.input_tokens) || 0;
  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;
  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;

  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;
  const totalOutputTokens = Number(usage.output_tokens) || 0;

  return {
    input_tokens: totalInputTokens,
    output_tokens: totalOutputTokens,
    total_tokens: totalInputTokens + totalOutputTokens,
  };
}

export type PruningResult = {
  context: BaseMessage[];
  remainingContextTokens: number;
  messagesToRefine: BaseMessage[];
  thinkingStartIndex?: number;
};

/**
 * Processes an array of messages and returns a context of messages that fit within a specified token limit.
 * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.
 *
 * @param options Configuration options for processing messages
 * @returns Object containing the message context, remaining tokens, messages not included, and summary index
 */
/**
 * Builds an emergency context from first and last message pairs when normal pruning
 * removes all AI messages. Preserves conversation continuity by keeping boundaries.
 */
function buildEmergencyContext(
  messages: BaseMessage[],
  endIndex: number
): BaseMessage[] {
  const result: BaseMessage[] = [];
  const instructions = endIndex > 0 ? messages[0] : undefined;
  if (instructions) {
    result.push(instructions);
  }

  /** Find first user+AI pair after system message */
  for (let i = endIndex; i < messages.length - 1; i++) {
    const current = messages[i];
    const next = messages[i + 1];
    if (current.getType() === 'human' && next.getType() === 'ai') {
      result.push(current, next);
      break;
    }
  }

  /** Find last user+AI pair */
  for (let i = messages.length - 1; i > endIndex; i--) {
    const current = messages[i];
    const prev = messages[i - 1];
    if (current.getType() === 'ai' && prev.getType() === 'human') {
      if (!result.includes(prev)) {
        result.push(prev, current);
      }
      break;
    }
  }

  return result;
}

export function getMessagesWithinTokenLimit({
  messages: _messages,
  maxContextTokens,
  indexTokenCountMap,
  startType: _startType,
  thinkingEnabled,
  tokenCounter,
  thinkingStartIndex: _thinkingStartIndex = -1,
  reasoningType = ContentTypes.THINKING,
  instructionTokens: _instructionTokens = 0,
}: {
  messages: BaseMessage[];
  maxContextTokens: number;
  indexTokenCountMap: Record<string, number | undefined>;
  startType?: string | string[];
  thinkingEnabled?: boolean;
  tokenCounter: TokenCounter;
  thinkingStartIndex?: number;
  reasoningType?: ContentTypes.THINKING | ContentTypes.REASONING_CONTENT;
  /**
   * Token overhead for instructions (system message + tool schemas + summary)
   * that are NOT included in `messages`. When messages[0] is already a
   * SystemMessage the budget is deducted from its indexTokenCountMap entry
   * as before; otherwise this value is subtracted from the available budget.
   */
  instructionTokens?: number;
}): PruningResult {
  // Every reply is primed with <|start|>assistant<|message|>, so we
  // start with 3 tokens for the label after all messages have been counted.
  let currentTokenCount = 3;
  const instructions =
    _messages[0]?.getType() === 'system' ? _messages[0] : undefined;
  const instructionsTokenCount =
    instructions != null ? (indexTokenCountMap[0] ?? 0) : _instructionTokens;
  const initialContextTokens = maxContextTokens - instructionsTokenCount;
  let remainingContextTokens = initialContextTokens;
  let startType = _startType;
  const originalLength = _messages.length;
  const messages = [..._messages];
  /**
   * IMPORTANT: this context array gets reversed at the end, since the latest messages get pushed first.
   *
   * This may be confusing to read, but it is done to ensure the context is in the correct order for the model.
   * */
  let context: Array<BaseMessage | undefined> = [];

  let thinkingStartIndex = _thinkingStartIndex;
  let thinkingEndIndex = -1;
  let thinkingBlock: ThinkingContentText | ReasoningContentText | undefined;
  const endIndex = instructions != null ? 1 : 0;
  const prunedMemory: BaseMessage[] = [];

  if (_thinkingStartIndex > -1) {
    const thinkingMessageContent = messages[_thinkingStartIndex]?.content;
    if (Array.isArray(thinkingMessageContent)) {
      thinkingBlock = thinkingMessageContent.find(
        (content) => content.type === reasoningType
      ) as ThinkingContentText | undefined;
    }
  }

  if (currentTokenCount < remainingContextTokens) {
    let currentIndex = messages.length;
    while (
      messages.length > 0 &&
      currentTokenCount < remainingContextTokens &&
      currentIndex > endIndex
    ) {
      currentIndex--;
      if (messages.length === 1 && instructions) {
        break;
      }
      const poppedMessage = messages.pop();
      if (!poppedMessage) continue;
      const messageType = poppedMessage.getType();
      if (
        thinkingEnabled === true &&
        thinkingEndIndex === -1 &&
        currentIndex === originalLength - 1 &&
        (messageType === 'ai' || messageType === 'tool')
      ) {
        thinkingEndIndex = currentIndex;
      }
      if (
        thinkingEndIndex > -1 &&
        !thinkingBlock &&
        thinkingStartIndex < 0 &&
        messageType === 'ai' &&
        Array.isArray(poppedMessage.content)
      ) {
        thinkingBlock = poppedMessage.content.find(
          (content) => content.type === reasoningType
        ) as ThinkingContentText | undefined;
        thinkingStartIndex = thinkingBlock != null ? currentIndex : -1;
      }
      /** False start, the latest message was not part of a multi-assistant/tool sequence of messages */
      if (
        thinkingEndIndex > -1 &&
        currentIndex === thinkingEndIndex - 1 &&
        messageType !== 'ai' &&
        messageType !== 'tool'
      ) {
        thinkingEndIndex = -1;
      }

      const tokenCount = indexTokenCountMap[currentIndex] ?? 0;

      if (
        prunedMemory.length === 0 &&
        currentTokenCount + tokenCount <= remainingContextTokens
      ) {
        context.push(poppedMessage);
        currentTokenCount += tokenCount;
      } else {
        prunedMemory.push(poppedMessage);
        if (thinkingEndIndex > -1 && thinkingStartIndex < 0) {
          continue;
        }
        break;
      }
    }

    if (context[context.length - 1]?.getType() === 'tool') {
      startType = ['ai', 'human'];
    }

    if (startType != null && startType.length > 0 && context.length > 0) {
      let requiredTypeIndex = -1;

      let totalTokens = 0;
      for (let i = context.length - 1; i >= 0; i--) {
        const currentType = context[i]?.getType() ?? '';
        if (
          Array.isArray(startType)
            ? startType.includes(currentType)
            : currentType === startType
        ) {
          requiredTypeIndex = i + 1;
          break;
        }
        const originalIndex = originalLength - 1 - i;
        totalTokens += indexTokenCountMap[originalIndex] ?? 0;
      }

      if (requiredTypeIndex > 0) {
        currentTokenCount -= totalTokens;
        context = context.slice(0, requiredTypeIndex);
      }
    }
  }

  if (instructions && originalLength > 0) {
    context.push(_messages[0] as BaseMessage);
    messages.shift();
  }

  remainingContextTokens -= currentTokenCount;
  const result: PruningResult = {
    remainingContextTokens,
    context: [] as BaseMessage[],
    messagesToRefine: prunedMemory,
  };

  if (thinkingStartIndex > -1) {
    result.thinkingStartIndex = thinkingStartIndex;
  }

  if (
    prunedMemory.length === 0 ||
    thinkingEndIndex < 0 ||
    (thinkingStartIndex > -1 &&
      isIndexInContext(_messages, context, thinkingStartIndex))
  ) {
    // we reverse at this step to ensure the context is in the correct order for the model, and we need to work backwards
    result.context = context.reverse() as BaseMessage[];
    return result;
  }

  if (thinkingEndIndex > -1 && thinkingStartIndex < 0) {
    throw new Error(
      'The payload is malformed. There is a thinking sequence but no "AI" messages with thinking blocks.'
    );
  }

  if (!thinkingBlock) {
    throw new Error(
      'The payload is malformed. There is a thinking sequence but no thinking block found.'
    );
  }

  // Since we have a thinking sequence, we need to find the last assistant message
  // in the latest AI/tool sequence to add the thinking block that falls outside of the current context
  // Latest messages are ordered first.
  let assistantIndex = -1;
  for (let i = 0; i < context.length; i++) {
    const currentMessage = context[i];
    const type = currentMessage?.getType();
    if (type === 'ai') {
      assistantIndex = i;
    }
    if (assistantIndex > -1 && (type === 'human' || type === 'system')) {
      break;
    }
  }

  if (assistantIndex === -1) {
    /** Emergency fallback: keep first and last message pairs instead of crashing */
    const emergencyContext = buildEmergencyContext(_messages, endIndex);
    return {
      remainingContextTokens: 0,
      context: emergencyContext,
      messagesToRefine: _messages.filter((m) => !emergencyContext.includes(m)),
    };
  }

  thinkingStartIndex = originalLength - 1 - assistantIndex;
  const thinkingTokenCount = tokenCounter(
    new AIMessage({ content: [thinkingBlock] })
  );
  const newRemainingCount = remainingContextTokens - thinkingTokenCount;
  const newMessage = addThinkingBlock(
    context[assistantIndex] as AIMessage,
    thinkingBlock
  );
  context[assistantIndex] = newMessage;
  if (newRemainingCount > 0) {
    result.context = context.reverse() as BaseMessage[];
    return result;
  }

  const thinkingMessage: AIMessage = context[assistantIndex] as AIMessage;
  // now we need to an additional round of pruning but making the thinking block fit
  const newThinkingMessageTokenCount =
    (indexTokenCountMap[thinkingStartIndex] ?? 0) + thinkingTokenCount;
  remainingContextTokens = initialContextTokens - newThinkingMessageTokenCount;
  currentTokenCount = 3;
  let newContext: BaseMessage[] = [];
  const secondRoundMessages = [..._messages];
  let currentIndex = secondRoundMessages.length;
  while (
    secondRoundMessages.length > 0 &&
    currentTokenCount < remainingContextTokens &&
    currentIndex > thinkingStartIndex
  ) {
    currentIndex--;
    const poppedMessage = secondRoundMessages.pop();
    if (!poppedMessage) continue;
    const tokenCount = indexTokenCountMap[currentIndex] ?? 0;
    if (currentTokenCount + tokenCount <= remainingContextTokens) {
      newContext.push(poppedMessage);
      currentTokenCount += tokenCount;
    } else {
      messages.push(poppedMessage);
      break;
    }
  }

  const firstMessage: AIMessage = newContext[newContext.length - 1];
  const firstMessageType = newContext[newContext.length - 1].getType();
  if (firstMessageType === 'tool') {
    startType = ['ai', 'human'];
  }

  if (startType != null && startType.length > 0 && newContext.length > 0) {
    let requiredTypeIndex = -1;

    let totalTokens = 0;
    for (let i = newContext.length - 1; i >= 0; i--) {
      const currentType = newContext[i]?.getType() ?? '';
      if (
        Array.isArray(startType)
          ? startType.includes(currentType)
          : currentType === startType
      ) {
        requiredTypeIndex = i + 1;
        break;
      }
      const originalIndex = originalLength - 1 - i;
      totalTokens += indexTokenCountMap[originalIndex] ?? 0;
    }

    if (requiredTypeIndex > 0) {
      currentTokenCount -= totalTokens;
      newContext = newContext.slice(0, requiredTypeIndex);
    }
  }

  if (firstMessageType === 'ai') {
    const newMessage = addThinkingBlock(firstMessage, thinkingBlock);
    newContext[newContext.length - 1] = newMessage;
  } else {
    newContext.push(thinkingMessage);
  }

  if (instructions && originalLength > 0) {
    newContext.push(_messages[0] as BaseMessage);
    secondRoundMessages.shift();
  }

  result.context = newContext.reverse();
  return result;
}

export function checkValidNumber(value: unknown): value is number {
  return typeof value === 'number' && !isNaN(value) && value > 0;
}

/**
 * Returns the set of tool_call ids referenced by an AIMessage — both via the
 * canonical `tool_calls` array and via inline `tool_use` / `tool_call` content
 * blocks (Anthropic-style content arrays).
 */
function getToolCallIds(message: BaseMessage): Set<string> {
  if (message.getType() !== 'ai') {
    return new Set<string>();
  }

  const ids = new Set<string>();
  const aiMessage = message as AIMessage;
  for (const toolCall of aiMessage.tool_calls ?? []) {
    if (typeof toolCall.id === 'string' && toolCall.id.length > 0) {
      ids.add(toolCall.id);
    }
  }

  if (Array.isArray(aiMessage.content)) {
    for (const part of aiMessage.content) {
      if (typeof part !== 'object') {
        continue;
      }
      const record = part as { type?: unknown; id?: unknown };
      if (
        (record.type === 'tool_use' || record.type === 'tool_call') &&
        typeof record.id === 'string' &&
        record.id.length > 0
      ) {
        ids.add(record.id);
      }
    }
  }

  return ids;
}

function getToolResultId(message: BaseMessage): string | null {
  if (message.getType() !== 'tool') {
    return null;
  }
  const toolMessage = message as ToolMessage & {
    tool_call_id?: unknown;
    toolCallId?: unknown;
  };
  if (
    typeof toolMessage.tool_call_id === 'string' &&
    toolMessage.tool_call_id.length > 0
  ) {
    return toolMessage.tool_call_id;
  }
  if (
    typeof toolMessage.toolCallId === 'string' &&
    toolMessage.toolCallId.length > 0
  ) {
    return toolMessage.toolCallId;
  }
  return null;
}

function resolveTokenCountForMessage({
  message,
  messageIndexMap,
  tokenCounter,
  indexTokenCountMap,
}: {
  message: BaseMessage;
  messageIndexMap: Map<BaseMessage, number>;
  tokenCounter: TokenCounter;
  indexTokenCountMap: Record<string, number | undefined>;
}): number {
  const originalIndex = messageIndexMap.get(message) ?? -1;
  if (originalIndex > -1 && indexTokenCountMap[originalIndex] != null) {
    return indexTokenCountMap[originalIndex] as number;
  }
  return tokenCounter(message);
}

/**
 * Repairs orphan tool messages in a pruned context. Drops `ToolMessage`s whose
 * matching `AIMessage` was pruned, and strips orphan `tool_use` content blocks
 * (and `tool_calls` entries) from `AIMessage`s whose corresponding tool results
 * are missing. Without this, Anthropic / Bedrock / OpenAI reject the request
 * with a structural validation 400.
 *
 * Intended to run on the post-prune `context` array (not the pre-prune
 * `allMessages`). The dropped messages are returned so the caller can append
 * them to `messagesToRefine` for summarization fidelity.
 */
export function repairOrphanedToolMessages({
  context,
  allMessages,
  tokenCounter,
  indexTokenCountMap,
}: {
  context: BaseMessage[];
  allMessages: BaseMessage[];
  tokenCounter: TokenCounter;
  indexTokenCountMap: Record<string, number | undefined>;
}): {
  context: BaseMessage[];
  reclaimedTokens: number;
  droppedOrphanCount: number;
  /** Messages removed from context during orphan repair.  These should be
   *  appended to `messagesToRefine` so that summarization can still see them
   *  (e.g. a ToolMessage whose parent AI was pruned). */
  droppedMessages: BaseMessage[];
} {
  const messageIndexMap = new Map<BaseMessage, number>();
  for (let i = 0; i < allMessages.length; i++) {
    messageIndexMap.set(allMessages[i], i);
  }

  const validToolCallIds = new Set<string>();
  const presentToolResultIds = new Set<string>();
  for (const message of context) {
    for (const id of getToolCallIds(message)) {
      validToolCallIds.add(id);
    }
    const resultId = getToolResultId(message);
    if (resultId != null) {
      presentToolResultIds.add(resultId);
    }
  }

  let reclaimedTokens = 0;
  let droppedOrphanCount = 0;
  const repairedContext: BaseMessage[] = [];
  const droppedMessages: BaseMessage[] = [];

  for (const message of context) {
    if (message.getType() === 'tool') {
      const toolResultId = getToolResultId(message);
      if (toolResultId == null || !validToolCallIds.has(toolResultId)) {
        droppedOrphanCount += 1;
        reclaimedTokens += resolveTokenCountForMessage({
          message,
          tokenCounter,
          messageIndexMap,
          indexTokenCountMap,
        });
        droppedMessages.push(message);
        continue;
      }
      repairedContext.push(message);
      continue;
    }

    if (message.getType() === 'ai' && message instanceof AIMessage) {
      const toolCallIds = getToolCallIds(message);
      if (toolCallIds.size > 0) {
        let hasOrphanToolCalls = false;
        for (const id of toolCallIds) {
          if (!presentToolResultIds.has(id)) {
            hasOrphanToolCalls = true;
            break;
          }
        }
        if (hasOrphanToolCalls) {
          const originalTokens = resolveTokenCountForMessage({
            message,
            messageIndexMap,
            tokenCounter,
            indexTokenCountMap,
          });
          const stripped = stripOrphanToolUseBlocks(
            message,
            presentToolResultIds
          );
          if (stripped != null) {
            const strippedTokens = tokenCounter(stripped);
            reclaimedTokens += originalTokens - strippedTokens;
            repairedContext.push(stripped);
          } else {
            droppedOrphanCount += 1;
            reclaimedTokens += originalTokens;
            droppedMessages.push(message);
          }
          continue;
        }
      }
    }

    repairedContext.push(message);
  }

  return {
    context: repairedContext,
    reclaimedTokens,
    droppedOrphanCount,
    droppedMessages,
  };
}

/**
 * Strips tool_use content blocks and tool_calls entries from an AI message
 * when their corresponding ToolMessages are not in the context.
 * Returns null if the message has no content left after stripping.
 */
function stripOrphanToolUseBlocks(
  message: AIMessage,
  presentToolResultIds: Set<string>
): AIMessage | null {
  const keptToolCalls = (message.tool_calls ?? []).filter(
    (tc) => typeof tc.id === 'string' && presentToolResultIds.has(tc.id)
  );

  let keptContent: MessageContentComplex[] | string;
  if (Array.isArray(message.content)) {
    const filtered = (message.content as MessageContentComplex[]).filter(
      (block) => {
        if (typeof block !== 'object') {
          return true;
        }
        const record = block as { type?: unknown; id?: unknown };
        if (
          (record.type === 'tool_use' || record.type === 'tool_call') &&
          typeof record.id === 'string'
        ) {
          return presentToolResultIds.has(record.id);
        }
        return true;
      }
    );

    if (filtered.length === 0) {
      return null;
    }
    keptContent = filtered;
  } else {
    keptContent = message.content;
  }

  return new AIMessage({
    ...message,
    content: keptContent,
    tool_calls: keptToolCalls.length > 0 ? keptToolCalls : undefined,
  });
}

/**
 * Lightweight structural cleanup: strips orphan tool_use blocks from AI messages
 * and drops orphan ToolMessages whose AI counterpart is missing.
 *
 * Unlike `repairOrphanedToolMessages`, this does NOT track tokens — it is
 * intended as a final safety net right before model invocation to prevent
 * Anthropic / Bedrock structural validation errors.
 *
 * Uses duck-typing instead of `getType()` because messages at this stage
 * may be plain objects (from LangGraph state serialization) rather than
 * proper BaseMessage class instances.
 *
 * Includes a fast-path: if every tool_call has a matching tool_result and
 * vice-versa, the original array is returned immediately with zero allocation.
 */
export function sanitizeOrphanToolBlocks(
  messages: BaseMessage[]
): BaseMessage[] {
  const allToolCallIds = new Set<string>();
  const allToolResultIds = new Set<string>();

  for (const msg of messages) {
    const msgAny = msg as unknown as Record<string, unknown>;
    const toolCalls = msgAny.tool_calls as Array<{ id?: string }> | undefined;
    if (Array.isArray(toolCalls)) {
      for (const tc of toolCalls) {
        if (
          typeof tc.id === 'string' &&
          tc.id.length > 0 &&
          !tc.id.startsWith(Constants.ANTHROPIC_SERVER_TOOL_PREFIX)
        ) {
          allToolCallIds.add(tc.id);
        }
      }
    }
    if (Array.isArray(msgAny.content)) {
      for (const block of msgAny.content as Array<Record<string, unknown>>) {
        if (
          typeof block === 'object' &&
          (block.type === 'tool_use' || block.type === 'tool_call') &&
          typeof block.id === 'string' &&
          !block.id.startsWith(Constants.ANTHROPIC_SERVER_TOOL_PREFIX)
        ) {
          allToolCallIds.add(block.id);
        }
      }
    }
    const toolCallId = msgAny.tool_call_id as string | undefined;
    if (typeof toolCallId === 'string' && toolCallId.length > 0) {
      allToolResultIds.add(toolCallId);
    }
  }

  let hasOrphans = false;
  for (const id of allToolCallIds) {
    if (!allToolResultIds.has(id)) {
      hasOrphans = true;
      break;
    }
  }
  if (!hasOrphans) {
    for (const id of allToolResultIds) {
      if (!allToolCallIds.has(id)) {
        hasOrphans = true;
        break;
      }
    }
  }
  if (!hasOrphans) {
    return messages;
  }

  const result: BaseMessage[] = [];
  const strippedAiIndices = new Set<number>();

  for (const msg of messages) {
    const msgAny = msg as unknown as Record<string, unknown>;
    const msgType =
      typeof (msg as { getType?: unknown }).getType === 'function'
        ? msg.getType()
        : ((msgAny.role as string | undefined) ??
          (msgAny._type as string | undefined));

    const toolCallId = msgAny.tool_call_id as string | undefined;
    if (
      (msgType === 'tool' || msg instanceof ToolMessage) &&
      typeof toolCallId === 'string' &&
      !allToolCallIds.has(toolCallId)
    ) {
      continue;
    }

    const toolCalls = msgAny.tool_calls as Array<{ id?: string }> | undefined;
    if (
      (msgType === 'ai' ||
        msgType === 'assistant' ||
        msg instanceof AIMessage) &&
      Array.isArray(toolCalls) &&
      toolCalls.length > 0
    ) {
      const hasOrphanCalls = toolCalls.some(
        (tc) => typeof tc.id === 'string' && !allToolResultIds.has(tc.id)
      );
      if (hasOrphanCalls) {
        if (msg instanceof AIMessage) {
          const stripped = stripOrphanToolUseBlocks(msg, allToolResultIds);
          if (stripped != null) {
            strippedAiIndices.add(result.length);
            result.push(stripped);
          }
          continue;
        }
        const keptToolCalls = toolCalls.filter(
          (tc) => typeof tc.id === 'string' && allToolResultIds.has(tc.id)
        );
        const keptContent = Array.isArray(msgAny.content)
          ? (msgAny.content as Array<Record<string, unknown>>).filter(
              (block) => {
                if (typeof block !== 'object') return true;
                if (
                  (block.type === 'tool_use' || block.type === 'tool_call') &&
                  typeof block.id === 'string'
                ) {
                  return allToolResultIds.has(block.id);
                }
                return true;
              }
            )
          : msgAny.content;
        if (
          keptToolCalls.length === 0 &&
          Array.isArray(keptContent) &&
          keptContent.length === 0
        ) {
          continue;
        }
        strippedAiIndices.add(result.length);
        const patched = Object.create(
          Object.getPrototypeOf(msg),
          Object.getOwnPropertyDescriptors(msg)
        );
        patched.tool_calls = keptToolCalls.length > 0 ? keptToolCalls : [];
        patched.content = keptContent;
        result.push(patched as BaseMessage);
        continue;
      }
    }

    result.push(msg);
  }

  // Bedrock/Anthropic require the conversation to end with a user message;
  // a stripped AI message (tool_use removed) represents a dead-end exchange.
  while (result.length > 0 && strippedAiIndices.has(result.length - 1)) {
    result.pop();
  }

  return result;
}

/**
 * Pre-flight truncation: truncates oversized tool result string content before
 * the budget walker runs. Older results get smaller budgets (recency factor),
 * with a 200-char floor so even ancient entries keep diagnostic signal. Returns
 * the number of ToolMessages truncated.
 */
export function preFlightTruncateToolResults(params: {
  messages: BaseMessage[];
  maxContextTokens: number;
  indexTokenCountMap: Record<string, number | undefined>;
  tokenCounter: TokenCounter;
}): number {
  const { messages, maxContextTokens, indexTokenCountMap, tokenCounter } =
    params;
  const baseMaxChars = calculateMaxToolResultChars(maxContextTokens);
  let truncatedCount = 0;

  const toolIndices: number[] = [];
  for (let i = 0; i < messages.length; i++) {
    if (messages[i].getType() === 'tool') {
      toolIndices.push(i);
    }
  }

  for (let t = 0; t < toolIndices.length; t++) {
    const i = toolIndices[t];
    const message = messages[i];
    const content = message.content;
    if (typeof content !== 'string') {
      continue;
    }

    const position = toolIndices.length > 1 ? t / (toolIndices.length - 1) : 1;
    const recencyFactor = 0.2 + 0.8 * position;
    const maxChars = Math.max(200, Math.floor(baseMaxChars * recencyFactor));

    if (content.length <= maxChars) {
      continue;
    }

    const truncated = truncateToolResultContent(content, maxChars);
    const cloned = new ToolMessage({
      content: truncated,
      tool_call_id: (message as ToolMessage).tool_call_id,
      name: message.name,
      id: message.id,
      additional_kwargs: message.additional_kwargs,
      response_metadata: message.response_metadata,
    });
    messages[i] = cloned;
    indexTokenCountMap[i] = tokenCounter(cloned);
    truncatedCount++;
  }

  return truncatedCount;
}

/**
 * Pre-flight truncation: truncates oversized `tool_use` input fields in AI messages.
 *
 * Tool call inputs (arguments) can be very large — e.g., code evaluation payloads from
 * MCP tools like chrome-devtools. Since these tool calls have already been executed,
 * the model only needs a summary of what was called, not the full arguments. Truncating
 * them before pruning can prevent entire messages from being dropped.
 *
 * Uses 15% of the context window (in estimated characters, ~4 chars/token) as the
 * per-input cap, capped at 200K chars.
 *
 * @returns The number of AI messages that had tool_use inputs truncated.
 */
export function preFlightTruncateToolCallInputs(params: {
  messages: BaseMessage[];
  maxContextTokens: number;
  indexTokenCountMap: Record<string, number | undefined>;
  tokenCounter: TokenCounter;
}): number {
  const { messages, maxContextTokens, indexTokenCountMap, tokenCounter } =
    params;
  const maxInputChars = Math.min(
    Math.floor(maxContextTokens * 0.15) * 4,
    200_000
  );
  let truncatedCount = 0;

  for (let i = 0; i < messages.length; i++) {
    const message = messages[i];
    if (message.getType() !== 'ai') {
      continue;
    }
    if (!Array.isArray(message.content)) {
      continue;
    }

    const originalContent = message.content as MessageContentComplex[];
    const state = { changed: false };
    const newContent = originalContent.map((block) => {
      if (typeof block !== 'object') {
        return block;
      }
      const record = block as Record<string, unknown>;
      if (record.type !== 'tool_use' && record.type !== 'tool_call') {
        return block;
      }

      const input = record.input;
      if (input == null) {
        return block;
      }
      const serialized =
        typeof input === 'string' ? input : JSON.stringify(input);
      if (serialized.length <= maxInputChars) {
        return block;
      }

      state.changed = true;
      // Replaces original input with { _truncated, _originalChars } —
      // safe because the tool call already executed in a prior turn.
      return {
        ...record,
        input: truncateToolInput(serialized, maxInputChars),
      };
    });

    if (!state.changed) {
      continue;
    }

    const aiMsg = message as AIMessage;
    const newToolCalls = (aiMsg.tool_calls ?? []).map((tc) => {
      const serializedArgs = JSON.stringify(tc.args);
      if (serializedArgs.length <= maxInputChars) {
        return tc;
      }
      // Replaces original args with { _truncated, _originalChars } —
      // safe because the tool call already executed in a prior turn.
      return {
        ...tc,
        args: truncateToolInput(serializedArgs, maxInputChars),
      };
    });

    messages[i] = new AIMessage({
      ...aiMsg,
      content: newContent,
      tool_calls: newToolCalls.length > 0 ? newToolCalls : undefined,
    });
    indexTokenCountMap[i] = tokenCounter(messages[i]);
    truncatedCount++;
  }

  return truncatedCount;
}

type ThinkingBlocks = {
  thinking_blocks?: Array<{
    type: 'thinking';
    thinking: string;
    signature: string;
  }>;
};

export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };
  let lastTurnStartIndex = factoryParams.startIndex;
  let lastCutOffIndex = 0;
  let totalTokens = Object.values(indexTokenCountMap).reduce(
    (a = 0, b = 0) => a + b,
    0
  ) as number;
  let runThinkingStartIndex = -1;
  /**
   * Original (pre-truncation) tool result content keyed by message index.
   * Lets the summarizer see full tool outputs even after pre-flight truncation
   * has shortened them in the live messages array. Cleared whenever the
   * factory is recreated (e.g. after a summarization checkpoint).
   */
  const originalToolContent = new Map<number, string>();
  return function pruneMessages(params: PruneMessagesParams): {
    context: BaseMessage[];
    indexTokenCountMap: Record<string, number | undefined>;
    messagesToRefine: BaseMessage[];
    originalToolContent?: Map<number, string>;
  } {
    if (
      factoryParams.provider === Providers.OPENAI &&
      factoryParams.thinkingEnabled === true
    ) {
      for (let i = lastTurnStartIndex; i < params.messages.length; i++) {
        const m = params.messages[i];
        if (
          m.getType() === 'ai' &&
          typeof m.additional_kwargs.reasoning_content === 'string' &&
          Array.isArray(
            (
              m.additional_kwargs.provider_specific_fields as
                | ThinkingBlocks
                | undefined
            )?.thinking_blocks
          ) &&
          (m as AIMessage).tool_calls &&
          ((m as AIMessage).tool_calls?.length ?? 0) > 0
        ) {
          const message = m as AIMessage;
          const thinkingBlocks = (
            message.additional_kwargs.provider_specific_fields as ThinkingBlocks
          ).thinking_blocks;
          const signature =
            thinkingBlocks?.[thinkingBlocks.length - 1].signature;
          const thinkingBlock: ThinkingContentText = {
            signature,
            type: ContentTypes.THINKING,
            thinking: message.additional_kwargs.reasoning_content as string,
          };

          params.messages[i] = new AIMessage({
            ...message,
            content: [thinkingBlock],
            additional_kwargs: {
              ...message.additional_kwargs,
              reasoning_content: undefined,
            },
          });
        }
      }
    }

    let currentUsage: UsageMetadata | undefined;
    if (
      params.usageMetadata &&
      (checkValidNumber(params.usageMetadata.input_tokens) ||
        (checkValidNumber(params.usageMetadata.input_token_details) &&
          (checkValidNumber(
            params.usageMetadata.input_token_details.cache_creation
          ) ||
            checkValidNumber(
              params.usageMetadata.input_token_details.cache_read
            )))) &&
      checkValidNumber(params.usageMetadata.output_tokens)
    ) {
      currentUsage = calculateTotalTokens(params.usageMetadata);
      totalTokens = currentUsage.total_tokens;
    }

    const newOutputs = new Set<number>();
    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {
      const message = params.messages[i];
      if (
        i === lastTurnStartIndex &&
        indexTokenCountMap[i] === undefined &&
        currentUsage
      ) {
        indexTokenCountMap[i] = currentUsage.output_tokens;
      } else if (indexTokenCountMap[i] === undefined) {
        indexTokenCountMap[i] = factoryParams.tokenCounter(message);
        if (currentUsage) {
          newOutputs.add(i);
        }
        totalTokens += indexTokenCountMap[i] ?? 0;
      }
    }

    // If `currentUsage` is defined, we need to distribute the current total tokens to our `indexTokenCountMap`,
    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,
    // relative the manually counted tokens in `indexTokenCountMap`.
    // EDGE CASE: when the resulting context gets pruned, we should not distribute the usage for messages that are not in the context.
    if (currentUsage) {
      let totalIndexTokens = 0;
      if (params.messages[0].getType() === 'system') {
        totalIndexTokens += indexTokenCountMap[0] ?? 0;
      }
      for (let i = lastCutOffIndex; i < params.messages.length; i++) {
        if (i === 0 && params.messages[0].getType() === 'system') {
          continue;
        }
        if (newOutputs.has(i)) {
          continue;
        }
        totalIndexTokens += indexTokenCountMap[i] ?? 0;
      }

      // Calculate ratio based only on messages that remain in the context
      const ratio = currentUsage.total_tokens / totalIndexTokens;
      const isRatioSafe = ratio >= 1 / 3 && ratio <= 2.5;

      // Apply the ratio adjustment only to messages at or after lastCutOffIndex, and only if the ratio is safe
      if (isRatioSafe) {
        if (
          params.messages[0].getType() === 'system' &&
          lastCutOffIndex !== 0
        ) {
          indexTokenCountMap[0] = Math.round(
            (indexTokenCountMap[0] ?? 0) * ratio
          );
        }

        for (let i = lastCutOffIndex; i < params.messages.length; i++) {
          if (newOutputs.has(i)) {
            continue;
          }
          indexTokenCountMap[i] = Math.round(
            (indexTokenCountMap[i] ?? 0) * ratio
          );
        }
      }
    }

    lastTurnStartIndex = params.messages.length;
    if (lastCutOffIndex === 0 && totalTokens <= factoryParams.maxTokens) {
      return {
        context: params.messages,
        indexTokenCountMap,
        messagesToRefine: [],
      };
    }

    /*
     * Pre-flight truncation: chop oversized tool result content and oversized
     * tool_use input args BEFORE the budget walker runs. Older results get
     * smaller per-message budgets (recencyFactor) so they degrade first while
     * recent turns stay intact. Saves the originals in `originalToolContent`
     * keyed by message index so the summarizer can later reconstruct full
     * fidelity even after the live array was shortened.
     */
    for (let i = 0; i < params.messages.length; i++) {
      const m = params.messages[i];
      if (
        m.getType() === 'tool' &&
        typeof m.content === 'string' &&
        !originalToolContent.has(i)
      ) {
        originalToolContent.set(i, m.content);
      }
    }
    preFlightTruncateToolResults({
      messages: params.messages,
      maxContextTokens: factoryParams.maxTokens,
      indexTokenCountMap,
      tokenCounter: factoryParams.tokenCounter,
    });
    preFlightTruncateToolCallInputs({
      messages: params.messages,
      maxContextTokens: factoryParams.maxTokens,
      indexTokenCountMap,
      tokenCounter: factoryParams.tokenCounter,
    });

    /*
     * Position-based content degradation: soft-trim then hard-clear older tool
     * results before the budget walker runs. Off by default — caller opts in
     * via `contextPruningConfig.enabled`.
     */
    if (factoryParams.contextPruningConfig?.enabled) {
      applyContextPruning({
        messages: params.messages,
        indexTokenCountMap,
        tokenCounter: factoryParams.tokenCounter,
        config: factoryParams.contextPruningConfig,
      });
    }

    const { context, thinkingStartIndex, messagesToRefine } =
      getMessagesWithinTokenLimit({
        maxContextTokens: factoryParams.maxTokens,
        messages: params.messages,
        indexTokenCountMap,
        startType: params.startType,
        thinkingEnabled: factoryParams.thinkingEnabled,
        tokenCounter: factoryParams.tokenCounter,
        reasoningType:
          factoryParams.provider === Providers.BEDROCK
            ? ContentTypes.REASONING_CONTENT
            : ContentTypes.THINKING,
        thinkingStartIndex:
          factoryParams.thinkingEnabled === true
            ? runThinkingStartIndex
            : undefined,
      });
    runThinkingStartIndex = thinkingStartIndex ?? -1;

    /*
     * Orphan repair pass — drops tool messages whose AIMessage was pruned and
     * strips orphan tool_use blocks from AIMessages whose tool results are
     * missing. Without this, the budget walker can split tool_use ↔ tool_result
     * pairs and the provider rejects the request with a structural 400.
     * Dropped messages flow into messagesToRefine for summarization.
     */
    const repaired = repairOrphanedToolMessages({
      context,
      allMessages: params.messages,
      tokenCounter: factoryParams.tokenCounter,
      indexTokenCountMap,
    });
    const finalContext = repaired.context;
    const finalMessagesToRefine =
      repaired.droppedMessages.length > 0
        ? [...messagesToRefine, ...repaired.droppedMessages]
        : messagesToRefine;

    /** The index is the first value of `context`, index relative to `params.messages` */
    lastCutOffIndex = Math.max(
      params.messages.length -
        (finalContext.length -
          (finalContext[0]?.getType() === 'system' ? 1 : 0)),
      0
    );

    return {
      context: finalContext,
      indexTokenCountMap,
      messagesToRefine: finalMessagesToRefine,
      originalToolContent,
    };
  };
}

/** Calculates context utilization as a percentage (0-100) */
export function getContextUtilization(
  indexTokenCountMap: Record<string, number | undefined>,
  instructionTokens: number,
  maxContextTokens: number
): number {
  if (maxContextTokens <= 0) return 0;
  let totalTokens = instructionTokens;
  for (const key in indexTokenCountMap) {
    totalTokens += indexTokenCountMap[key] ?? 0;
  }
  return (totalTokens / maxContextTokens) * 100;
}