import { assert, assertEquals, assertRejects } from "@std/assert";
import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import { fauxAssistantMessage, fauxProvider } from "@earendil-works/pi-ai/providers/faux";
import { createModels } from "@earendil-works/pi-ai/models";
import {
  type ConversationId,
  type Cursor,
  type EntryRecord,
  MemoryStorage,
} from "@earendil-works/pi-durable";
import { type Berry, Conversations, openBerry } from "./berry.ts";
import { script, setup, turn } from "./fixture.ts";
import { conversationTree, type TreeEntry } from "./tree.ts";

async function history(berry: Berry, id: number): Promise<EntryRecord[]> {
  const conversation = (await berry.harness.conversation(id as ConversationId, context))!;
  const entries: EntryRecord[] = [];
  let cursor: Cursor | undefined;
  do {
    const page = await conversation.entries({ order: "ascending" }, 100, cursor, context);
    entries.push(...page.items);
    cursor = page.next;
  } while (cursor !== undefined);
  return entries;
}

/** The paths from each branch's tip back to the root, which must be its history. */
async function assertPaths(berry: Berry, tree: Awaited<ReturnType<typeof conversationTree>>) {
  const byId = new Map(tree.entries.map((entry) => [entry.id, entry]));
  for (const branch of tree.branches) {
    const path: TreeEntry[] = [];
    for (let id = branch.tip; id !== null; id = byId.get(id)!.parent) path.unshift(byId.get(id)!);
    assertEquals(
      path.map((entry) => [entry.id, entry.kind]),
      (await history(berry, branch.id)).map((entry) => [entry.id, entry.kind]),
    );
  }
}

Deno.test("a conversation's branches join into one tree, whose paths are their histories", async () => {
  await using t = await setup();
  const { berry } = t;
  const tid = await berry.create({ extensions: {}, model: t.model });
  const info = async () => (await berry.harness.snapshot(Conversations, context))![tid];
  await turn(t, tid, [script("1")]);
  await turn(t, tid, []);
  const first = (await info()).id;
  const [start, ...rest] = await history(berry, first);
  const firstAnswer = rest.find((entry) => entry.kind === "pi.assistant")!;

  // a new first message, then a branch after the first answer
  await berry.branch(tid, start.id);
  await turn(t, tid, []);
  await berry.branch(tid, firstAnswer.id);
  await turn(t, tid, []);

  const tree = await conversationTree(berry, tid);
  assertEquals(tree.branches.length, 3);
  assertEquals(tree.head, tree.branches[2].id);
  assertEquals(await berry.tidOf(tree.head as ConversationId), tid);
  await assertPaths(berry, tree);
  assertEquals(tree.entries.filter((entry) => entry.parent === null), [
    { id: start.id, parent: null, branch: first, kind: "berry.start", text: "" },
  ]);

  // back to where it began, which needs no new branch
  await berry.branch(tid, rest.at(-1)!.id);
  assertEquals((await info()).head, undefined);
  assertEquals((await conversationTree(berry, tid)).branches.length, 3);
});

Deno.test("a fork's tree holds the history it inherited, but none of its parent's other branches", async () => {
  await using t = await setup();
  const { berry } = t;
  const tid = await berry.create({ extensions: {}, model: t.model });
  await turn(t, tid, []);
  const parent = (await berry.conversation(tid))!;
  const entries = await history(berry, parent.id);
  await berry.branch(tid, entries[0].id);
  await turn(t, tid, []);

  const fork = await parent.fork(entries.at(-1)!.id, { ownership: { kind: "ownerless" } }, context);
  const forkTid = await berry.tidOf(fork.id);
  await turn(t, forkTid, []);
  const tree = await conversationTree(berry, forkTid);
  assertEquals(tree.branches.map((branch) => branch.id), [fork.id]);
  await assertPaths(berry, tree);
});

Deno.test("regenerating an agent's message runs it again from just before it, keeping the first one", async () => {
  await using t = await setup();
  const { berry } = t;
  const tid = await berry.create({ extensions: {}, model: t.model });
  await turn(t, tid, [script("1")]);
  const first = (await berry.conversation(tid))!.id;
  const entries = await history(berry, first);
  const answer = entries.at(-1)!;
  const message = entries.find((entry) => entry.kind === "pi.user")!;

  t.faux.setResponses([fauxAssistantMessage("again")]);
  await berry.regenerate(tid, answer.id);
  await (await berry.conversation(tid))!.waitForIdle(context);

  const tree = await conversationTree(berry, tid);
  await assertPaths(berry, tree);
  const byId = new Map(tree.entries.map((entry) => [entry.id, entry]));
  const path = (branch: number) => {
    const entries: TreeEntry[] = [];
    let id = tree.branches.find((of) => of.id === branch)!.tip;
    for (; id !== null; id = byId.get(id)!.parent) entries.unshift(byId.get(id)!);
    return entries.filter((entry) => entry.kind !== "pi.system");
  };
  const was = path(first);
  const now = path(tree.head);
  assertEquals(
    was.map((entry) => entry.kind),
    ["berry.start", "pi.user", "pi.assistant", "pi.tool-result", "pi.assistant"],
  );
  assertEquals(now.slice(0, -1), was.slice(0, -1));
  assertEquals([was.at(-1)!.text, now.at(-1)!.text], ["done", "again"]);
  await assertRejects(() => berry.regenerate(tid, message.id), Error, "not the agent's message");
});

async function rewound(provider?: Parameters<typeof fauxProvider>[0]) {
  const faux = fauxProvider(provider);
  const models = createModels();
  models.setProvider(faux.provider);
  const sessions: (string | undefined)[] = [];
  const payloads: Promise<unknown>[] = [];
  const original = models.streamSimple.bind(models);
  models.streamSimple = (...args) => {
    const [model, , options] = args;
    sessions.push(options?.sessionId);
    // every provider exposes the request it would send through onPayload
    const next = options?.onPayload?.({ prompt_cache_key: "sentry" }, model);
    if (next !== undefined) payloads.push(Promise.resolve(next));
    return original(...args);
  };
  const berry = await openBerry(new MemoryStorage(), models);
  try {
    const model = { provider: "faux", modelId: faux.getModel().id };
    await berry.enableModels([model]);
    const tid = await berry.create({ extensions: {}, model });
    const harness = { faux, berry };
    await turn(harness, tid, []);
    const first = (await berry.conversation(tid))!;
    const [start] = (await first.entries({ order: "ascending" }, 10, undefined, context)).items;
    await berry.branch(tid, start.id);
    await turn(harness, tid, []);
    return { tid, sessions, payloads: await Promise.all(payloads) };
  } finally {
    await berry.close();
  }
}

Deno.test("a rewind keeps the provider session so the new branch starts on a warm prompt cache", async () => {
  const { tid, sessions } = await rewound();
  assert(sessions.length > 1);
  assertEquals(new Set(sessions), new Set([tid]));
});

Deno.test("codex keeps a session per branch, and its body carries the conversation's cache key", async () => {
  const { tid, sessions, payloads } = await rewound({ api: "openai-codex-responses" });
  assert(new Set(sessions).size > 1);
  assert(payloads.length > 0);
  for (const payload of payloads) {
    assertEquals((payload as { prompt_cache_key: unknown }).prompt_cache_key, tid);
  }
});

Deno.test("tree entries preview their messages on one line", async () => {
  await using t = await setup();
  const tid = await t.berry.create({ extensions: {}, model: t.model });
  await turn(t, tid, [script("throw new Error('no\\nway')")]);
  const { entries } = await conversationTree(t.berry, tid);
  assertEquals(
    entries.filter((entry) => entry.kind !== "pi.system").map(({ kind, text, tools, error }) => ({
      kind,
      text: text.split(" ").slice(0, 2).join(" "),
      tools,
      error,
    })),
    [
      { kind: "berry.start", text: "", tools: undefined, error: undefined },
      { kind: "pi.user", text: "go", tools: undefined, error: undefined },
      { kind: "pi.assistant", text: "", tools: ["codemode"], error: undefined },
      { kind: "pi.tool-result", text: "Script failed", tools: ["codemode"], error: true },
      { kind: "pi.assistant", text: "done", tools: undefined, error: undefined },
    ],
  );
});