import { env } from "cloudflare:workers";
import { accountDb, stmt } from "../scoped-db";
import type { ScopedStatement } from "../scoped-db";
import { locate, registerLocator } from "../scoped-db/locators";
import { randomId } from "../crypto";
import type { MeterId } from "~/lib/meters";
import { parsePlans } from "~/lib/plans";
import { CREDIT_POOLS, METER_PRICES, creditPoolFor } from "~/lib/pricing";
import type { CreditPoolId } from "~/lib/pricing";

/**
 * Prepaid credits: a local ledger per account and credit pool (see CREDIT_POOLS in src/lib/pricing.ts).
 * Compute credit is held in standard-1 runner seconds and spent by every priced compute meter at its price
 * relative to standard-1; AI token credit is spent 1:1 by ai.tokens.
 *
 * Draw-down order for a creditable meter is plan-included first, then credits, then overage (paid plans) or
 * the hard cap (free). Draws happen in the same account batch that inserts the usage event, keyed by the event's
 * idempotency key, so a retried event never draws twice. reserveUsage counts the unspent balance as extra
 * headroom inside its own atomic batch, so concurrent holds can never commit more than plan + credits.
 */
/** Whether usage on `meter` can be paid for with prepaid credit. */
export function isCreditMeter(meter: string): meter is MeterId {
  return creditPoolFor(meter) !== null;
}

export function isCreditPool(pool: string): pool is CreditPoolId {
  return Object.hasOwn(CREDIT_POOLS, pool);
}

/** Older ledgers and pack configs named the compute pool after its base meter. */
export function normalizePool(value: string): CreditPoolId | null {
  if (isCreditPool(value)) return value;
  return value.startsWith("compute.") && isCreditMeter(value) ? "compute" : null;
}

export class CreditError extends Error {}

/** Balances below this are rounding dust and count as empty. */
export const CREDIT_EPSILON = 1e-6;

/** Every meter that spends from `pool`, with its factor (pool units per meter unit). */
export function poolMeters(pool: CreditPoolId): { meter: string; factor: number }[] {
  return Object.keys(METER_PRICES).flatMap((meter) => {
    const credit = creditPoolFor(meter);
    return credit?.pool === pool ? [{ meter, factor: credit.factor }] : [];
  });
}

/**
 * SQL for the pool units already claimed by in-flight holds on the pool's *other* meters (account ?A, meter ?M
 * excluded). Holds are counted in full, which is conservative: some may land inside a plan's included amount.
 * Subtracting this keeps concurrent holds on different runner sizes from all counting the same credit.
 */
export function otherHoldsSql(
  pool: CreditPoolId,
  accountParam: string,
  meterParam: string,
  period: string,
  now = Date.now(),
): string {
  if (!/^\d{4}-\d{2}$/.test(period)) throw new CreditError("Bad period");
  const cases = poolMeters(pool)
    .map(({ meter, factor }) => `WHEN '${meter}' THEN ${factor}`)
    .join(" ");
  // Only live holds of this period: expired ones are waiting for the sweep and will never commit.
  return `COALESCE((SELECT SUM(r.quantity * CASE r.meter ${cases} ELSE 0 END) FROM usage_reservations r
    WHERE r.account_id = ${accountParam} AND r.meter != ${meterParam} AND r.period = '${period}'
      AND r.expires_at > ${Math.floor(now)}), 0)`;
}

/** Unspent credit in a pool, in pool units. */
export async function creditBalance(accountId: string, pool: CreditPoolId): Promise<number> {
  const row = await accountDb(accountId).first<{ balance: number }>(
    "SELECT balance FROM credit_balances WHERE account_id = ? AND meter = ?",
    accountId,
    pool,
  );
  return row?.balance ?? 0;
}

export async function creditBalances(accountId: string): Promise<Record<string, number>> {
  const rows = await accountDb(accountId).all<{ meter: string; balance: number }>(
    `SELECT meter, balance FROM credit_balances WHERE account_id = ? AND balance > ${CREDIT_EPSILON}`,
    accountId,
  );
  return Object.fromEntries(rows.map((row) => [row.meter, row.balance]));
}

/**
 * Statements to append after a usage event insert (same batch) that draw credit for the part of the event
 * above the plan's included amount and not already covered. The usage_periods trigger has already added the
 * event quantity when these run. Everything, including the balance and the account's plan, is read inside
 * the batch, so a purchase landing concurrently is never missed and no extra round trip is made.
 */
export function creditDrawStatements(input: {
  idempotencyKey: string;
  accountId: string;
  meter: MeterId;
  period: string;
  quantity: number;
  at: number;
  /** The nonce written by this batch's event insert. */
  eventNonce: string;
}): ScopedStatement[] {
  const credit = creditPoolFor(input.meter);
  if (!credit || input.quantity <= 0) return [];
  const plans = parsePlans(env.METERING_PLANS);
  const scratch = input.accountId.startsWith("scratch:");
  const freeIncluded = (scratch ? plans.scratch : plans.free).included[input.meter] ?? 0;
  const proIncluded = scratch ? freeIncluded : (plans.pro.included[input.meter] ?? 0);
  const nonce = randomId();
  const drawn = "(SELECT quantity FROM credit_draws WHERE idempotency_key = ?1 AND nonce = ?2)";
  return [
    stmt(
      `INSERT OR IGNORE INTO credit_draws (idempotency_key, account_id, meter, period, quantity, nonce, created_at)
       SELECT ?1, ?2, ?3, ?4, d.amount, ?5, ?6 FROM (
         SELECT MIN(b.balance / ?11, MIN(?7, MAX(0, p.quantity - p.credit_covered
           - CASE (SELECT plan FROM account_billing WHERE account_id = ?2) WHEN 'pro' THEN ?10 ELSE ?8 END))) AS amount
         FROM credit_balances b JOIN usage_periods p ON p.account_id = b.account_id AND p.meter = ?3 AND p.period = ?4
         WHERE b.account_id = ?2 AND b.meter = ?12 AND b.balance > ${CREDIT_EPSILON}
           AND EXISTS (SELECT 1 FROM usage_events WHERE idempotency_key = ?1 AND nonce = ?9)
       ) d WHERE d.amount > 0`,
      input.idempotencyKey,
      input.accountId,
      input.meter,
      input.period,
      nonce,
      input.at,
      input.quantity,
      freeIncluded,
      input.eventNonce,
      proIncluded,
      credit.factor,
      credit.pool,
    ),
    stmt(
      `UPDATE credit_balances
       SET balance = CASE WHEN balance - ${drawn} * ?6 < ${CREDIT_EPSILON} THEN 0 ELSE balance - ${drawn} * ?6 END, updated_at = ?5
       WHERE account_id = ?3 AND meter = ?4 AND ${drawn} IS NOT NULL`,
      input.idempotencyKey,
      nonce,
      input.accountId,
      credit.pool,
      input.at,
      credit.factor,
    ),
    stmt(
      `UPDATE usage_periods SET credit_covered = credit_covered + ${drawn}
       WHERE account_id = ?3 AND meter = ?4 AND period = ?5 AND ${drawn} IS NOT NULL`,
      input.idempotencyKey,
      nonce,
      input.accountId,
      input.meter,
      input.period,
    ),
  ];
}

export type CreditMovement = {
  ref: string;
  accountId: string;
  /** The credit pool (CREDIT_POOLS); the column is still called meter. */
  meter: CreditPoolId;
  kind: "purchase" | "refund" | "dispute" | "adjust";
  /** Positive to add credit. Reversals pass a positive amount to remove; it is clamped to the unspent balance. */
  quantity: number;
  /** Reversals: the purchase being reversed and the unclamped amount owed back for it. */
  purchaseRef?: string | null;
  requested?: number | null;
  pack?: string | null;
  paymentIntent?: string | null;
  amountCents?: number | null;
};

/**
 * Applies one ledger movement exactly once (keyed by `ref`). Purchases and positive adjustments add the full
 * quantity; refunds and disputes remove at most the unspent balance, since spent credit can't be clawed back.
 * Returns the quantity actually applied (signed), or 0 when `ref` was already applied.
 */
export async function applyCreditMovement(movement: CreditMovement): Promise<number> {
  if (!isCreditPool(movement.meter)) throw new CreditError(`${movement.meter} is not a credit pool`);
  if (!Number.isFinite(movement.quantity) || movement.quantity < 0) throw new CreditError("Bad credit quantity");
  const removes = movement.kind === "refund" || movement.kind === "dispute";
  const nonce = randomId();
  const now = Date.now();
  const delta = removes
    ? "-MIN(?5, COALESCE((SELECT balance FROM credit_balances WHERE account_id = ?2 AND meter = ?3), 0))"
    : "?5";
  const applied = "(SELECT quantity FROM credit_ledger WHERE ref = ?1 AND nonce = ?2)";
  const db = accountDb(movement.accountId);
  await db.batch([
    stmt(
      "INSERT OR IGNORE INTO credit_balances (account_id, meter, balance, updated_at) VALUES (?, ?, 0, ?)",
      movement.accountId,
      movement.meter,
      now,
    ),
    stmt(
      `INSERT OR IGNORE INTO credit_ledger (ref, account_id, meter, kind, quantity, pack, payment_intent, amount_cents,
         nonce, created_at, requested, purchase_ref)
       VALUES (?1, ?2, ?3, ?4, ${delta}, ?6, ?7, ?8, ?9, ?10, ?11, ?12)`,
      movement.ref,
      movement.accountId,
      movement.meter,
      movement.kind,
      movement.quantity,
      movement.pack ?? null,
      movement.paymentIntent ?? null,
      movement.amountCents ?? null,
      nonce,
      now,
      movement.requested ?? null,
      movement.purchaseRef ?? null,
    ),
    stmt(
      `UPDATE credit_balances SET balance = MAX(0, balance + ${applied}), updated_at = ?5
       WHERE account_id = ?3 AND meter = ?4 AND ${applied} IS NOT NULL`,
      movement.ref,
      nonce,
      movement.accountId,
      movement.meter,
      now,
    ),
  ]);
  const row = await db.first<{ quantity: number; nonce: string }>(
    "SELECT quantity, nonce FROM credit_ledger WHERE ref = ?",
    movement.ref,
  );
  // Refunds and disputes name only the payment intent; this finds the account that bought with it.
  if (movement.kind === "purchase" && movement.paymentIntent) {
    await registerLocator("credit_payment", movement.paymentIntent, movement.accountId);
  }
  return row && row.nonce === nonce ? row.quantity : 0;
}

export type CreditPurchase = {
  ref: string;
  accountId: string;
  meter: CreditPoolId;
  quantity: number;
  /** Sum of amounts already owed back by earlier reversals of this purchase (before clamping). */
  reversed: number;
  /** What is left of this purchase, assuming draws spend the oldest credit first. */
  unspent: number;
};

type PurchaseRow = { ref: string; account_id: string; meter: CreditPoolId; quantity: number; created_at: number };

async function reversedFor(accountId: string, refs: string[]): Promise<Map<string, number>> {
  const out = new Map<string, number>();
  if (refs.length === 0) return out;
  const rows = await accountDb(accountId).all<{ purchase_ref: string; requested: number }>(
    `SELECT purchase_ref, COALESCE(SUM(requested), 0) AS requested FROM credit_ledger
     WHERE kind IN ('refund', 'dispute', 'adjust') AND purchase_ref IN (${refs.map(() => "?").join(", ")}) GROUP BY purchase_ref`,
    ...refs,
  );
  for (const row of rows) out.set(row.purchase_ref, row.requested);
  return out;
}

/**
 * Purchases paid through `paymentIntent`, each with how much of it is still unspent. Draws spend the oldest
 * credit first, so a pack's unspent share is what the balance holds after every newer pack keeps its own.
 * Reversing one pack can therefore never take credit that was bought with another.
 */
export async function purchasesForPayment(paymentIntent: string): Promise<CreditPurchase[]> {
  const accountId = await locate("credit_payment", paymentIntent);
  if (!accountId) return [];
  const db = accountDb(accountId);
  const targets = await db.all<PurchaseRow>(
    "SELECT ref, account_id, meter, quantity, created_at FROM credit_ledger WHERE payment_intent = ? AND kind = 'purchase'",
    paymentIntent,
  );
  const out: CreditPurchase[] = [];
  for (const target of targets) {
    const siblings = await db.all<PurchaseRow>(
      `SELECT ref, account_id, meter, quantity, created_at FROM credit_ledger
       WHERE account_id = ? AND meter = ? AND kind = 'purchase' ORDER BY created_at DESC, ref DESC`,
      target.account_id,
      target.meter,
    );
    const reversed = await reversedFor(
      accountId,
      siblings.map((row) => row.ref),
    );
    let pool = await creditBalance(target.account_id, target.meter);
    let unspent = 0;
    for (const row of siblings) {
      const owed = Math.max(0, row.quantity - (reversed.get(row.ref) ?? 0));
      const share = Math.min(owed, pool);
      pool -= share;
      if (row.ref === target.ref) unspent = share;
    }
    out.push({
      ref: target.ref,
      accountId: target.account_id,
      meter: target.meter,
      quantity: target.quantity,
      reversed: reversed.get(target.ref) ?? 0,
      unspent,
    });
  }
  return out;
}

export type LedgerEntry = {
  ref: string;
  meter: string;
  kind: string;
  quantity: number;
  pack: string | null;
  createdAt: number;
};

export async function creditLedger(accountId: string, limit = 50): Promise<LedgerEntry[]> {
  const rows = await accountDb(accountId).all<{
    ref: string;
    meter: string;
    kind: string;
    quantity: number;
    pack: string | null;
    created_at: number;
  }>(
    "SELECT ref, meter, kind, quantity, pack, created_at FROM credit_ledger WHERE account_id = ? ORDER BY created_at DESC LIMIT ?",
    accountId,
    limit,
  );
  return rows.map((row) => ({
    ref: row.ref,
    meter: row.meter,
    kind: row.kind,
    quantity: row.quantity,
    pack: row.pack,
    createdAt: row.created_at,
  }));
}