import { DurableObject } from "cloudflare:workers";
import { batch, first, stmt as d1Stmt } from "./db";
import { backoffMs, ISOLATE_AFTER } from "./do-flush";
import { accountDb, repoDb, stmt, type ScopedDb, type ScopedStatement } from "./scoped-db";
import { locate } from "./scoped-db/locators";
import { afterHostingFlush } from "./sites/caps";
import { REPO_PACK_BYTES, REPO_STARS } from "./counter-keys";
import { forwardFlushed } from "./counters-forward";

/** 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"
  | "platform-ai"
  | "platform-sites";
export type Totals = [number, number, number];
type Row = { key: string; a: number; b: number; c: number; meta: string | null };

/**
 * Pending writes of one flush, grouped by store: D1 directory statements (star counts) and statements for one
 * repository's or account's object. Each group is its own transaction.
 */
type Writes = { d1: D1PreparedStatement[]; scoped: Array<{ db: ScopedDb; statements: ScopedStatement[] }> };

async function count(read: Promise<{ value: number } | null>): Promise<number> {
  return (await read)?.value ?? 0;
}

async function repoExists(repoId: string): Promise<boolean> {
  return (await first<{ ok: number }>("SELECT 1 AS ok FROM repos WHERE id = ?", repoId)) !== null;
}

/**
 * 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`.
 */
async function repoCountWrites(repoId: string, rows: Row[]): Promise<Writes> {
  const d1: D1PreparedStatement[] = [];
  const statements: ScopedStatement[] = [];
  for (const row of rows) {
    if (row.key === REPO_STARS) {
      d1.push(
        d1Stmt("UPDATE repos SET stars_count = (SELECT COUNT(*) FROM stars WHERE repo_id = ?1) WHERE id = ?1", repoId),
      );
    } else if (row.key === REPO_PACK_BYTES) {
      statements.push(
        stmt(
          `INSERT INTO repo_storage_estimates (repo_id, pack_bytes, updated_at) VALUES (?1, ?2, ?3)
           ON CONFLICT (repo_id) DO UPDATE SET pack_bytes = excluded.pack_bytes, updated_at = excluded.updated_at`,
          repoId,
          row.a,
          Date.now(),
        ),
      );
    } else {
      const table = row.key.startsWith("d:") ? "discussions" : "discussion_comments";
      statements.push(
        stmt(
          `UPDATE ${table} SET upvotes = (SELECT COUNT(*) FROM discussion_votes WHERE subject_id = ?1) WHERE id = ?1`,
          row.key.slice(2),
        ),
      );
    }
  }
  const live = statements.length > 0 && (await repoExists(repoId));
  return { d1, scoped: live ? [{ db: repoDb(repoId), statements }] : [] };
}

type Kind = {
  /** Current stored 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[]): Promise<Writes>;
  /** 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;
  /**
   * A running total elsewhere that each flushed delta is added to (best effort), so platform-wide sums read one
   * object instead of every scope.
   */
  forward?: { kind: CounterKind; scope: string };
};

const NO_WRITES: Writes = { d1: [], scoped: [] };

/** Sites live in their repository's RepoState; a site no longer registered is gone and its usage is dropped. */
async function siteDb(siteId: string): Promise<ScopedDb | null> {
  const repoId = await locate("site", siteId);
  return repoId ? repoDb(repoId) : null;
}

const KINDS: Record<CounterKind, Kind> = {
  /** scope = site id, key = UTC day, totals = requests / cpu ms / bytes out, meta = billed account id. */
  "site-usage": {
    forward: { kind: "platform-sites", scope: "global" },
    async seed(siteId, day) {
      const db = await siteDb(siteId);
      const row = db
        ? await db.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,
          )
        : null;
      return [row?.requests ?? 0, row?.cpu_ms ?? 0, row?.bytes_out ?? 0];
    },
    async flush(siteId, rows) {
      const db = await siteDb(siteId);
      if (!db) return NO_WRITES;
      const statements = 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,
        ),
      );
      return { d1: [], scoped: [{ db, statements }] };
    },
    afterFlush: afterHostingFlush,
  },
  /** scope = repo id, key = release asset id, totals[0] = downloads. */
  "release-downloads": {
    async seed(repoId, assetId) {
      const row = await repoDb(repoId).first<{ download_count: number }>(
        "SELECT download_count FROM release_assets WHERE id = ?",
        assetId,
      );
      return [row?.download_count ?? 0, 0, 0];
    },
    async flush(repoId, rows) {
      if (!(await repoExists(repoId))) return NO_WRITES;
      return {
        d1: [],
        scoped: [
          {
            db: repoDb(repoId),
            statements: 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 accountDb(accountId).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];
    },
    async flush(accountId, rows) {
      // A deleted account's object was wiped; writing would bring it back.
      if (!(await first<{ ok: number }>("SELECT 1 AS ok FROM accounts WHERE id = ?", accountId))) return NO_WRITES;
      const statements = rows.map((row) => {
        const [day, tool] = splitUsageKey(row.key);
        return stmt(
          `INSERT INTO mcp_usage_daily (account_id, day, tool, calls, errors) VALUES (?1, ?2, ?3, ?4, ?5)
           ON CONFLICT (account_id, day, tool) DO UPDATE SET calls = excluded.calls, errors = excluded.errors`,
          accountId,
          day,
          tool,
          row.a,
          row.b,
        );
      });
      return { d1: [], scoped: [{ db: accountDb(accountId), statements }] };
    },
  },
  /**
   * scope = "global", key = UTC day, totals = platform AI cost in microcents. Lives only in this object: it is
   * the running total the daily AI ceiling checks, so no call sums usage over every account.
   */
  "platform-ai": {
    async seed() {
      return [0, 0, 0];
    },
    async flush() {
      return NO_WRITES;
    },
  },
  /** scope = "global", key = UTC day, totals = requests / cpu ms / bytes out of every site, fed by site-usage. */
  "platform-sites": {
    async seed() {
      return [0, 0, 0];
    },
    async flush() {
      return NO_WRITES;
    },
  },
  /** 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(first("SELECT COUNT(*) AS value FROM stars WHERE repo_id = ?", repoId)), 0, 0];
      const db = repoDb(repoId);
      if (key === REPO_PACK_BYTES)
        return [
          await count(db.first("SELECT pack_bytes AS value FROM repo_storage_estimates WHERE repo_id = ?", repoId)),
          0,
          0,
        ];
      return [
        await count(db.first("SELECT COUNT(*) AS value FROM discussion_votes WHERE subject_id = ?", key.slice(2))),
        0,
        0,
      ];
    },
    flush: repoCountWrites,
    seedIncludesDelta: (key) => key !== REPO_PACK_BYTES,
  },
};

async function applyWrites(writes: Writes): Promise<void> {
  if (writes.d1.length > 0) await batch(writes.d1);
  for (const group of writes.scoped) await group.db.batch(group.statements);
}

/**
 * Flushes rows together. Once `isolate` is set (the flush has failed repeatedly), rows are retried one by one
 * and those that still fail are returned as poison so the caller can dead-letter them instead of blocking
 * every later flush. If several rows all fail individually the store itself is the problem, so it throws.
 */
async function flushIsolating(
  rows: Row[],
  write: (rows: Row[]) => Promise<void>,
  isolate: boolean,
): Promise<Array<{ item: Row; error: string }>> {
  if (rows.length === 0) return [];
  try {
    await write(rows);
    return [];
  } catch (error) {
    if (!isolate) throw error;
  }
  const poison: Array<{ item: Row; error: string }> = [];
  for (const row of rows) {
    try {
      await write([row]);
    } catch (error) {
      poison.push({ item: row, error: String(error) });
    }
  }
  if (rows.length > 1 && poison.length === rows.length) throw new Error(`All ${rows.length} rows failed to flush`);
  return poison;
}

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,
        async (group) => applyWrites(await kind.flush(target.scope, group)),
        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");
    if (kind.forward) await forwardFlushed(this.ctx.storage.sql, this.env.COUNTERS, kind.forward, rows);
    // 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);
  }
}