import type { ModelMessage } from "ai";

// Keep complete assistant/tool batches together. Never remove just one side of a tool call.
export function splitCompactableHistory(messages: ModelMessage[], keepUserTurns = 3) {
  const userIndices = messages.flatMap((message, index) =>
    message.role === "user" ? [index] : [],
  );
  let cut =
    userIndices.length > keepUserTurns ? userIndices[userIndices.length - keepUserTurns] : 0;
  if (!cut) {
    const openCalls = new Set<string>();
    for (let index = 0; index < messages.length - 6; index++) {
      const message = messages[index];
      if (message.role === "assistant" && Array.isArray(message.content)) {
        for (const part of message.content)
          if (part.type === "tool-call") openCalls.add(part.toolCallId);
      }
      if (message.role === "tool")
        for (const part of message.content) openCalls.delete(part.toolCallId);
      if (!openCalls.size && message.role !== "user") cut = index + 1;
    }
  }
  return { older: messages.slice(0, cut), recent: messages.slice(cut) };
}

export function estimateContextTokens(messages: ModelMessage[]) {
  // Conservatively count CJK characters and media rather than assuming four characters per token.
  const text = JSON.stringify(messages);
  const cjk = (text.match(/[^\x00-\x7f]/g) || []).length;
  return Math.ceil((text.length - cjk) / 3) + cjk;
}

export async function compactAgentContext(params: {
  messages: ModelMessage[];
  previousSummary: string;
  summarize: (text: string) => Promise<string>;
  threshold?: number;
}) {
  if (estimateContextTokens(params.messages) < (params.threshold ?? 18_000)) return params;
  const { older, recent } = splitCompactableHistory(params.messages);
  if (!older.length) throw new Error("当前任务上下文过长，已保留结果，请分成更小的任务继续。");
  const summary = (
    await params.summarize(JSON.stringify({ previousSummary: params.previousSummary, older }))
  ).trim();
  if (!summary || summary.length > 12_000)
    throw new Error("上下文整理未完成，已保留已完成的步骤。");
  return { ...params, previousSummary: summary, messages: recent };
}