import type { Plan } from "@/lib/db/schema";

import { and, count, eq, gte, isNotNull } from "drizzle-orm";

import { PLAN_LIMITS, type QuotaLimits } from "@/lib/billing/plans";
import { db } from "@/lib/db";
import { member, orgSubscription, sites, siteTests } from "@/lib/db/schema";

export type QuotaKind = "tests" | "sites" | "seats";

export type QuotaResult = {
  kind: QuotaKind;
  used: number;
  limit: number;
  ok: boolean;
  plan: Plan;
  percent: number;
};

export class QuotaExceededError extends Error {
  constructor(public readonly result: QuotaResult) {
    super(
      `Quota exceeded for ${result.kind} on plan ${result.plan}: ${result.used}/${result.limit}`,
    );
    this.name = "QuotaExceededError";
  }
}

export async function getOrgPlan(organizationId: string): Promise<Plan> {
  const rows = await db
    .select({ plan: orgSubscription.plan })
    .from(orgSubscription)
    .where(eq(orgSubscription.organizationId, organizationId))
    .limit(1);
  return rows[0]?.plan ?? "free";
}

export function getCurrentMonthStart(): Date {
  const now = new Date();
  return new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1, 0, 0, 0));
}

async function countUsage(organizationId: string, kind: QuotaKind): Promise<number> {
  if (kind === "sites") {
    const rows = await db
      .select({ n: count() })
      .from(sites)
      .where(eq(sites.organizationId, organizationId));
    return rows[0]?.n ?? 0;
  }
  if (kind === "seats") {
    const rows = await db
      .select({ n: count() })
      .from(member)
      .where(eq(member.organizationId, organizationId));
    return rows[0]?.n ?? 0;
  }
  const monthStart = getCurrentMonthStart();
  const rows = await db
    .select({ n: count() })
    .from(siteTests)
    .where(
      and(
        eq(siteTests.organizationId, organizationId),
        gte(siteTests.createdAt, monthStart.toISOString()),
        isNotNull(siteTests.organizationId),
      ),
    );
  return rows[0]?.n ?? 0;
}

export async function checkQuota(organizationId: string, kind: QuotaKind): Promise<QuotaResult> {
  const plan = await getOrgPlan(organizationId);
  const limits: QuotaLimits = PLAN_LIMITS[plan];
  const limitKey: keyof QuotaLimits =
    kind === "tests" ? "testsPerMonth" : kind === "sites" ? "sites" : "seats";
  const limit = limits[limitKey];
  const used = await countUsage(organizationId, kind);
  const percent = limit === Number.POSITIVE_INFINITY ? 0 : used / Math.max(limit, 1);
  return {
    kind,
    used,
    limit,
    ok: used < limit,
    plan,
    percent,
  };
}

export async function assertQuota(organizationId: string, kind: QuotaKind): Promise<QuotaResult> {
  const result = await checkQuota(organizationId, kind);
  if (!result.ok) {
    throw new QuotaExceededError(result);
  }
  return result;
}