import { beforeEach, describe, expect, mock, test } from "bun:test";
import { fakeD1 } from "../metering/fake-d1.test-helper";
import { memoryScopes } from "../scoped-db/testing";

let d1 = fakeD1();
const stopped: string[] = [];
const fakeEnv: Record<string, unknown> = {
  get DB() {
    return d1.binding;
  },
  KV: { get: async () => null, put: async () => undefined, delete: async () => undefined },
  EDITOR_CONTAINER: { getByName: (name: string) => ({ stop: async () => stopped.push(name) && 0 }) },
  METERING_PLANS: "",
  PUBLIC_URL: "https://forge.test",
};
mock.module("cloudflare:workers", () => ({
  env: fakeEnv,
  waitUntil: () => undefined,
}));
let scopes = memoryScopes();
mock.module("../scoped-db", () => ({
  repoDb: (id: string) => scopes.repoDb(id),
  accountDb: (id: string) => scopes.accountDb(id),
  stmt: scopes.module.stmt,
}));
const locators = new Map<string, string>();
const actualLocators = await import("../scoped-db/locators");
mock.module("../scoped-db/locators", () => ({
  ...actualLocators,
  locate: async (kind: string, key: string) => locators.get(`${kind}:${key}`) ?? null,
  registerLocator: async (kind: string, key: string, scopeId: string) => {
    locators.set(`${kind}:${key}`, scopeId);
  },
  locatedScopes: async (kind: string) => [
    ...new Set([...locators].filter(([key]) => key.startsWith(`${kind}:`)).map(([, scope]) => scope)),
  ],
  locatedByPrefix: async (kind: string, prefix: string) => [
    ...new Set([...locators].filter(([key]) => key.startsWith(`${kind}:${prefix}`)).map(([, scope]) => scope)),
  ],
  forgetLocator: async (kind: string, key: string) => {
    locators.delete(`${kind}:${key}`);
  },
}));

const { closeContainerSession, closeExpiredContainerSessions, closeOrphanSession, openContainerSession } =
  await import("./container-meter");
const { setPlan } = await import("../metering/quota");

const account = () => scopes.raw("account", "u1");

function billed(): number {
  const row = account().query("SELECT SUM(quantity) AS q FROM usage_events WHERE meter = 'compute.editor'").get() as {
    q: number | null;
  } | null;
  return row?.q ?? 0;
}

function backdate(id: string, ms: number) {
  account().query("UPDATE editor_container_sessions SET started_at = started_at - ? WHERE id = ?").run(ms, id);
}

beforeEach(async () => {
  d1 = fakeD1();
  scopes = memoryScopes();
  locators.clear();
  d1.db.exec(
    "INSERT INTO accounts (id, name, kind) VALUES ('u1', 'alice', 'user'); INSERT INTO repos VALUES ('r1', 'u1', 'web');",
  );
  await setPlan("u1", "pro", "billing");
});

describe("container sessions", () => {
  test("a disabled account cannot start a container", async () => {
    d1.db.exec("UPDATE accounts SET disabled_at = 1, disabled_reason = 'payment_dispute' WHERE id = 'u1'");
    await expect(openContainerSession("u1", "r1", "main", 60)).rejects.toThrow("payment dispute");
  });

  test("bill at most what was held, and a retried close never bills twice", async () => {
    const hold = await openContainerSession("u1", "r1", "main", 60);
    backdate(hold.id, 600_000);
    expect(await closeContainerSession(hold, "stopped")).toBe(60);
    expect(await closeContainerSession(hold, "stopped")).toBe(60);
    expect(billed()).toBe(60);
  });

  test("a close retried after a failed commit still bills the original stop time", async () => {
    const hold = await openContainerSession("u1", "r1", "main", 3600);
    const stoppedAt = Date.now() + 30_000;
    account()
      .query("UPDATE editor_container_sessions SET stopped_at = ?, stop_reason = 'idle' WHERE id = ?")
      .run(stoppedAt, hold.id);
    const seconds = await closeContainerSession(hold, "stopped", stoppedAt + 600_000);
    expect(Math.round(seconds)).toBe(30);
    expect(Math.round(billed())).toBe(30);
  });

  test("the account alarm closes rows whose hold expired, billing only the hold", async () => {
    const hold = await openContainerSession("u1", "r1", "main", 120);
    backdate(hold.id, 10 * 3600_000);
    account()
      .query("UPDATE usage_reservations SET expires_at = ? WHERE id = ?")
      .run(Date.now() - 1, hold.id);
    expect(await closeExpiredContainerSessions("u1")).toBe(1);
    expect(stopped).toEqual(["u1:r1:main"]);
    expect(billed()).toBe(120);
    const open = account().query("SELECT COUNT(*) AS n FROM editor_container_sessions WHERE stopped_at IS NULL").get();
    expect(open).toEqual({ n: 0 });
  });

  test("opening a session indexes its account under the repository", async () => {
    await openContainerSession("u1", "r1", "main", 60);
    expect(locators.get("editor_repo_account:r1:u1")).toBe("u1");
  });

  test("orphans with a hold bill at most the hold", async () => {
    const hold = await openContainerSession("u1", "r1", "main", 90);
    backdate(hold.id, 3600_000);
    expect(await closeOrphanSession(hold.id, "stopped")).toBe(90);
    expect(billed()).toBe(90);
  });

  test("orphans without a hold are capped at the session cap", async () => {
    const hold = await openContainerSession("u1", "r1", "main", 60);
    account().query("DELETE FROM usage_reservations WHERE id = ?").run(hold.id);
    backdate(hold.id, 100 * 3600_000);
    expect(await closeOrphanSession(hold.id, "stopped")).toBe(8 * 3600);
  });
});