import { env, waitUntil } from "cloudflare:workers";
import { batch, first, stmt } from "../db";
import { creditDrawStatements } from "./credits";
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;
  const claim = await first<{ account_id: string }>(
    "SELECT account_id FROM usage_claims WHERE scratch_id = ?",
    accountId,
  );
  return claim?.account_id ?? 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,
): D1PreparedStatement {
  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 D1 batch. 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,
): D1PreparedStatement[] {
  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 D1 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);
  const statements = eventStatements(input, accountId, period, at);
  const changes =
    statements.length > 1
      ? ((await batch(statements))[0]?.meta.changes ?? 0)
      : (await statements[0].run()).meta.changes;
  const inserted = changes > 0;
  if (inserted) writeAnalytics(input, accountId, period);
  return inserted;
}

/** Whether an event with this idempotency key was already recorded (a retry of settled work). */
export async function isRecorded(idempotencyKey: string): Promise<boolean> {
  return (
    (await 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): D1PreparedStatement {
  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");
  await gaugeStatement(await resolveAccount(accountId), meter, value, at).run();
}

/** Moves anonymous scratch usage to the account that claimed it, and redirects future scratch events. */
export async function claimScratchUsage(scratchId: string, accountId: string): Promise<void> {
  if (!scratchId.startsWith("scratch:")) throw new MeteringError("Not a scratch account");
  const existing = await first<{ account_id: string }>(
    "SELECT account_id FROM usage_claims WHERE scratch_id = ?",
    scratchId,
  );
  if (existing && existing.account_id !== accountId) throw new MeteringError("Already claimed by another account");
  await batch([
    stmt(
      "INSERT OR IGNORE INTO usage_claims (scratch_id, account_id, claimed_at) VALUES (?, ?, ?)",
      scratchId,
      accountId,
      Date.now(),
    ),
    stmt(
      `INSERT INTO usage_periods (account_id, meter, period, quantity, reserved, updated_at)
       SELECT ?1, meter, period, quantity, reserved, updated_at FROM usage_periods WHERE account_id = ?2 AND true
       ON CONFLICT (account_id, meter, period) DO UPDATE SET quantity = quantity + excluded.quantity,
         reserved = reserved + excluded.reserved`,
      accountId,
      scratchId,
    ),
    stmt(
      `INSERT INTO usage_repo_periods (account_id, repo_id, meter, period, quantity)
       SELECT ?1, repo_id, meter, period, quantity FROM usage_repo_periods WHERE account_id = ?2 AND true
       ON CONFLICT (account_id, repo_id, meter, period) DO UPDATE SET quantity = quantity + excluded.quantity`,
      accountId,
      scratchId,
    ),
    stmt(
      `INSERT INTO usage_days (account_id, meter, day, quantity)
       SELECT ?1, meter, day, quantity FROM usage_days WHERE account_id = ?2 AND true
       ON CONFLICT (account_id, meter, day) DO UPDATE SET quantity = quantity + excluded.quantity`,
      accountId,
      scratchId,
    ),
    stmt(
      `INSERT INTO usage_gauges (account_id, meter, period, current, peak, updated_at)
       SELECT ?1, meter, period, current, peak, updated_at FROM usage_gauges WHERE account_id = ?2 AND true
       ON CONFLICT (account_id, meter, period) DO UPDATE SET current = current + excluded.current,
         peak = MAX(peak, current + excluded.current)`,
      accountId,
      scratchId,
    ),
    stmt("UPDATE usage_events SET account_id = ? WHERE account_id = ?", accountId, scratchId),
    stmt("UPDATE usage_reservations SET account_id = ? WHERE account_id = ?", accountId, scratchId),
    stmt("DELETE FROM usage_periods WHERE account_id = ?", scratchId),
    stmt("DELETE FROM usage_repo_periods WHERE account_id = ?", scratchId),
    stmt("DELETE FROM usage_days WHERE account_id = ?", scratchId),
    stmt("DELETE FROM usage_gauges WHERE account_id = ?", scratchId),
  ]);
}