import { first } from "../db";
import { backfillOnce, backfillRepos } from "../backfill";
import { locatedScopes } from "../scoped-db/locators";
import { accountDb, repoDb, stmt } from "../scoped-db";
import type { ScopedStatement } from "../scoped-db";
import { noteScratchAccount } from "./directory";
import { gaugeStatement, resolveAccount } from "./record";
import type { MeterId } from "~/lib/meters";

/**
 * Storage gauges without a cron over every repository: each repository or site scope (RepoState) computes what
 * it stores per billed account and writes it to that account's `scope_gauges`; the account's gauge for a meter
 * is the sum over its scopes. A scope recomputes after a debounce when its storage tables change (RepoState
 * alarm task `gauges`), and an account re-sums when its `scope_gauges` change (AccountState task `scope-gauges`).
 */
export const SCOPE_METERS: MeterId[] = [
  "storage.ci_cache",
  "storage.artifacts",
  "storage.r2",
  "storage.sites",
  "domains.custom",
];

type GaugeRow = { accountId: string; meter: MeterId; value: number };

async function scopeOwner(scopeId: string): Promise<string | null> {
  if (scopeId.startsWith("site:")) return null;
  return (await first<{ owner_id: string }>("SELECT owner_id FROM repos WHERE id = ?", scopeId))?.owner_id ?? null;
}

/** What this scope stores per billed account (claimed scratch ids resolved to their claimer). */
export async function computeScopeGauges(scopeId: string): Promise<GaugeRow[]> {
  const db = repoDb(scopeId);
  const raw: GaugeRow[] = [];
  const owner = await scopeOwner(scopeId);
  if (owner) {
    const [cache, artifacts, stored] = await Promise.all([
      // Cache objects are content-addressed within the repository: bill each stored object once.
      db.first<{ value: number | null }>(
        "SELECT SUM(size) AS value FROM (SELECT DISTINCT r2_key, size FROM action_caches WHERE size > 0)",
      ),
      db.first<{ value: number | null }>("SELECT SUM(pack_bytes) AS value FROM repo_storage_estimates"),
      db.first<{ value: number | null }>(
        `SELECT SUM(size) AS value FROM (
           SELECT a.size FROM release_assets a JOIN releases l ON l.id = a.release_id
           UNION ALL SELECT size FROM attachments)`,
      ),
    ]);
    raw.push(
      { accountId: owner, meter: "storage.ci_cache", value: cache?.value ?? 0 },
      { accountId: owner, meter: "storage.artifacts", value: artifacts?.value ?? 0 },
      { accountId: owner, meter: "storage.r2", value: stored?.value ?? 0 },
    );
  }
  const [sites, domains] = await Promise.all([
    db.all<{ account_id: string; value: number | null }>(
      `SELECT COALESCE(s.owner_id, 'scratch:site:' || s.id) AS account_id, SUM(d.total_bytes) AS value
       FROM sites s JOIN site_deployments d ON d.site_id = s.id WHERE d.status = 'active'
       GROUP BY COALESCE(s.owner_id, 'scratch:site:' || s.id)`,
    ),
    db.all<{ account_id: string; value: number }>(
      `SELECT owner_id AS account_id, COUNT(*) AS value FROM sites
       WHERE owner_id IS NOT NULL AND custom_domain IS NOT NULL GROUP BY owner_id`,
    ),
  ]);
  for (const row of sites) raw.push({ accountId: row.account_id, meter: "storage.sites", value: row.value ?? 0 });
  for (const row of domains) raw.push({ accountId: row.account_id, meter: "domains.custom", value: row.value });
  const merged = new Map<string, GaugeRow>();
  for (const row of raw) {
    const accountId = await resolveAccount(row.accountId);
    const key = `${accountId}\u0000${row.meter}`;
    const current = merged.get(key);
    merged.set(key, { accountId, meter: row.meter, value: (current?.value ?? 0) + row.value });
  }
  return [...merged.values()];
}

/**
 * Writes the scope's current storage to every account it bills, and clears it from accounts it no longer bills
 * (a transfer, a site changing hands). `gauge_accounts` remembers who was billed last time.
 */
export async function publishScopeGauges(scopeId: string): Promise<void> {
  const rows = await computeScopeGauges(scopeId);
  const db = repoDb(scopeId);
  const previous = (await db.all<{ account_id: string }>("SELECT account_id FROM gauge_accounts")).map(
    (row) => row.account_id,
  );
  const current = new Map<string, GaugeRow[]>();
  for (const row of rows) current.set(row.accountId, [...(current.get(row.accountId) ?? []), row]);
  for (const accountId of new Set([...previous, ...current.keys()])) {
    await noteScratchAccount(accountId);
    const statements: ScopedStatement[] = [stmt("DELETE FROM scope_gauges WHERE scope_id = ?", scopeId)];
    for (const row of current.get(accountId) ?? []) {
      statements.push(
        stmt("INSERT INTO scope_gauges (scope_id, meter, value) VALUES (?, ?, ?)", scopeId, row.meter, row.value),
      );
    }
    await accountDb(accountId).batch(statements);
  }
  await db.batch([
    stmt("DELETE FROM gauge_accounts"),
    ...[...current.keys()].map((accountId) => stmt("INSERT INTO gauge_accounts (account_id) VALUES (?)", accountId)),
  ]);
}

/** Before a scope's object is destroyed: its storage stops counting for everyone it billed. */
export async function retractScopeGauges(scopeId: string): Promise<void> {
  const accounts = await repoDb(scopeId).all<{ account_id: string }>("SELECT account_id FROM gauge_accounts");
  await Promise.all(
    accounts.map((row) => accountDb(row.account_id).run("DELETE FROM scope_gauges WHERE scope_id = ?", scopeId)),
  );
}

type Sql = {
  exec<T extends Record<string, string | number | null>>(
    sql: string,
    ...params: Array<string | number | null>
  ): { toArray(): T[] };
};

/** Inside the account's object: its scope gauges summed into its billed gauges. */
export function sumScopeGauges(sql: Sql, accountId: string, now: number): void {
  const sums = new Map(
    sql
      .exec<{ meter: string; value: number | null }>(
        "SELECT meter, SUM(value) AS value FROM scope_gauges GROUP BY meter",
      )
      .toArray()
      .map((row) => [row.meter, row.value ?? 0]),
  );
  for (const meter of SCOPE_METERS) {
    const value = sums.get(meter) ?? 0;
    // Zero is only written over an existing gauge, so accounts that never stored anything get no rows.
    if (
      value === 0 &&
      sql.exec("SELECT 1 AS one FROM usage_gauges WHERE meter = ? LIMIT 1", meter).toArray().length === 0
    )
      continue;
    const statement = gaugeStatement(accountId, meter, value, now);
    sql.exec(statement.sql, ...statement.params);
  }
}

/**
 * Cron until done: publishes every repository's gauges once (later ones follow writes), then the repo-less site
 * scopes (`site:<id>`), which have no directory row to walk.
 */
export async function backfillScopeGauges(): Promise<void> {
  if (!(await backfillRepos("scope-gauges", publishScopeGauges))) return;
  await backfillOnce("site-scope-gauges", async () => {
    for (const scope of await locatedScopes("site")) if (scope.startsWith("site:")) await publishScopeGauges(scope);
  });
}