import { DurableObject } from "cloudflare:workers";
import { first, stmt } from "./db";
import { backoffMs, flushIsolating, ISOLATE_AFTER } from "./do-flush";
import { afterHostingFlush } from "./sites/caps";
import { REPO_PACK_BYTES, REPO_STARS } from "./counter-keys";

/** How long increments sit in the object before D1 sees them; reads that must be exact go to the object. */
const FLUSH_DELAY_MS = 10_000;
const RETRY_DELAY_MS = 60_000;
const FLUSH_BATCH = 200;
/** Rows untouched this long are dropped once flushed, so an object only holds its working set. */
const IDLE_ROW_MS = 3 * 24 * 60 * 60 * 1000;

export type CounterKind = "site-usage" | "release-downloads" | "mcp-usage" | "repo-counts";
export type Totals = [number, number, number];
type Row = { key: string; a: number; b: number; c: number; meta: string | null };

async function count(sql: string, ...params: string[]): Promise<number> {
  return (await first<{ value: number }>(sql, ...params))?.value ?? 0;
}

/**
 * Stars and upvotes flush as a recount of their source rows, not the object's total: rows also disappear
 * without passing through the counter (account deletion cascades, access-revocation purges, discussion
 * deletes), so writing the object's total would put a stale number back. Pack bytes have no writer other
 * than this counter, so its absolute total is authoritative; a future recompute must call the object's `reset`.
 */
function repoCountStatement(repoId: string, row: Row): ReturnType<typeof stmt> {
  if (row.key === REPO_STARS)
    return stmt("UPDATE repos SET stars_count = (SELECT COUNT(*) FROM stars WHERE repo_id = ?1) WHERE id = ?1", repoId);
  if (row.key === REPO_PACK_BYTES)
    return stmt(
      `INSERT INTO repo_storage_estimates (repo_id, pack_bytes, updated_at)
       SELECT ?1, ?2, ?3 WHERE EXISTS (SELECT 1 FROM repos WHERE id = ?1)
       ON CONFLICT (repo_id) DO UPDATE SET pack_bytes = excluded.pack_bytes, updated_at = excluded.updated_at`,
      repoId,
      row.a,
      Date.now(),
    );
  const table = row.key.startsWith("d:") ? "discussions" : "discussion_comments";
  return stmt(
    `UPDATE ${table} SET upvotes = (SELECT COUNT(*) FROM discussion_votes WHERE subject_id = ?1) WHERE id = ?1`,
    row.key.slice(2),
  );
}

type Kind = {
  /** Current D1 totals for a key the object has not seen yet, so taking over a live counter loses nothing. */
  seed(scope: string, key: string): Promise<Totals>;
  /** Writes absolute totals, so a retried flush is idempotent. */
  flush(scope: string, rows: Row[]): ReturnType<typeof stmt>[];
  /** Runs after a successful flush with the billed accounts it touched; reads D1 and writes KV only. */
  afterFlush?(accountIds: string[]): Promise<void>;
  /** The seed counts source rows the caller already wrote, so the triggering delta is already in it. */
  seedIncludesDelta?(key: string): boolean;
};

const KINDS: Record<CounterKind, Kind> = {
  /** scope = site id, key = UTC day, totals = requests / cpu ms / bytes out, meta = billed account id. */
  "site-usage": {
    async seed(siteId, day) {
      const row = await first<{ requests: number; cpu_ms: number; bytes_out: number }>(
        "SELECT requests, cpu_ms, bytes_out FROM site_usage WHERE site_id = ? AND day = ?",
        siteId,
        day,
      );
      return [row?.requests ?? 0, row?.cpu_ms ?? 0, row?.bytes_out ?? 0];
    },
    flush(siteId, rows) {
      return rows.map((row) =>
        stmt(
          `INSERT INTO site_usage (site_id, day, account_id, requests, cpu_ms, bytes_out)
           SELECT ?1, ?2, ?3, ?4, ?5, ?6 WHERE EXISTS (SELECT 1 FROM sites WHERE id = ?1)
           ON CONFLICT (site_id, day) DO UPDATE SET requests = excluded.requests, cpu_ms = excluded.cpu_ms,
             bytes_out = excluded.bytes_out, account_id = excluded.account_id`,
          siteId,
          row.key,
          row.meta,
          row.a,
          row.b,
          row.c,
        ),
      );
    },
    afterFlush: afterHostingFlush,
  },
  /** scope = repo id, key = release asset id, totals[0] = downloads. */
  "release-downloads": {
    async seed(_repoId, assetId) {
      const row = await first<{ download_count: number }>(
        "SELECT download_count FROM release_assets WHERE id = ?",
        assetId,
      );
      return [row?.download_count ?? 0, 0, 0];
    },
    flush(_repoId, rows) {
      return rows.map((row) => stmt("UPDATE release_assets SET download_count = ? WHERE id = ?", row.a, row.key));
    },
  },
  /** scope = account id, key = `${UTC day}|${tool}`, totals = calls / errors / unused. */
  "mcp-usage": {
    async seed(accountId, key) {
      const [day, tool] = splitUsageKey(key);
      const row = await first<{ calls: number; errors: number }>(
        "SELECT calls, errors FROM mcp_usage_daily WHERE account_id = ? AND day = ? AND tool = ?",
        accountId,
        day,
        tool,
      );
      return [row?.calls ?? 0, row?.errors ?? 0, 0];
    },
    flush(accountId, rows) {
      return rows.map((row) => {
        const [day, tool] = splitUsageKey(row.key);
        return stmt(
          `INSERT INTO mcp_usage_daily (account_id, day, tool, calls, errors)
           SELECT ?1, ?2, ?3, ?4, ?5 WHERE EXISTS (SELECT 1 FROM accounts WHERE id = ?1)
           ON CONFLICT (account_id, day, tool) DO UPDATE SET calls = excluded.calls, errors = excluded.errors`,
          accountId,
          day,
          tool,
          row.a,
          row.b,
        );
      });
    },
  },
  /** scope = repo id; see REPO_STARS etc. Seeds count source rows so a lost increment heals on reseed. */
  "repo-counts": {
    async seed(repoId, key) {
      if (key === REPO_STARS)
        return [await count("SELECT COUNT(*) AS value FROM stars WHERE repo_id = ?", repoId), 0, 0];
      if (key === REPO_PACK_BYTES)
        return [await count("SELECT pack_bytes AS value FROM repo_storage_estimates WHERE repo_id = ?", repoId), 0, 0];
      return [await count("SELECT COUNT(*) AS value FROM discussion_votes WHERE subject_id = ?", key.slice(2)), 0, 0];
    },
    flush(repoId, rows) {
      return rows.map((row) => repoCountStatement(repoId, row));
    },
    seedIncludesDelta: (key) => key !== REPO_PACK_BYTES,
  },
};

function splitUsageKey(key: string): [string, string] {
  const at = key.indexOf("|");
  return [key.slice(0, at), key.slice(at + 1)];
}

type Scope = { kind: CounterKind; scope: string };

/**
 * Write-heavy counters (per-request site usage, asset downloads) accumulate in one object per scope instead
 * of being an UPDATE on D1's single writer each time. D1 receives absolute totals at most every
 * FLUSH_DELAY_MS per object, so its load is bounded by the number of active scopes rather than traffic.
 */
export class Counters extends DurableObject<Env> {
  private ready = false;
  private scopeSaved = false;

  private init(): void {
    if (this.ready) return;
    this.ctx.storage.sql.exec(`CREATE TABLE IF NOT EXISTS counters (
      key TEXT PRIMARY KEY,
      a INTEGER NOT NULL,
      b INTEGER NOT NULL,
      c INTEGER NOT NULL,
      meta TEXT,
      dirty INTEGER NOT NULL DEFAULT 0,
      touched_at INTEGER NOT NULL
    )`);
    // base_* = the totals D1 is known to hold (seed or last flush); reseed = 1 after `reset`, when a/b/c hold
    // only the unflushed delta and the D1 total must be re-read before the row is used or flushed.
    const columns = new Set(
      this.ctx.storage.sql
        .exec<{ name: string }>("PRAGMA table_info(counters)")
        .toArray()
        .map((column) => column.name),
    );
    for (const column of ["base_a INTEGER", "base_b INTEGER", "base_c INTEGER", "reseed INTEGER NOT NULL DEFAULT 0"]) {
      if (!columns.has(column.split(" ")[0])) this.ctx.storage.sql.exec(`ALTER TABLE counters ADD COLUMN ${column}`);
    }
    this.ctx.storage.sql.exec(`CREATE TABLE IF NOT EXISTS dead_letters (
      key TEXT NOT NULL,
      a INTEGER NOT NULL,
      b INTEGER NOT NULL,
      c INTEGER NOT NULL,
      meta TEXT,
      error TEXT NOT NULL,
      at INTEGER NOT NULL
    )`);
    this.ready = true;
  }

  private read(key: string): Row | null {
    return (
      this.ctx.storage.sql.exec<Row>("SELECT key, a, b, c, meta FROM counters WHERE key = ?", key).toArray()[0] ?? null
    );
  }

  /** Adds to a key's totals and returns the new totals. */
  async add(kind: CounterKind, scope: string, key: string, delta: Totals, meta: string | null): Promise<Totals> {
    this.init();
    if (!this.scopeSaved) {
      await this.ctx.storage.put<Scope>("scope", { kind, scope });
      this.scopeSaved = true;
    }
    let applied = delta;
    await this.reseedIfReset(kind, scope, key);
    if (!this.read(key)) {
      const seeded = await KINDS[kind].seed(scope, key);
      const inserted = this.ctx.storage.sql
        .exec(
          `INSERT OR IGNORE INTO counters (key, a, b, c, meta, touched_at, base_a, base_b, base_c)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING key`,
          key,
          ...seeded,
          meta,
          Date.now(),
          ...seeded,
        )
        .toArray().length;
      if (inserted && KINDS[kind].seedIncludesDelta?.(key)) applied = [0, 0, 0];
    }
    const row = this.ctx.storage.sql
      .exec<Row>(
        `UPDATE counters SET a = a + ?, b = b + ?, c = c + ?, meta = ?, dirty = 1, touched_at = ?
         WHERE key = ? RETURNING key, a, b, c, meta`,
        ...applied,
        meta,
        Date.now(),
        key,
      )
      .one();
    if ((await this.ctx.storage.getAlarm()) === null) await this.ctx.storage.setAlarm(Date.now() + FLUSH_DELAY_MS);
    return [row.a, row.b, row.c];
  }

  /**
   * Re-reads a key's total from D1 after something recomputed it, keeping increments the object has not flushed
   * yet: the row keeps only its unflushed delta (totals minus what D1 last received) and is marked for reseed,
   * so the fresh D1 total is added back before the row is next read or flushed.
   */
  async reset(key: string): Promise<void> {
    this.init();
    const row = this.ctx.storage.sql
      .exec<Row & { dirty: number; base_a: number | null; base_b: number | null; base_c: number | null }>(
        "SELECT key, a, b, c, meta, dirty, base_a, base_b, base_c FROM counters WHERE key = ? AND reseed = 0",
        key,
      )
      .toArray()[0];
    if (!row) return;
    if (!row.dirty) {
      this.ctx.storage.sql.exec("DELETE FROM counters WHERE key = ?", key);
      return;
    }
    if (row.base_a === null || row.base_b === null || row.base_c === null) {
      // A row written before base tracking existed: its unflushed share is unknown, so flush it as is.
      console.warn("Counter reset skipped for a row without a flush baseline", key);
      return;
    }
    this.ctx.storage.sql.exec(
      "UPDATE counters SET a = a - base_a, b = b - base_b, c = c - base_c, reseed = 1 WHERE key = ?",
      key,
    );
  }

  private async reseedIfReset(kind: CounterKind, scope: string, key: string): Promise<void> {
    const pending = this.ctx.storage.sql
      .exec<{ key: string }>("SELECT key FROM counters WHERE key = ? AND reseed = 1", key)
      .toArray();
    if (pending.length === 0) return;
    const seeded = await KINDS[kind].seed(scope, key);
    const keepDelta = KINDS[kind].seedIncludesDelta?.(key) ? 0 : 1;
    this.ctx.storage.sql.exec(
      `UPDATE counters SET a = ? + a * ?, b = ? + b * ?, c = ? + c * ?, base_a = ?, base_b = ?, base_c = ?, reseed = 0
       WHERE key = ? AND reseed = 1`,
      seeded[0],
      keepDelta,
      seeded[1],
      keepDelta,
      seeded[2],
      keepDelta,
      ...seeded,
      key,
    );
  }

  /** Exact current totals (D1 may lag by up to FLUSH_DELAY_MS). */
  async totals(kind: CounterKind, scope: string, key: string): Promise<Totals> {
    this.init();
    await this.reseedIfReset(kind, scope, key);
    const row = this.read(key);
    return row ? [row.a, row.b, row.c] : KINDS[kind].seed(scope, key);
  }

  /** Rows D1 rejected on their own; kept for inspection and replay instead of blocking later flushes. */
  async deadLetters(): Promise<(Row & { error: string; at: number })[]> {
    this.init();
    return this.ctx.storage.sql
      .exec<Row & { error: string; at: number }>("SELECT key, a, b, c, meta, error, at FROM dead_letters ORDER BY at")
      .toArray();
  }

  override async alarm(): Promise<void> {
    this.init();
    const target = await this.ctx.storage.get<Scope>("scope");
    if (!target) return;
    const sql = this.ctx.storage.sql;
    for (const { key } of sql.exec<{ key: string }>("SELECT key FROM counters WHERE reseed = 1").toArray())
      await this.reseedIfReset(target.kind, target.scope, key);
    const rows = sql
      .exec<Row>("SELECT key, a, b, c, meta FROM counters WHERE dirty = 1 AND reseed = 0 LIMIT ?", FLUSH_BATCH)
      .toArray();
    const kind = KINDS[target.kind];
    const attempts = (await this.ctx.storage.get<number>("attempts")) ?? 0;
    try {
      const poison = await flushIsolating(rows, (row) => kind.flush(target.scope, [row]), attempts >= ISOLATE_AFTER);
      for (const { item, error } of poison) {
        console.error("Counter row dead-lettered", target, item.key, error);
        sql.exec(
          "INSERT INTO dead_letters (key, a, b, c, meta, error, at) VALUES (?, ?, ?, ?, ?, ?, ?)",
          item.key,
          item.a,
          item.b,
          item.c,
          item.meta,
          error,
          Date.now(),
        );
        sql.exec("DELETE FROM counters WHERE key = ?", item.key);
      }
    } catch (error) {
      console.error("Counter flush failed", target, 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");
    // D1 now holds what was read; only rows unchanged since then are clean, anything added meanwhile stays dirty.
    for (const row of rows) {
      sql.exec("UPDATE counters SET base_a = ?, base_b = ?, base_c = ? WHERE key = ?", row.a, row.b, row.c, row.key);
      sql.exec(
        "UPDATE counters SET dirty = 0 WHERE key = ? AND a = ? AND b = ? AND c = ? AND meta IS ?",
        row.key,
        row.a,
        row.b,
        row.c,
        row.meta,
      );
    }
    const accounts = [...new Set(rows.map((row) => row.meta).filter((meta): meta is string => Boolean(meta)))];
    await kind
      .afterFlush?.(accounts)
      .catch((error: unknown) => console.error("Counter post-flush check failed", target, error));
    sql.exec("DELETE FROM counters WHERE dirty = 0 AND touched_at < ?", Date.now() - IDLE_ROW_MS);
    const pending = sql.exec<{ count: number }>("SELECT COUNT(*) AS count FROM counters WHERE dirty = 1").one().count;
    if (pending > 0) await this.ctx.storage.setAlarm(Date.now() + FLUSH_DELAY_MS);
  }
}