import type { ModelMessage } from 'ai';
import { TokenCounter } from './token-counter.js';
import {
  COMPACTED_PLACEHOLDER,
  SUMMARY_MARKER,
  COMPACTED_MARKER,
  type CompressionConfig,
  DEFAULT_COMPRESSION_CONFIG,
} from './types.js';

// 扩展的工具结果类型，支持压缩标记
interface CompactedToolResult {
  type: 'tool-result';
  toolCallId: string;
  result: unknown;
  [COMPACTED_MARKER]?: {
    compactedAt: number;
    originalSize: number;
  };
}

/**
 * 检查消息是否为摘要消息
 */
function isSummaryMessage(message: ModelMessage): boolean {
  if (typeof message.content === 'string') {
    return message.content.includes(SUMMARY_MARKER);
  }
  if (Array.isArray(message.content)) {
    return message.content.some(
      (part) =>
        typeof part === 'object' &&
        part !== null &&
        'text' in part &&
        typeof (part as { text?: unknown }).text === 'string' &&
        ((part as { text: string }).text).includes(SUMMARY_MARKER)
    );
  }
  return false;
}

/**
 * 检查是否为工具结果
 */
function isToolResult(part: unknown): part is CompactedToolResult {
  return (
    typeof part === 'object' &&
    part !== null &&
    (part as { type?: unknown }).type === 'tool-result'
  );
}

/**
 * 检查工具结果是否已压缩
 */
function isCompactedResult(part: unknown): boolean {
  if (!isToolResult(part)) return false;
  return COMPACTED_MARKER in part;
}

/**
 * 获取工具结果的 token 数量
 */
function getToolResultTokens(part: unknown): number {
  if (!isToolResult(part)) return 0;
  return TokenCounter.estimateText(JSON.stringify(part.result));
}

/**
 * 压缩工具结果
 */
function compactToolResult(part: CompactedToolResult): CompactedToolResult {
  const originalSize = TokenCounter.estimateText(JSON.stringify(part.result));
  return {
    ...part,
    result: COMPACTED_PLACEHOLDER,
    [COMPACTED_MARKER]: {
      compactedAt: Date.now(),
      originalSize,
    },
  };
}

/**
 * Prune 策略：压缩旧的工具调用结果
 *
 * 逻辑：
 * 1. 从后往前遍历消息
 * 2. 跳过最近 pruneProtect tokens 的工具结果
 * 3. 遇到 summary 消息停止
 * 4. 将超出保护范围的 tool-result 替换为占位符
 *
 * @param messages 消息数组
 * @param config 压缩配置
 * @returns 处理后的消息数组和释放的 tokens
 */
export function prune(
  messages: ModelMessage[],
  config: CompressionConfig = DEFAULT_COMPRESSION_CONFIG
): { messages: ModelMessage[]; freedTokens: number } {
  const { pruneProtect, pruneMinimum } = config;

  // 深拷贝消息数组
  const result = JSON.parse(JSON.stringify(messages)) as ModelMessage[];

  let protectedTokens = 0;
  let freedTokens = 0;
  const toPrune: Array<{ msgIndex: number; partIndex: number; tokens: number }> = [];

  // 从后往前遍历
  for (let msgIndex = result.length - 1; msgIndex >= 0; msgIndex--) {
    const message = result[msgIndex];

    // 遇到摘要消息停止
    if (isSummaryMessage(message)) {
      break;
    }

    // 只处理包含工具结果的消息
    if (!Array.isArray(message.content)) continue;

    for (let partIndex = message.content.length - 1; partIndex >= 0; partIndex--) {
      const part = message.content[partIndex];

      // 跳过非工具结果
      if (!isToolResult(part)) continue;

      // 跳过已压缩的
      if (isCompactedResult(part)) {
        break; // 遇到已压缩的，说明之前已经 prune 过，停止
      }

      const tokens = getToolResultTokens(part);
      protectedTokens += tokens;

      // 超出保护范围的标记为待压缩
      if (protectedTokens > pruneProtect) {
        toPrune.push({ msgIndex, partIndex, tokens });
        freedTokens += tokens;
      }
    }
  }

  // 如果释放的 tokens 不够最小量，不执行压缩
  if (freedTokens < pruneMinimum) {
    return { messages, freedTokens: 0 };
  }

  // 执行压缩
  for (const { msgIndex, partIndex } of toPrune) {
    const message = result[msgIndex];
    if (Array.isArray(message.content)) {
      const part = message.content[partIndex];
      if (isToolResult(part)) {
        // 使用 any 来绕过严格类型检查，因为我们在运行时知道这是安全的
        (message.content as unknown[])[partIndex] = compactToolResult(part);
      }
    }
  }

  return { messages: result, freedTokens };
}

/**
 * 过滤已压缩的内容（用于发送给模型前）
 * 将已压缩的工具结果替换为占位符文本
 */
export function filterCompacted(messages: ModelMessage[]): ModelMessage[] {
  return messages.map((message) => {
    if (!Array.isArray(message.content)) return message;

    const filteredContent = message.content.map((part) => {
      if (isCompactedResult(part) && isToolResult(part)) {
        return {
          ...part,
          result: COMPACTED_PLACEHOLDER,
        };
      }
      return part;
    });

    return {
      ...message,
      content: filteredContent,
    } as ModelMessage;
  });
}
