import test from "node:test";
import assert from "node:assert/strict";
import { resolvePolicy } from "./config.ts";
import { pruneContextMessages } from "./prune.ts";

const bulky = "line\n".repeat(300);
const boundaryBulky = "boundary\n".repeat(300);
const thresholdWithTrailingNewline = "threshold\n".repeat(150);

function buildPolicy(recentUserTurns = 4) {
  return {
    ...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
    recentUserTurns,
  };
}

test("pruneContextMessages replaces old bulky tool results with distilled summaries instead of deleting them", () => {
  const policy = buildPolicy(2);
  const bulkyFailure = [
    "Build failed while compiling focus parser",
    "Error: missing export createFocusMatcher from ./summary-focus.ts",
    ...Array.from({ length: 220 }, () => "stack frame"),
  ].join("\n");
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "toolResult", toolName: "bash", content: bulkyFailure },
    { role: "assistant", content: "observed turn 1" },
    { role: "user", content: "turn 2" },
    { role: "assistant", content: "observed turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  const distilled = pruned.find((message) => message.role === "toolResult");
  assert.ok(distilled);
  assert.match(distilled!.content, /missing export createFocusMatcher/);
  assert.doesNotMatch(distilled!.content, /stack frame\nstack frame\nstack frame/);
});

test("aggressive mode distills an older bulky tool result sooner than conservative mode", () => {
  const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
  const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "toolResult", toolName: "read", content: bulky },
    { role: "assistant", content: "after turn 1" },
    { role: "user", content: "turn 2" },
    { role: "assistant", content: "after turn 2" },
    { role: "user", content: "turn 3" },
    { role: "assistant", content: "after turn 3" },
    { role: "user", content: "turn 4" },
  ];

  const conservativePruned = pruneContextMessages(messages, conservative);
  const aggressivePruned = pruneContextMessages(messages, aggressive);

  assert.equal(conservativePruned[1]?.content, bulky);
  assert.notEqual(aggressivePruned[1]?.content, bulky);
  assert.match(aggressivePruned[1]?.content ?? "", /^\[distilled read output\]/);
});

test("pruneContextMessages keeps recent bulky tool results inside the recent-turn window", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "assistant", content: "observed turn 1" },
    { role: "user", content: "turn 2" },
    { role: "toolResult", toolName: "read", content: bulky },
    { role: "assistant", content: "observed turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(pruned, messages);
});

test("pruneContextMessages keeps old non-bulky tool results outside the recent-turn window", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "toolResult", toolName: "read", content: "short output" },
    { role: "assistant", content: "observed turn 1" },
    { role: "user", content: "turn 2" },
    { role: "assistant", content: "observed turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(pruned, messages);
});

test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline },
    { role: "assistant", content: "after threshold output" },
    { role: "user", content: "turn 2" },
    { role: "assistant", content: "after turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(pruned, messages);
});

test("pruneContextMessages honors the recent-user-turn boundary", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "toolResult", toolName: "read", content: bulky },
    { role: "assistant", content: "after turn 1" },
    { role: "user", content: "turn 2" },
    { role: "toolResult", toolName: "read", content: boundaryBulky },
    { role: "assistant", content: "after turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.equal(pruned[1]?.role, "toolResult");
  assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/);
  assert.deepEqual(
    pruned.map((message) => message.content),
    ["turn 1", pruned[1]!.content, "after turn 1", "turn 2", boundaryBulky, "after turn 2", "turn 3"]
  );
});