import { env, waitUntil } from "cloudflare:workers";
import { accountDb, stmt } from "../scoped-db";
import type { ScopedStatement } from "../scoped-db";
import { locate } from "../scoped-db/locators";
import { creditDrawStatements } from "./credits";
import { claimLocator, noteScratchAccount } from "./directory";
import { randomId } from "../crypto";
import { meterDef } from "~/lib/meters";
import type { MeterId } from "~/lib/meters";

/** Billing periods are UTC calendar months, `YYYY-MM`. */
export function periodOf(now = Date.now()): string {
  return new Date(now).toISOString().slice(0, 7);
}

export type UsageInput = {
  /** The account billed: a user or org id, or a `scratch:` id for anonymous usage. */
  accountId: string;
  meter: MeterId;
  /** Exact quantity in the meter's unit; compute is per second (fractions allowed), never rounded up. */
  quantity: number;
  /** Optional sanity check against the catalog unit. */
  unit?: string;
  repoId?: string | null;
  attributes?: Record<string, string | number | boolean | null>;
  /** Stable for the logical event: a retry with the same key is ignored. Prefix with your subsystem. */
  idempotencyKey: string;
  /** Event time, defaults to now; decides the billing period. */
  at?: number;
};

const MAX_KEY_LENGTH = 200;

export class MeteringError extends Error {}

export function validate(input: UsageInput): void {
  const def = meterDef(input.meter);
  if (!def) throw new MeteringError(`Unknown meter ${input.meter}`);
  if (def.kind !== "counter") throw new MeteringError(`${input.meter} is a gauge; use setGauge`);
  if (input.unit && input.unit !== def.unit) throw new MeteringError(`${input.meter} is measured in ${def.unit}`);
  if (!Number.isFinite(input.quantity) || input.quantity < 0) throw new MeteringError("Quantity must be >= 0");
  if (!input.idempotencyKey || input.idempotencyKey.length > MAX_KEY_LENGTH)
    throw new MeteringError("Bad idempotency key");
  if (!input.accountId) throw new MeteringError("Missing account");
}

/** Anonymous usage that was claimed is redirected to the claimer from then on. */
export async function resolveAccount(accountId: string): Promise<string> {
  if (!accountId.startsWith("scratch:")) return accountId;
  return (await locate("scratch_claim", accountId)) ?? accountId;
}

export function writeAnalytics(input: UsageInput, accountId: string, period: string): void {
  const dataset = env.USAGE as AnalyticsEngineDataset | undefined;
  if (!dataset) return;
  try {
    dataset.writeDataPoint({
      indexes: [accountId],
      blobs: [input.meter, period, input.repoId ?? "", input.idempotencyKey, JSON.stringify(input.attributes ?? {})],
      doubles: [input.quantity],
    });
  } catch {
    // Analytics Engine is best-effort; D1 is the billing source of truth.
  }
}

export function eventStatement(
  input: UsageInput,
  accountId: string,
  period: string,
  at: number,
  nonce: string | null = null,
): ScopedStatement {
  return stmt(
    `INSERT OR IGNORE INTO usage_events (idempotency_key, account_id, meter, period, quantity, repo_id, attributes, created_at, nonce)
     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
    input.idempotencyKey,
    accountId,
    input.meter,
    period,
    input.quantity,
    input.repoId ?? null,
    input.attributes ? JSON.stringify(input.attributes) : null,
    at,
    nonce,
  );
}

/**
 * The event insert followed by its prepaid-credit draw, for one batch on `accountDb(accountId)`. The first
 * statement is the event; the draw only runs if that insert landed in this batch, so retries never draw.
 */
export function eventStatements(input: UsageInput, accountId: string, period: string, at: number): ScopedStatement[] {
  const nonce = randomId();
  const draws = creditDrawStatements({ ...input, accountId, period, at, eventNonce: nonce });
  return [eventStatement(input, accountId, period, at, nonce), ...draws];
}

/**
 * Records a usage event. Idempotent: the insert is keyed by `idempotencyKey` and the period
 * aggregates are maintained by triggers on that insert, so retries never double count.
 * Returns false when the key was already recorded.
 */
export async function recordUsage(input: UsageInput): Promise<boolean> {
  validate(input);
  if (input.quantity === 0) return false;
  const at = input.at ?? Date.now();
  const period = periodOf(at);
  const accountId = await resolveAccount(input.accountId);
  await noteScratchAccount(accountId);
  const statements = eventStatements(input, accountId, period, at);
  const changes = (await accountDb(accountId).batch(statements))[0]?.meta.changes ?? 0;
  const inserted = changes > 0;
  if (inserted) writeAnalytics(input, accountId, period);
  return inserted;
}

/** Whether `accountId` already recorded an event with this idempotency key (a retry of settled work). */
export async function isRecorded(accountId: string, idempotencyKey: string): Promise<boolean> {
  const resolved = await resolveAccount(accountId);
  return (
    (await accountDb(resolved).first<{ one: number }>(
      "SELECT 1 AS one FROM usage_events WHERE idempotency_key = ?",
      idempotencyKey,
    )) !== null
  );
}

/** Fire-and-forget wrapper for hot paths: metering failures must never fail the user's request. */
export function recordUsageSoon(ctx: { waitUntil(promise: Promise<unknown>): void } | null, input: UsageInput): void {
  const promise = recordUsage(input).catch((error: unknown) => console.warn("recordUsage failed", input.meter, error));
  (ctx ?? { waitUntil }).waitUntil(promise);
}

/** The last millisecond of a `YYYY-MM` period. */
export function periodEnd(period: string): number {
  const [year, month] = period.split("-").map(Number);
  return Date.UTC(year, month, 1) - 1;
}

const GAUGE_UPSERT = `INSERT INTO usage_gauges (account_id, meter, period, current, peak, updated_at) VALUES (?1, ?2, ?3, ?4, ?4, ?5)
  ON CONFLICT (account_id, meter, period) DO UPDATE SET current = excluded.current, peak = MAX(peak, excluded.current),
    updated_at = CASE WHEN current != excluded.current THEN excluded.updated_at ELSE updated_at END`;

/** Upserts a gauge level; `updated_at` only moves when the value changes, so idle accounts don't look active. */
export function gaugeStatement(accountId: string, meter: MeterId, value: number, at: number): ScopedStatement {
  return stmt(GAUGE_UPSERT, accountId, meter, periodOf(at), value, at);
}

/** Sets a storage gauge's current level; the period peak is what gets billed. */
export async function setGauge(accountId: string, meter: MeterId, value: number, at = Date.now()): Promise<void> {
  const def = meterDef(meter);
  if (!def || def.kind !== "gauge") throw new MeteringError(`${meter} is not a gauge`);
  if (!Number.isFinite(value) || value < 0) throw new MeteringError("Gauge value must be >= 0");
  const resolved = await resolveAccount(accountId);
  await noteScratchAccount(resolved);
  await accountDb(resolved).batch([gaugeStatement(resolved, meter, value, at)]);
}

type ScratchRows = {
  periods: { meter: string; period: string; quantity: number; reserved: number; updated_at: number }[];
  repoPeriods: { repo_id: string; meter: string; period: string; quantity: number }[];
  days: { meter: string; day: string; quantity: number }[];
  gauges: { meter: string; period: string; current: number; peak: number; updated_at: number }[];
  events: {
    idempotency_key: string;
    meter: string;
    period: string;
    quantity: number;
    repo_id: string | null;
    attributes: string | null;
    created_at: number;
  }[];
  reservations: {
    id: string;
    meter: string;
    period: string;
    quantity: number;
    expires_at: number;
    created_at: number;
    nonce: string | null;
  }[];
};

async function scratchRows(scratchId: string): Promise<ScratchRows> {
  const db = accountDb(scratchId);
  const [periods, repoPeriods, days, gauges, events, reservations] = await Promise.all([
    db.all<ScratchRows["periods"][number]>(
      "SELECT meter, period, quantity, reserved, updated_at FROM usage_periods WHERE account_id = ?",
      scratchId,
    ),
    db.all<ScratchRows["repoPeriods"][number]>(
      "SELECT repo_id, meter, period, quantity FROM usage_repo_periods WHERE account_id = ?",
      scratchId,
    ),
    db.all<ScratchRows["days"][number]>("SELECT meter, day, quantity FROM usage_days WHERE account_id = ?", scratchId),
    db.all<ScratchRows["gauges"][number]>(
      "SELECT meter, period, current, peak, updated_at FROM usage_gauges WHERE account_id = ?",
      scratchId,
    ),
    db.all<ScratchRows["events"][number]>(
      `SELECT idempotency_key, meter, period, quantity, repo_id, attributes, created_at FROM usage_events
       WHERE account_id = ?`,
      scratchId,
    ),
    db.all<ScratchRows["reservations"][number]>(
      `SELECT id, meter, period, quantity, expires_at, created_at, nonce FROM usage_reservations
       WHERE account_id = ?`,
      scratchId,
    ),
  ]);
  return { periods, repoPeriods, days, gauges, events, reservations };
}

const MOVED_DAY = "strftime('%Y-%m-%d', created_at / 1000, 'unixepoch')";

/**
 * The claimer's half of a scratch claim, as one batch. Moved events are tagged with `marker` so each rollup
 * adds the scratch total minus what the rollup triggers already added for the events that landed: the result
 * is exactly the scratch totals, even if an event key was somehow already present.
 */
function claimStatements(accountId: string, scratchId: string, rows: ScratchRows, marker: string): ScopedStatement[] {
  const landed = (where: string) =>
    `COALESCE((SELECT SUM(quantity) FROM usage_events WHERE account_id = ?1 AND nonce = ?2 AND ${where}), 0)`;
  return [
    stmt(
      "INSERT OR IGNORE INTO usage_claims (scratch_id, account_id, claimed_at) VALUES (?, ?, ?)",
      scratchId,
      accountId,
      Date.now(),
    ),
    ...rows.events.map((event) =>
      stmt(
        `INSERT OR IGNORE INTO usage_events (idempotency_key, account_id, meter, period, quantity, repo_id, attributes, created_at, nonce)
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
        event.idempotency_key,
        accountId,
        event.meter,
        event.period,
        event.quantity,
        event.repo_id,
        event.attributes,
        event.created_at,
        marker,
      ),
    ),
    ...rows.periods.map((row) =>
      stmt(
        `INSERT INTO usage_periods (account_id, meter, period, quantity, reserved, updated_at)
         VALUES (?1, ?3, ?4, ?5 - ${landed("meter = ?3 AND period = ?4")}, ?6, ?7)
         ON CONFLICT (account_id, meter, period) DO UPDATE SET quantity = quantity + excluded.quantity,
           reserved = reserved + excluded.reserved`,
        accountId,
        marker,
        row.meter,
        row.period,
        row.quantity,
        row.reserved,
        row.updated_at,
      ),
    ),
    ...rows.repoPeriods.map((row) =>
      stmt(
        `INSERT INTO usage_repo_periods (account_id, repo_id, meter, period, quantity)
         VALUES (?1, ?3, ?4, ?5, ?6 - ${landed("repo_id = ?3 AND meter = ?4 AND period = ?5")})
         ON CONFLICT (account_id, repo_id, meter, period) DO UPDATE SET quantity = quantity + excluded.quantity`,
        accountId,
        marker,
        row.repo_id,
        row.meter,
        row.period,
        row.quantity,
      ),
    ),
    ...rows.days.map((row) =>
      stmt(
        `INSERT INTO usage_days (account_id, meter, day, quantity)
         VALUES (?1, ?3, ?4, ?5 - ${landed(`meter = ?3 AND ${MOVED_DAY} = ?4`)})
         ON CONFLICT (account_id, meter, day) DO UPDATE SET quantity = quantity + excluded.quantity`,
        accountId,
        marker,
        row.meter,
        row.day,
        row.quantity,
      ),
    ),
    ...rows.gauges.map((row) =>
      stmt(
        `INSERT INTO usage_gauges (account_id, meter, period, current, peak, updated_at) VALUES (?, ?, ?, ?, ?, ?)
         ON CONFLICT (account_id, meter, period) DO UPDATE SET current = current + excluded.current,
           peak = MAX(peak, current + excluded.current)`,
        accountId,
        row.meter,
        row.period,
        row.current,
        row.peak,
        row.updated_at,
      ),
    ),
    ...rows.reservations.map((row) =>
      stmt(
        `INSERT OR IGNORE INTO usage_reservations (id, account_id, meter, period, quantity, expires_at, created_at, nonce)
         VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
        row.id,
        accountId,
        row.meter,
        row.period,
        row.quantity,
        row.expires_at,
        row.created_at,
        row.nonce,
      ),
    ),
  ];
}

/**
 * Moves anonymous scratch usage to the account that claimed it, and redirects future scratch events. The claim
 * is registered first (so new events already go to the claimer), then the claimer's batch lands, then the
 * scratch object is emptied; a crash before the last step would move the rows again on retry.
 */
export async function claimScratchUsage(scratchId: string, accountId: string): Promise<void> {
  if (!scratchId.startsWith("scratch:")) throw new MeteringError("Not a scratch account");
  const existing = await locate("scratch_claim", scratchId);
  if (existing && existing !== accountId) throw new MeteringError("Already claimed by another account");
  if ((await claimLocator("scratch_claim", scratchId, accountId)) !== accountId) {
    throw new MeteringError("Already claimed by another account");
  }
  const rows = await scratchRows(scratchId);
  await accountDb(accountId).batch(claimStatements(accountId, scratchId, rows, `claim:${randomId()}`));
  await accountDb(scratchId).batch(
    ["usage_events", "usage_reservations", "usage_periods", "usage_repo_periods", "usage_days", "usage_gauges"].map(
      (table) => stmt(`DELETE FROM ${table} WHERE account_id = ?`, scratchId),
    ),
  );
}