import { assert, assertEquals, assertStringIncludes } from "@std/assert";
import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
import { fauxAssistantMessage, fauxProvider } from "@earendil-works/pi-ai/providers/faux";
import { createModels } from "@earendil-works/pi-ai/models";
import { MemoryStorage } from "@earendil-works/pi-durable";
import { openBerry } from "./berry.ts";
import { setup, turn } from "./fixture.ts";

Deno.test("a connection body read failure retries generation", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const tid = await berry.create({ extensions: {}, model });
  faux.setResponses([
    fauxAssistantMessage("", {
      stopReason: "error",
      errorMessage: "error reading a body from connection",
    }),
    fauxAssistantMessage("recovered"),
  ]);
  const conversation = (await berry.conversation(tid))!;
  const settled = await (await conversation.submit({ type: "input", content: "go" }, context))
    .wait(context);
  await conversation.waitForIdle(context);
  assertEquals(settled.status, "done");
  assertEquals(faux.state.callCount, 2);
  const stored = await conversation.context(context);
  assertEquals(stored.messages.at(-1)?.content, [{ type: "text", text: "recovered" }]);
});

for (const delivery of ["a setup exception", "a result-only stream"] as const) {
  Deno.test(`generation retries EOF delivered through ${delivery}`, async () => {
    const faux = fauxProvider();
    const models = createModels();
    models.setProvider(faux.provider);
    const original = models.streamSimple.bind(models);
    let attempts = 0;
    models.streamSimple = (...args) => {
      if (++attempts === 1) {
        if (delivery === "a setup exception") throw new Error("Unexpected EOF");
        const stream = createAssistantMessageEventStream();
        stream.end(
          fauxAssistantMessage("", { stopReason: "error", errorMessage: "Unexpected EOF" }),
        );
        return stream;
      }
      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 });
      faux.setResponses([fauxAssistantMessage("recovered")]);
      const conversation = (await berry.conversation(tid))!;
      const settled = await (await conversation.submit({ type: "input", content: "go" }, context))
        .wait(context);
      assertEquals(settled.status, "done");
      assertEquals(attempts, 2);
    } finally {
      await berry.close();
    }
  });
}

Deno.test("compaction also retries EOF and places the recovered summary", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const tid = await berry.create({ extensions: {}, model });
  await turn(harness, tid, []);
  const conversation = (await berry.conversation(tid))!;
  faux.setResponses([fauxAssistantMessage("done")]);
  await (await conversation.submit({ type: "input", content: "x".repeat(100_000) }, context))
    .wait(context);
  await conversation.waitForIdle(context);
  faux.setResponses([
    fauxAssistantMessage("", { stopReason: "error", errorMessage: "Unexpected EOF" }),
    fauxAssistantMessage("Recovered summary"),
  ]);
  const task = await conversation.compact(undefined, context);
  const { outcome } = (await berry.harness.waitForTask(task, context)).state;
  assert(outcome.status === "completed");
  await conversation.waitForIdle(context);
  assertEquals(faux.state.callCount, 4);
  const entries =
    (await conversation.entries({ order: "ascending" }, 100, undefined, context)).items;
  const summary = entries.find((entry) => entry.kind === "pi.compaction");
  assert(summary);
  const text = JSON.stringify(summary.model);
  assertStringIncludes(text, "Recovered summary");
});