import { assert, assertEquals, assertRejects } from "@std/assert";
import { createModels, type Provider } from "@earendil-works/pi-ai";
import { fauxProvider } from "@earendil-works/pi-ai/providers/faux";
import { MemoryStorage } from "@earendil-works/pi-durable";
import { defineExtension } from "@berry/api";
import { until } from "@berry/util";
import { openBerry } from "./berry.ts";

Deno.test("providers registered during a failed setup never replace the running load", async () => {
  const models = createModels();
  const berry = await openBerry(new MemoryStorage(), models);
  const original = { ...fauxProvider().provider, id: "custom" };
  const extension = (provider: Provider) =>
    defineExtension({
      name: "provider",
      initial: () => null,
      setup(ext) {
        ext.provider(provider);
      },
    });
  try {
    await berry.install(extension(original));
    await assertRejects(() =>
      berry.install(defineExtension({
        ...extension({ ...original }),
        setup(ext) {
          ext.provider({ ...original });
          ext.provider({ ...original, id: "new" });
          throw new Error("setup failed");
        },
      }))
    );
    assertEquals(models.getProvider("custom"), original);
    assertEquals(models.getProvider("new"), undefined);
  } finally {
    await berry.close();
  }
});

Deno.test("builtin overrides survive reload and restore originals on release", async () => {
  const models = createModels();
  const builtin = fauxProvider().provider;
  const second = { ...builtin, id: "second" };
  models.setProvider(builtin);
  models.setProvider(second);
  const berry = await openBerry(new MemoryStorage(), models);
  const override = { ...builtin, name: "override" };
  const replacement = { ...builtin, name: "replacement" };
  const extension = (providers: Provider[]) =>
    defineExtension({
      name: "provider",
      initial: () => null,
      setup(ext) {
        for (const provider of providers) ext.provider(provider);
      },
    });
  try {
    await berry.install(extension([override, { ...second }]));
    assertEquals(models.getProvider(builtin.id), override);
    assert(models.getProvider("second") !== second);
    await berry.install(extension([replacement]));
    assertEquals(models.getProvider(builtin.id), replacement);
    assertEquals(models.getProvider("second"), second);
    berry.uninstall("provider");
    assertEquals(models.getProvider(builtin.id), builtin);
    await berry.install(extension([override]));
  } finally {
    await berry.close();
  }
  assertEquals(models.getProvider(builtin.id), builtin);
});

Deno.test("provider conflicts leave staged overrides and the running load untouched", async () => {
  const models = createModels();
  const builtin = fauxProvider().provider;
  models.setProvider(builtin);
  const berry = await openBerry(new MemoryStorage(), models);
  const owned = { ...builtin, id: "custom" };
  try {
    await berry.install(defineExtension({
      name: "owner",
      initial: () => null,
      setup: (ext) => ext.provider(owned),
    }));
    await assertRejects(() =>
      berry.install(defineExtension({
        name: "conflict",
        initial: () => null,
        setup(ext) {
          ext.provider({ ...builtin });
          ext.provider({ ...owned });
        },
      }))
    );
    assertEquals(models.getProvider(builtin.id), builtin);
    assertEquals(models.getProvider("custom"), owned);
    await assertRejects(() =>
      berry.install(defineExtension({
        name: "owner",
        initial: () => null,
        setup(ext) {
          ext.provider({ ...builtin });
          ext.provider({ ...builtin });
        },
      }))
    );
    assertEquals(models.getProvider(builtin.id), builtin);
    assertEquals(models.getProvider("custom"), owned);
  } finally {
    await berry.close();
  }
});

Deno.test("reload and uninstall release only the providers owned by that extension load", async () => {
  const models = createModels();
  const berry = await openBerry(new MemoryStorage(), models);
  const original = { ...fauxProvider().provider, id: "custom" };
  const replacement = { ...original, name: "replacement" };
  let disposed = false;
  try {
    await berry.install(defineExtension({
      name: "provider",
      initial: () => null,
      setup(ext) {
        ext.provider(original);
        ext.provider({ ...original, id: "removed" });
        ext.onDispose(() => void (disposed = true));
      },
    }));
    await berry.install(defineExtension({
      name: "provider",
      initial: () => null,
      setup(ext) {
        ext.provider(replacement);
      },
    }));
    await until(() => disposed);
    assertEquals(models.getProvider("custom"), replacement);
    assertEquals(models.getProvider("removed"), undefined);
    await assertRejects(() =>
      berry.install(defineExtension({
        name: "other",
        initial: () => null,
        setup(ext) {
          ext.provider({ ...replacement });
        },
      }))
    );
    assertEquals(models.getProvider("custom"), replacement);
    berry.uninstall("provider");
    assertEquals(models.getProvider("custom"), undefined);

    await berry.install(defineExtension({
      name: "provider",
      initial: () => null,
      setup(ext) {
        ext.provider(replacement);
      },
    }));
    assert(models.getProvider("custom"));
  } finally {
    await berry.close();
  }
  assertEquals(models.getProvider("custom"), undefined);
});