import { and, count, eq, gte } from "drizzle-orm";
import db from "../db";
import { usage } from "../db/schema";
import { env } from "../env";

const DAY_MS = 24 * 60 * 60 * 1000;

const BURST_LIMIT = 5;
const BURST_WINDOW_MS = 60 * 1000;

const recentRequests = new Map<string, number[]>();
const inFlight = new Set<string>();

export type LimitResult =
  | { ok: true }
  | { ok: false; reason: "quota" | "burst" | "busy"; message: string };

export const isAdmin = (did: string) => did === env.ADMIN_DID;

/** Requests the user has made in the last 24 hours. */
export async function dailyUsage(did: string) {
  const [row] = await db
    .select({ count: count(usage.id) })
    .from(usage)
    .where(
      and(eq(usage.did, did), gte(usage.createdAt, new Date(Date.now() - DAY_MS))),
    );
  return row?.count ?? 0;
}

export async function remainingQueries(did: string) {
  if (isAdmin(did)) return Infinity;
  return Math.max(0, env.DAILY_QUERY_LIMIT - await dailyUsage(did));
}

export const quotaMessage = () =>
  `You've used all ${env.DAILY_QUERY_LIMIT} of your queries for the last 24 hours. Your quota frees up gradually, so please try again a little later.`;

/** Checks every limit without counting the request. Admins are exempt. */
export async function checkLimits(did: string): Promise<LimitResult> {
  if (isAdmin(did)) return { ok: true };

  if (inFlight.has(did)) {
    return {
      ok: false,
      reason: "busy",
      message: "I'm still working on your last message. One moment!",
    };
  }

  const now = Date.now();
  const recent = (recentRequests.get(did) ?? []).filter((t) =>
    now - t < BURST_WINDOW_MS
  );
  recentRequests.set(did, recent);
  if (recent.length >= BURST_LIMIT) {
    return {
      ok: false,
      reason: "burst",
      message: "You're sending messages quite fast. Give me a minute to catch up.",
    };
  }

  if (await remainingQueries(did) <= 0) {
    return { ok: false, reason: "quota", message: quotaMessage() };
  }

  return { ok: true };
}

/** Counts a request against the user's limits and marks it as in progress. */
export async function beginRequest(did: string) {
  recentRequests.set(did, [...(recentRequests.get(did) ?? []), Date.now()]);
  inFlight.add(did);
  await db.insert(usage).values({ did });
}

export function endRequest(did: string) {
  inFlight.delete(did);
}