import { DurableObject } from "cloudflare:workers";
import { backoffMs, flushIsolating, ISOLATE_AFTER } from "../do-flush";
import { accountDb } from "../scoped-db";
import { eventStatements, periodOf } from "./record";
import type { UsageInput } from "./record";

const FLUSH_DELAY_MS = 5_000;
const RETRY_DELAY_MS = 30_000;
/** D1 batches are one transaction; this keeps each well under the statement limits. */
const FLUSH_BATCH = 100;
/** Flushed keys are remembered this long so a retried request is still recognised as a duplicate. */
const LEDGER_MS = 24 * 60 * 60 * 1000;

type Pending = { key: string; input: string; at: number };

/**
 * One per billed account: buffers high-volume usage events (git operations, downloads) and lands them in
 * the account's AccountState as one batched transaction every few seconds instead of one write per request. D1 stays the billing
 * source of truth: events keep their idempotency keys and are inserted with INSERT OR IGNORE, so a flush
 * that is retried after a partial failure never double counts. Quota checks read D1 and therefore see
 * buffered usage up to FLUSH_DELAY_MS late; anything that must be exact uses `recordUsage` directly.
 */
export class MeterBuffer extends DurableObject<Env> {
  private ready = false;

  private init(): void {
    if (this.ready) return;
    this.ctx.storage.sql.exec(`CREATE TABLE IF NOT EXISTS events (
      key TEXT PRIMARY KEY,
      input TEXT NOT NULL,
      at INTEGER NOT NULL,
      flushed_at INTEGER
    )`);
    this.ctx.storage.sql.exec("CREATE INDEX IF NOT EXISTS events_pending ON events(flushed_at, at)");
    this.ctx.storage.sql.exec(`CREATE TABLE IF NOT EXISTS dead_letters (
      key TEXT PRIMARY KEY,
      input TEXT NOT NULL,
      error TEXT NOT NULL,
      at INTEGER NOT NULL
    )`);
    this.ready = true;
  }

  /** Buffers an event; false when this key was already seen (a retry). */
  async add(input: UsageInput, accountId: string, at: number): Promise<boolean> {
    this.init();
    const stored: UsageInput = { ...input, accountId, at };
    const inserted =
      this.ctx.storage.sql
        .exec(
          "INSERT OR IGNORE INTO events (key, input, at) VALUES (?, ?, ?) RETURNING key",
          input.idempotencyKey,
          JSON.stringify(stored),
          at,
        )
        .toArray().length > 0;
    if (inserted) {
      const flushBy = Date.now() + FLUSH_DELAY_MS;
      const alarm = await this.ctx.storage.getAlarm();
      if (alarm === null || alarm > flushBy) await this.ctx.storage.setAlarm(flushBy);
    }
    return inserted;
  }

  async deadLetters(): Promise<{ key: string; input: string; error: string; at: number }[]> {
    this.init();
    return this.ctx.storage.sql
      .exec<{ key: string; input: string; error: string; at: number }>(
        "SELECT key, input, error, at FROM dead_letters ORDER BY at",
      )
      .toArray();
  }

  async has(key: string): Promise<boolean> {
    this.init();
    return this.ctx.storage.sql.exec("SELECT 1 FROM events WHERE key = ?", key).toArray().length > 0;
  }

  private accountOf(pending: Pending[]): string {
    return (JSON.parse(pending[0]?.input ?? "{}") as Partial<UsageInput>).accountId ?? "";
  }

  override async alarm(): Promise<void> {
    this.init();
    const sql = this.ctx.storage.sql;
    const pending = sql
      .exec<Pending>("SELECT key, input, at FROM events WHERE flushed_at IS NULL ORDER BY at LIMIT ?", FLUSH_BATCH)
      .toArray();
    if (pending.length > 0) {
      const attempts = (await this.ctx.storage.get<number>("attempts")) ?? 0;
      try {
        const poison = await flushIsolating(
          pending,
          (row) => {
            const input: UsageInput = JSON.parse(row.input);
            return eventStatements(input, input.accountId, periodOf(row.at), row.at);
          },
          attempts >= ISOLATE_AFTER,
          // One buffer per billed account, so every pending row lands in that account's object.
          (statements) => accountDb(this.accountOf(pending)).batch(statements),
        );
        for (const { item, error } of poison) {
          console.error("Meter event dead-lettered", item.key, error);
          sql.exec(
            "INSERT OR REPLACE INTO dead_letters (key, input, error, at) VALUES (?, ?, ?, ?)",
            item.key,
            item.input,
            error,
            Date.now(),
          );
        }
      } catch (error) {
        console.error("Meter buffer flush failed", pending.length, attempts + 1, error);
        await this.ctx.storage.put("attempts", attempts + 1);
        await this.ctx.storage.setAlarm(Date.now() + backoffMs(attempts + 1, RETRY_DELAY_MS));
        return;
      }
      if (attempts > 0) await this.ctx.storage.delete("attempts");
      const now = Date.now();
      this.ctx.storage.transactionSync(() => {
        // Dead-lettered events are marked flushed too, so the ledger still rejects their retries.
        for (const row of pending) sql.exec("UPDATE events SET flushed_at = ? WHERE key = ?", now, row.key);
      });
    }
    sql.exec("DELETE FROM events WHERE flushed_at IS NOT NULL AND flushed_at < ?", Date.now() - LEDGER_MS);
    const left = sql.exec<{ count: number }>("SELECT COUNT(*) AS count FROM events WHERE flushed_at IS NULL").one();
    if (left.count > 0)
      await this.ctx.storage.setAlarm(Date.now() + (pending.length === FLUSH_BATCH ? 0 : FLUSH_DELAY_MS));
    else if (sql.exec<{ count: number }>("SELECT COUNT(*) AS count FROM events").one().count > 0) {
      await this.ctx.storage.setAlarm(Date.now() + LEDGER_MS);
    }
  }
}