import { DurableObject } from "cloudflare:workers";
import type { RateLimitVerdict } from "./rate-limit-shard";

/**
 * Fixed-window counters, one row per key, kept in SQLite so an eviction or restart does not hand out a fresh
 * budget. Expired rows are pruned by an alarm at the earliest window end.
 */
export class RateLimiter extends DurableObject<Env> {
  private ready = false;

  private init(): void {
    if (this.ready) return;
    this.ctx.storage.sql.exec(`CREATE TABLE IF NOT EXISTS windows (
      key TEXT PRIMARY KEY,
      reset_at INTEGER NOT NULL,
      hits INTEGER NOT NULL
    )`);
    this.ctx.storage.sql.exec("CREATE INDEX IF NOT EXISTS windows_reset ON windows (reset_at)");
    this.ready = true;
  }

  /** Counts one hit against `key` and reports whether it is within `limit` for the current window. */
  async hit(key: string, limit: number, periodMs: number, now = Date.now()): Promise<RateLimitVerdict> {
    this.init();
    const resetAt = Math.floor(now / periodMs) * periodMs + periodMs;
    const row = this.ctx.storage.sql
      .exec<{ hits: number; reset_at: number }>(
        `INSERT INTO windows (key, reset_at, hits) VALUES (?1, ?2, 1)
         ON CONFLICT (key) DO UPDATE SET
           hits = CASE WHEN reset_at = ?2 THEN hits + 1 ELSE 1 END,
           reset_at = ?2
         RETURNING hits, reset_at`,
        key,
        resetAt,
      )
      .one();
    const alarm = await this.ctx.storage.getAlarm();
    if (alarm === null || alarm > resetAt) await this.ctx.storage.setAlarm(resetAt);
    return { success: row.hits <= limit, resetAt: row.reset_at };
  }

  async alarm(): Promise<void> {
    this.init();
    const now = Date.now();
    this.ctx.storage.sql.exec("DELETE FROM windows WHERE reset_at <= ?", now);
    const next = this.ctx.storage.sql
      .exec<{ at: number | null }>("SELECT MIN(reset_at) AS at FROM windows")
      .toArray()[0]?.at;
    if (next !== null && next !== undefined) await this.ctx.storage.setAlarm(next);
  }
}