import { assertEquals } from "@std/assert";
import { type CompactionCap, defineExtension } from "@berry/api";
import { envApiKeyAuth, type Model, type Provider } from "@earendil-works/pi-ai";
import { setup } from "./fixture.ts";

const model = (id: string, contextWindow: number): Model<"openai-completions"> => ({
  id,
  name: id,
  api: "openai-completions",
  provider: "sized",
  baseUrl: "https://example.invalid",
  input: ["text"],
  reasoning: false,
  contextWindow,
  maxTokens: 4096,
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
});

const provider: Provider = {
  id: "sized",
  name: "sized",
  auth: { apiKey: envApiKeyAuth("sized key", ["SIZED_KEY"]) },
  getModels: () => [model("big", 1_048_576), model("small", 20_000)],
  stream: () => {
    throw new Error("unused");
  },
  streamSimple: () => {
    throw new Error("unused");
  },
};

const sized = () =>
  defineExtension({
    name: "sized",
    initial: () => null,
    setup(ext) {
      ext.provider(provider);
    },
  });

const cap = (name: string, value: CompactionCap | undefined) =>
  defineExtension({
    name,
    initial: () => null,
    setup(ext) {
      ext.compactionCap(() => value);
    },
  });

Deno.test("generation sees windows capped at the threshold plus reserve, smaller ones alone", async () => {
  await using harness = await setup();
  const { berry, faux, models } = harness;
  await berry.install(sized());
  await berry.install(cap("cap", { maxContextTokens: 262_144, backgroundTokens: 32_768 }));
  const seen = berry.effectiveModels;
  assertEquals(seen.getModel("sized", "big")?.contextWindow, 262_144 + 16_384);
  assertEquals(seen.getModel("sized", "small")?.contextWindow, 20_000);
  assertEquals(seen.getModel("faux", faux.getModel().id)?.contextWindow, 128_000);
  // the stored catalog keeps its own windows
  assertEquals(models.getModel("sized", "big")?.contextWindow, 1_048_576);
  assertEquals(
    berry.compactionCap(),
    { maxContextTokens: 262_144, backgroundTokens: 32_768 },
  );
});

Deno.test("the tightest cap wins and lapses on uninstall", async () => {
  await using harness = await setup();
  const { berry } = harness;
  const seen = berry.effectiveModels;
  await berry.install(sized());
  await berry.install(cap("loose", { maxContextTokens: 500_000, backgroundTokens: 1_000 }));
  await berry.install(cap("tight", { maxContextTokens: 100_000, backgroundTokens: 2_000 }));
  assertEquals(
    berry.compactionCap(),
    { maxContextTokens: 100_000, backgroundTokens: 2_000 },
  );
  assertEquals(seen.getModel("sized", "big")?.contextWindow, 100_000 + 16_384);
  berry.uninstall("tight");
  assertEquals(
    berry.compactionCap(),
    { maxContextTokens: 500_000, backgroundTokens: 1_000 },
  );
  assertEquals(seen.getModel("sized", "big")?.contextWindow, 500_000 + 16_384);
  berry.uninstall("loose");
  assertEquals(berry.compactionCap(), undefined);
  assertEquals(seen.getModel("sized", "big")?.contextWindow, 1_048_576);
});

Deno.test("invalid and failing caps read as no cap", async () => {
  await using harness = await setup();
  const { berry } = harness;
  await berry.install(sized());
  await berry.install(cap("zero", { maxContextTokens: 0, backgroundTokens: 0 }));
  await berry.install(cap("negative", { maxContextTokens: 100_000, backgroundTokens: -1 }));
  await berry.install(defineExtension({
    name: "failing",
    initial: () => null,
    setup(ext) {
      ext.compactionCap(() => {
        throw new Error("no cap here");
      });
    },
  }));
  assertEquals(berry.compactionCap(), undefined);
  assertEquals(berry.effectiveModels.getModel("sized", "big")?.contextWindow, 1_048_576);
});