// src/lib/mcp/quota.test.ts — monthly MCP quota (ELEVA-110).
// Pure tests (isOverMcpQuota) always run; the usage-SUM and month-rollover
// tests need Postgres + the test DB migrated (0032) and are gated on DB_TESTS=1
// like auth.test.ts.
import { beforeEach, describe, expect, mock, test } from "bun:test";

// server-only throws outside Next.js; stub it.
mock.module("server-only", () => ({}));
mock.module("next/cache", () => ({
  revalidatePath: () => {},
  revalidateTag: () => {},
  updateTag: () => {},
  unstable_cache: (fn: unknown) => fn,
}));

const RUN_DB = process.env.DB_TESTS === "1";

describe("isOverMcpQuota (pure)", () => {
  test("under the limit → not over", async () => {
    const { isOverMcpQuota } = await import("./quota");
    expect(isOverMcpQuota(0, 1000)).toBe(false);
    expect(isOverMcpQuota(999, 1000)).toBe(false);
  });

  test("usage == limit → over (SUM ≥ limit rejects)", async () => {
    const { isOverMcpQuota } = await import("./quota");
    expect(isOverMcpQuota(1000, 1000)).toBe(true);
  });

  test("usage > limit → over", async () => {
    const { isOverMcpQuota } = await import("./quota");
    expect(isOverMcpQuota(1001, 1000)).toBe(true);
  });

  test("-1 = unlimited, never over", async () => {
    const { isOverMcpQuota } = await import("./quota");
    expect(isOverMcpQuota(0, -1)).toBe(false);
    expect(isOverMcpQuota(10_000_000, -1)).toBe(false);
  });

  test("limit 0 (expired org) → any usage incl. 0 is over", async () => {
    const { isOverMcpQuota } = await import("./quota");
    expect(isOverMcpQuota(0, 0)).toBe(true);
    expect(isOverMcpQuota(5, 0)).toBe(true);
  });
});

describe("resolveMcpAllowance (pure)", () => {
  test("within quota → within", async () => {
    const { resolveMcpAllowance } = await import("./quota");
    expect(resolveMcpAllowance(5, 1000, 0)).toBe("within");
  });

  test("over quota, credits remain → credit", async () => {
    const { resolveMcpAllowance } = await import("./quota");
    expect(resolveMcpAllowance(1000, 1000, 3)).toBe("credit");
  });

  test("over quota, no credits → blocked", async () => {
    const { resolveMcpAllowance } = await import("./quota");
    expect(resolveMcpAllowance(1000, 1000, 0)).toBe("blocked");
  });

  test("unlimited plan never touches credits, always within", async () => {
    const { resolveMcpAllowance } = await import("./quota");
    expect(resolveMcpAllowance(10_000_000, -1, 0)).toBe("within");
  });
});

describe("resolveEffectiveMonthlyUsage (pure)", () => {
  test("current month → the stored counter", async () => {
    const { resolveEffectiveMonthlyUsage } = await import("./quota");
    expect(resolveEffectiveMonthlyUsage("2026-07", 42, "2026-07")).toBe(42);
  });

  test("stale month → 0 (unspent this month)", async () => {
    const { resolveEffectiveMonthlyUsage } = await import("./quota");
    expect(resolveEffectiveMonthlyUsage("2026-06", 999, "2026-07")).toBe(0);
  });

  test("null month (never used) → 0", async () => {
    const { resolveEffectiveMonthlyUsage } = await import("./quota");
    expect(resolveEffectiveMonthlyUsage(null, 0, "2026-07")).toBe(0);
  });
});

describe.skipIf(!RUN_DB)("assertMcpQuota — atomic UPDATE (postgres)", () => {
  let db: typeof import("@/lib/db/client").db;
  let authSchema: typeof import("@/lib/db/auth-schema");

  const ORG = "atomic-quota-org";

  async function seedOrg(opts: {
    mcpMonthlyUsage?: number;
    mcpUsageMonth?: string | null;
    mcpCreditBalance?: number;
  }): Promise<void> {
    await db
      .insert(authSchema.organization)
      .values({
        id: ORG,
        name: "Atomic Quota Org",
        slug: ORG,
        createdAt: new Date(),
        mcpMonthlyUsage: opts.mcpMonthlyUsage ?? 0,
        mcpUsageMonth: opts.mcpUsageMonth ?? null,
        mcpCreditBalance: opts.mcpCreditBalance ?? 0,
      })
      .onConflictDoUpdate({
        target: authSchema.organization.id,
        set: {
          mcpMonthlyUsage: opts.mcpMonthlyUsage ?? 0,
          mcpUsageMonth: opts.mcpUsageMonth ?? null,
          mcpCreditBalance: opts.mcpCreditBalance ?? 0,
        },
      });
  }

  async function readOrg(): Promise<{
    mcpMonthlyUsage: number;
    mcpUsageMonth: string | null;
    mcpCreditBalance: number;
  }> {
    const { eq } = await import("drizzle-orm");
    const [row] = await db
      .select({
        mcpMonthlyUsage: authSchema.organization.mcpMonthlyUsage,
        mcpUsageMonth: authSchema.organization.mcpUsageMonth,
        mcpCreditBalance: authSchema.organization.mcpCreditBalance,
      })
      .from(authSchema.organization)
      .where(eq(authSchema.organization.id, ORG));
    return row;
  }

  beforeEach(async () => {
    ({ db } = await import("@/lib/db/client"));
    authSchema = await import("@/lib/db/auth-schema");
    const { assertTestDb } = await import("@/lib/db/assert-test-db");
    await assertTestDb();
  });

  test("fresh org, under limit → resolves and initializes the counter to 1", async () => {
    const { assertMcpQuota, currentMonth } = await import("./quota");
    await seedOrg({});
    await assertMcpQuota(ORG, async () => 1000);
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(1);
    expect(row.mcpUsageMonth).toBe(currentMonth());
  });

  test("at limit, no credits → throws and leaves the counter unchanged", async () => {
    const { assertMcpQuota, McpQuotaError, currentMonth } = await import(
      "./quota"
    );
    await seedOrg({ mcpMonthlyUsage: 1000, mcpUsageMonth: currentMonth() });
    await expect(assertMcpQuota(ORG, async () => 1000)).rejects.toThrow(
      McpQuotaError,
    );
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(1000); // not incremented — request was rejected
  });

  test("at limit, one credit remaining → resolves, increments usage, spends the credit", async () => {
    const { assertMcpQuota, currentMonth } = await import("./quota");
    await seedOrg({
      mcpMonthlyUsage: 1000,
      mcpUsageMonth: currentMonth(),
      mcpCreditBalance: 1,
    });
    await assertMcpQuota(ORG, async () => 1000);
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(1001);
    expect(row.mcpCreditBalance).toBe(0);
  });

  test("within quota with credits available → credits untouched", async () => {
    const { assertMcpQuota, currentMonth } = await import("./quota");
    await seedOrg({
      mcpMonthlyUsage: 5,
      mcpUsageMonth: currentMonth(),
      mcpCreditBalance: 10,
    });
    await assertMcpQuota(ORG, async () => 1000);
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(6);
    expect(row.mcpCreditBalance).toBe(10); // untouched — request was within plan quota
  });

  test("stale usage month rolls over: resets to 1 even though the old count was over limit", async () => {
    const { assertMcpQuota } = await import("./quota");
    await seedOrg({ mcpMonthlyUsage: 999_999, mcpUsageMonth: "2020-01" });
    await assertMcpQuota(ORG, async () => 1000);
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(1);
  });

  test("unlimited plan (-1) always resolves and still counts usage", async () => {
    const { assertMcpQuota, currentMonth } = await import("./quota");
    await seedOrg({
      mcpMonthlyUsage: 10_000_000,
      mcpUsageMonth: currentMonth(),
    });
    await assertMcpQuota(ORG, async () => -1);
    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(10_000_001);
  });

  // M1 regression: the old design read the balance, THEN wrote a
  // fire-and-forget decrement in a separate statement — N concurrent callers
  // could all read the same pre-write balance and all pass. This proves the
  // atomic UPDATE...WHERE...RETURNING actually serializes concurrent racers
  // on Postgres's row lock instead of trusting a stale JS-side read.
  test("concurrency: N requests racing the LAST unit of plan quota → exactly 1 passes", async () => {
    const { assertMcpQuota, McpQuotaError, currentMonth } = await import(
      "./quota"
    );
    const LIMIT = 1000;
    const RACERS = 12;
    await seedOrg({
      mcpMonthlyUsage: LIMIT - 1,
      mcpUsageMonth: currentMonth(),
    });

    const results = await Promise.allSettled(
      Array.from({ length: RACERS }, () =>
        assertMcpQuota(ORG, async () => LIMIT),
      ),
    );

    const fulfilled = results.filter((r) => r.status === "fulfilled");
    const rejected = results.filter(
      (r) =>
        r.status === "rejected" &&
        (r as PromiseRejectedResult).reason instanceof McpQuotaError,
    );
    expect(fulfilled.length).toBe(1);
    expect(rejected.length).toBe(RACERS - 1);

    const row = await readOrg();
    expect(row.mcpMonthlyUsage).toBe(LIMIT); // exactly one increment landed, not RACERS
  });

  // Same race, but over the plan quota with exactly ONE credit left — proves
  // the credit decrement (the other TOCTOU surface: the old code's
  // fire-and-forget `GREATEST(balance - 1, 0)`) is equally race-proof.
  test("concurrency: N requests racing the LAST credit → exactly 1 passes, balance floors at 0", async () => {
    const { assertMcpQuota, McpQuotaError, currentMonth } = await import(
      "./quota"
    );
    const LIMIT = 1000;
    const RACERS = 12;
    await seedOrg({
      mcpMonthlyUsage: LIMIT,
      mcpUsageMonth: currentMonth(),
      mcpCreditBalance: 1,
    });

    const results = await Promise.allSettled(
      Array.from({ length: RACERS }, () =>
        assertMcpQuota(ORG, async () => LIMIT),
      ),
    );

    const fulfilled = results.filter((r) => r.status === "fulfilled");
    const rejected = results.filter(
      (r) =>
        r.status === "rejected" &&
        (r as PromiseRejectedResult).reason instanceof McpQuotaError,
    );
    expect(fulfilled.length).toBe(1);
    expect(rejected.length).toBe(RACERS - 1);

    const row = await readOrg();
    expect(row.mcpCreditBalance).toBe(0); // spent exactly once, never negative
    expect(row.mcpMonthlyUsage).toBe(LIMIT + 1); // only the winner's request counted
  });
});

describe.skipIf(!RUN_DB)("getOrgMcpMonthlyUsage (postgres)", () => {
  let db: typeof import("@/lib/db/client").db;
  let authSchema: typeof import("@/lib/db/auth-schema");
  const ORG = "display-usage-org";

  beforeEach(async () => {
    ({ db } = await import("@/lib/db/client"));
    authSchema = await import("@/lib/db/auth-schema");
    const { assertTestDb } = await import("@/lib/db/assert-test-db");
    await assertTestDb();
    await db
      .insert(authSchema.organization)
      .values({
        id: ORG,
        name: "Display Usage Org",
        slug: ORG,
        createdAt: new Date(),
      })
      .onConflictDoNothing();
  });

  test("reads the current-month counter", async () => {
    const { eq } = await import("drizzle-orm");
    const { getOrgMcpMonthlyUsage, currentMonth } = await import("./quota");
    await db
      .update(authSchema.organization)
      .set({ mcpMonthlyUsage: 77, mcpUsageMonth: currentMonth() })
      .where(eq(authSchema.organization.id, ORG));
    expect(await getOrgMcpMonthlyUsage(ORG)).toBe(77);
  });

  test("stale month → 0, no unmigrated org → 0", async () => {
    const { eq } = await import("drizzle-orm");
    const { getOrgMcpMonthlyUsage } = await import("./quota");
    await db
      .update(authSchema.organization)
      .set({ mcpMonthlyUsage: 500, mcpUsageMonth: "2020-01" })
      .where(eq(authSchema.organization.id, ORG));
    expect(await getOrgMcpMonthlyUsage(ORG)).toBe(0);
    expect(await getOrgMcpMonthlyUsage("no-such-org-at-all")).toBe(0);
  });
});

describe.skipIf(!RUN_DB)("MCP monthly metering (postgres)", () => {
  let db: typeof import("@/lib/db/client").db;
  let schema: typeof import("@/lib/db/domain-schema");
  let authSchema: typeof import("@/lib/db/auth-schema");

  const ORG = "quota-org";
  const OTHER_ORG = "quota-other-org";
  const USER = "seed-owner";

  function lastMonth(): string {
    const d = new Date();
    d.setUTCDate(1);
    d.setUTCMonth(d.getUTCMonth() - 1);
    return d.toISOString().slice(0, 7);
  }

  async function ensureOrg(id: string, name: string, slug: string) {
    await db
      .insert(authSchema.organization)
      .values({ id, name, slug, createdAt: new Date() })
      .onConflictDoNothing();
  }

  async function ensureUser(id: string) {
    await db
      .insert(authSchema.user)
      .values({
        id,
        name: "Seed User",
        email: `${id}@seo-insight.local`,
        emailVerified: true,
        createdAt: new Date(),
        updatedAt: new Date(),
      })
      .onConflictDoNothing();
  }

  async function seedKey(opts: {
    orgId: string;
    requestMonth: string | null;
    monthlyCount: number;
  }): Promise<string> {
    const { generateApiKey } = await import("./keys");
    const { prefix, hash } = generateApiKey();
    const id = `qkey-${prefix}-${Math.random().toString(36).slice(2)}`;
    await db.insert(schema.mcpApiKeys).values({
      id,
      orgId: opts.orgId,
      userId: USER,
      name: "Quota test key",
      prefix,
      keyHash: hash,
      requestMonth: opts.requestMonth,
      monthlyCount: opts.monthlyCount,
    });
    return id;
  }

  beforeEach(async () => {
    ({ db } = await import("@/lib/db/client"));
    schema = await import("@/lib/db/domain-schema");
    authSchema = await import("@/lib/db/auth-schema");
    const { assertTestDb } = await import("@/lib/db/assert-test-db");
    await assertTestDb();
    await db.delete(schema.mcpApiKeys);
    await ensureUser(USER);
    await ensureOrg(ORG, "Quota Org", "quota-org");
    await ensureOrg(OTHER_ORG, "Quota Other Org", "quota-other-org");
  });

  test("month rollover: increment-update resets monthlyCount to 1 and bumps requestMonth", async () => {
    const { currentMonth } = await import("./quota");
    const { eq, sql } = await import("drizzle-orm");
    const keyId = await seedKey({
      orgId: ORG,
      requestMonth: lastMonth(),
      monthlyCount: 55,
    });

    // The exact fire-and-forget update resolveMcpOrg runs (src/lib/mcp/auth.ts).
    const month = currentMonth();
    const bump = () =>
      db
        .update(schema.mcpApiKeys)
        .set({
          lastUsedAt: new Date(),
          requestCount: sql`${schema.mcpApiKeys.requestCount} + 1`,
          monthlyCount: sql`CASE WHEN ${schema.mcpApiKeys.requestMonth} = ${month} THEN ${schema.mcpApiKeys.monthlyCount} + 1 ELSE 1 END`,
          requestMonth: month,
        })
        .where(eq(schema.mcpApiKeys.id, keyId));

    await bump();
    let [row] = await db
      .select({
        requestMonth: schema.mcpApiKeys.requestMonth,
        monthlyCount: schema.mcpApiKeys.monthlyCount,
        requestCount: schema.mcpApiKeys.requestCount,
      })
      .from(schema.mcpApiKeys)
      .where(eq(schema.mcpApiKeys.id, keyId));
    expect(row.requestMonth).toBe(month);
    expect(row.monthlyCount).toBe(1); // reset, not 56
    expect(row.requestCount).toBe(1); // cumulative keeps counting

    // Same-month second hit → plain increment.
    await bump();
    [row] = await db
      .select({
        requestMonth: schema.mcpApiKeys.requestMonth,
        monthlyCount: schema.mcpApiKeys.monthlyCount,
        requestCount: schema.mcpApiKeys.requestCount,
      })
      .from(schema.mcpApiKeys)
      .where(eq(schema.mcpApiKeys.id, keyId));
    expect(row.requestMonth).toBe(month);
    expect(row.monthlyCount).toBe(2);
    expect(row.requestCount).toBe(2);
  });
});