// src/messages/dedup.ts
import type { BaseMessage } from '@langchain/core/messages';
import { MessageTypes } from '@/common';
import { DEDUP_MAX_CONTENT_LENGTH } from '@/common/constants';

/**
 * Deduplicates consecutive identical system messages in the context window.
 *
 * Problem: In long tool-use chains, the same system messages (e.g., post-prune notes,
 * conversation summaries) can accumulate when the context is rebuilt on each iteration.
 * These duplicates waste tokens without adding information.
 *
 * Strategy: Only deduplicate system messages that appear consecutively or are exact
 * duplicates of an earlier system message. The FIRST occurrence is always kept.
 * Non-system messages (human, ai, tool) are never touched.
 *
 * Important constraints:
 * - The first system message (index 0) is ALWAYS preserved (it's the main system prompt)
 * - Only system messages are candidates for deduplication
 * - Messages with content longer than DEDUP_MAX_CONTENT_LENGTH are skipped (too expensive to compare)
 * - Content comparison is by string equality (fast and deterministic)
 *
 * @param messages - The message array to deduplicate (not mutated)
 * @returns A new array with duplicate system messages removed, and the count of removed messages
 */
export function deduplicateSystemMessages(messages: BaseMessage[]): {
  messages: BaseMessage[];
  removedCount: number;
} {
  if (messages.length <= 1) {
    return { messages, removedCount: 0 };
  }

  const seenSystemContents = new Set<string>();
  const result: BaseMessage[] = [];
  let removedCount = 0;

  for (let i = 0; i < messages.length; i++) {
    const msg = messages[i];
    const type = msg.getType();

    // Non-system messages are always kept
    if (type !== MessageTypes.SYSTEM) {
      result.push(msg);
      continue;
    }

    // First system message (main prompt) is always kept
    if (i === 0) {
      result.push(msg);
      // Track its content for dedup of later duplicates
      const content = getContentString(msg);
      if (content != null) {
        seenSystemContents.add(content);
      }
      continue;
    }

    // Get string content for comparison
    const content = getContentString(msg);

    // Skip dedup for very long or non-string content
    if (content == null) {
      result.push(msg);
      continue;
    }

    // Check if this exact system message was already seen
    if (seenSystemContents.has(content)) {
      removedCount++;
      continue; // Skip this duplicate
    }

    // New unique system message — keep it and track
    seenSystemContents.add(content);
    result.push(msg);
  }

  return { messages: result, removedCount };
}

/**
 * Extracts a comparable string from a message's content.
 * Returns null if the content is too large or non-string (skip dedup for those).
 */
function getContentString(msg: BaseMessage): string | null {
  if (typeof msg.content === 'string') {
    if (msg.content.length > DEDUP_MAX_CONTENT_LENGTH) {
      return null;
    }
    return msg.content;
  }
  // Array content (e.g., Anthropic cache_control blocks) — serialize for comparison
  // but only if reasonably sized
  try {
    const serialized = JSON.stringify(msg.content);
    if (serialized.length > DEDUP_MAX_CONTENT_LENGTH) {
      return null;
    }
    return serialized;
  } catch {
    return null;
  }
}