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

let d1 = fakeD1();
const kv = new Map<string, string>();
const fakeEnv: Record<string, unknown> = {
  get DB() {
    return d1.binding;
  },
  KV: {
    get: async (key: string) => kv.get(key) ?? null,
    put: async (key: string, value: string) => void kv.set(key, value),
    delete: async (key: string) => void kv.delete(key),
  },
  USAGE: { writeDataPoint: () => undefined },
  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 accountSql = (id: string) => scopes.raw("account", id);
const counterRows = new Map<string, [number, number, number]>();
mock.module("../counters", () => ({
  addCounter: async (kind: string, scope: string, key: string, delta: [number, number, number]) => {
    const id = `${kind}|${scope}|${key}`;
    const current = counterRows.get(id) ?? [0, 0, 0];
    const next: [number, number, number] = [current[0] + delta[0], current[1] + delta[1], current[2] + delta[2]];
    counterRows.set(id, next);
    return next;
  },
  counterTotals: async (kind: string, scope: string, key: string) =>
    counterRows.get(`${kind}|${scope}|${key}`) ?? [0, 0, 0],
}));

const { setGauge } = await import("./record");
const {
  AccountDisabledError,
  PlatformLimitError,
  QuotaExceededError,
  assertWithinQuota,
  billingSettings,
  planBlock,
  reserveUsage,
  setPlan,
  setSpendingLimit,
  withinSpendingLimit,
} = await import("./quota");
const { AI_TOKEN_HOLD, aiMeter, assertPlatformAiBudget, dailyAiCeilingCents } = await import("./ai");

const GB = 1024 ** 3;

beforeEach(() => {
  d1 = fakeD1();
  scopes = memoryScopes();
  kv.clear();
  counterRows.clear();
  fakeEnv.METERING_PLANS = "";
  delete fakeEnv.AI_DAILY_SPEND_CEILING_CENTS;
  delete fakeEnv.PAST_DUE_GRACE_DAYS;
  d1.db.exec(`
    INSERT INTO accounts (id, name, kind) VALUES ('u1', 'alice', 'user'), ('u2', 'bob', 'user'), ('o1', 'acme', 'org');
    INSERT INTO org_members VALUES ('o1', 'u1', 'owner', 0), ('o1', 'u2', 'member', 0);
  `);
});

function pastDue(accountId: string, at: number) {
  accountSql(accountId)
    .query(
      `INSERT INTO billing_customers (account_id, stripe_customer_id, subscription_id, subscription_status,
         subscription_event_at, past_due_since, updated_at) VALUES (?, 'cus_x', 'sub_x', 'past_due', ?, ?, 0)`,
    )
    .run(accountId, Date.now(), at);
}

describe("spending limits", () => {
  test("Pro without a chosen limit gets the default; an explicit limit wins", async () => {
    await setPlan("u1", "pro", "billing");
    expect((await billingSettings("u1")).spendingLimitCents).toBe(DEFAULT_PRO_SPENDING_LIMIT_CENTS);
    await setSpendingLimit("u1", 20_000);
    expect((await billingSettings("u1")).spendingLimitCents).toBe(20_000);
    expect((await billingSettings("u2")).spendingLimitCents).toBeNull();
  });

  test("a limit of 0 still allows included usage; the boundary is reach-but-not-exceed", async () => {
    await setPlan("u1", "pro", "billing");
    await setSpendingLimit("u1", 0);
    expect(await withinSpendingLimit("u1", "ai.calls")).toBe(true);
    await reserveUsage({ id: "r1", accountId: "u1", meter: "ai.calls", quantity: 1000 });
    expect(await withinSpendingLimit("u1", "ai.calls")).toBe(true);
    await expect(reserveUsage({ id: "r2", accountId: "u1", meter: "ai.calls", quantity: 1 })).rejects.toBeInstanceOf(
      QuotaExceededError,
    );
  });
});

describe("past_due grace", () => {
  test("Pro falls back to Free once past_due outlasts the grace period", async () => {
    await setPlan("u1", "pro", "billing");
    pastDue("u1", Date.now() - 3 * 86_400_000);
    expect((await billingSettings("u1")).plan.id).toBe("pro");
    fakeEnv.PAST_DUE_GRACE_DAYS = "2";
    expect((await billingSettings("u1")).plan.id).toBe("free");
    // A past_due row with no recorded start (predating migration 0047) starts its grace now.
    accountSql("u1").exec(
      "UPDATE billing_customers SET past_due_since = NULL, subscription_event_at = 0 WHERE account_id = 'u1'",
    );
    expect((await billingSettings("u1")).plan.id).toBe("pro");
    accountSql("u1").exec("UPDATE billing_customers SET past_due_since = 1 WHERE account_id = 'u1'");
    await setPlan("u1", "pro", "admin");
    expect((await billingSettings("u1")).plan.id).toBe("pro");
  });
});

describe("disabled accounts", () => {
  test("a disabled account and its org's members cannot meter anything", async () => {
    d1.db.exec("UPDATE accounts SET disabled_at = 1 WHERE id = 'o1'");
    await expect(assertWithinQuota("o1", "ai.calls")).rejects.toBeInstanceOf(AccountDisabledError);
    await expect(
      reserveUsage({ id: "d1", accountId: "o1", meter: "compute.editor", quantity: 1 }),
    ).rejects.toBeInstanceOf(AccountDisabledError);
    // A member's work on the org's repos bills the org, so it is refused too.
    const meter = aiMeter({ accountId: "o1", feature: "mergefix", scope: "member-run" });
    await expect(meter.beforeCall("glm")).rejects.toBeInstanceOf(AccountDisabledError);
    expect((await planBlock("o1", "storage.lfs", 1))?.toHttpError().status).toBe(403);
    await expect(assertWithinQuota("u2", "ai.calls")).resolves.toBeUndefined();
  });
});

describe("storage gauges (BILL-P9)", () => {
  test("concurrent writes cannot both fit into the same headroom", async () => {
    await setGauge("u1", "storage.lfs", 0.6 * GB);
    const results = await Promise.all([
      planBlock("u1", "storage.lfs", 0.3 * GB),
      planBlock("u1", "storage.lfs", 0.3 * GB),
    ]);
    expect(results.filter((blocked) => blocked === null)).toHaveLength(1);
    expect(await planBlock("u1", "storage.lfs", 0.1 * GB)).toBeNull();
  });
});

describe("AI", () => {
  test("token holds are reserved atomically and settled at the real count", async () => {
    fakeEnv.METERING_PLANS = JSON.stringify({ free: { hardCaps: { "ai.tokens": AI_TOKEN_HOLD + 100 } } });
    const first = aiMeter({ accountId: "u1", feature: "mergefix", scope: "a" });
    const second = aiMeter({ accountId: "u1", feature: "mergefix", scope: "b" });
    const index = await first.beforeCall("glm");
    await second.beforeCall("glm");
    await expect(
      aiMeter({ accountId: "u1", feature: "mergefix", scope: "c" }).beforeCall("glm"),
    ).rejects.toBeInstanceOf(QuotaExceededError);
    await first.afterCall(index, "glm", { inputTokens: 40, outputTokens: 10 });
    const row = accountSql("u1")
      .query("SELECT quantity, reserved FROM usage_periods WHERE account_id = 'u1' AND meter = 'ai.tokens'")
      .get() as { quantity: number; reserved: number };
    expect(row).toEqual({ quantity: 50, reserved: 100 });
  });

  test("the platform daily ceiling stops every account and fails closed", async () => {
    expect(dailyAiCeilingCents("off")).toBeNull();
    expect(dailyAiCeilingCents(undefined)).toBeNull();
    fakeEnv.AI_DAILY_SPEND_CEILING_CENTS = "nonsense";
    await expect(assertPlatformAiBudget()).rejects.toBeInstanceOf(PlatformLimitError);
    fakeEnv.AI_DAILY_SPEND_CEILING_CENTS = "1";
    await assertPlatformAiBudget();
    await setPlan("u2", "pro", "billing");
    // Another account's calls move the platform total; the check reads that one counter, not every account.
    const heavy = aiMeter({ accountId: "u2", feature: "mergefix", scope: "big" });
    const call = await heavy.beforeCall("glm");
    await heavy.afterCall(call, "glm", { inputTokens: 10_000_000, outputTokens: 0 });
    const meter = aiMeter({ accountId: "u1", feature: "mergefix", scope: "x" });
    await expect(meter.beforeCall("glm")).rejects.toBeInstanceOf(PlatformLimitError);
    await meter.beforeCall("byok-model", true);
  });
});