import { AsyncLocalStorage } from "node:async_hooks";

/**
 * Per-unit-of-work cost accounting for hotspot detection. Every request, queue batch, cron task and Durable
 * Object alarm runs inside `withPerf`, which counts the D1 statements and scoped-object (RepoState /
 * AccountState) calls it makes and how long they took, then reports one PostHog event:
 *
 * - `server_request`: `route` (a template, never names; lib/route-template.ts), `method`, `status`, `surface`.
 * - `server_job`: `job` (cron task, queue, alarm name).
 *
 * Both carry `duration_ms`, `d1_queries`, `d1_ms`, `d1_replica_queries`, `d1_rows_read`, `d1_rows_written`,
 * `do_calls`, `do_ms`, `do_rows_read`, `do_rows_written`, `cold_start` and `sample_weight` (rows are what D1 and
 * Durable Object SQL bill by). Slow or expensive units are always sent; the rest are sampled (PERF_SAMPLE_RATE, default 2%)
 * and carry `sample_weight` = 1 / rate so sums and counts in PostHog estimate the true totals. No person
 * profile, no account ids, no paths.
 */
export type PerfCounters = {
  d1Queries: number;
  d1Ms: number;
  d1ReplicaQueries: number;
  doCalls: number;
  doMs: number;
  doRowsRead: number;
  doRowsWritten: number;
  d1RowsRead: number;
  d1RowsWritten: number;
};

type PerfScope = PerfCounters & { startedAt: number; coldStart: boolean };

const storage = new AsyncLocalStorage<PerfScope>();

const SLOW_MS = 1000;
const HEAVY_D1 = 25;
const HEAVY_DO = 25;
const DEFAULT_SAMPLE_RATE = 0.02;

let coldStart = true;

type Reporter = (event: string, properties: Record<string, string | number | boolean | null>) => void;

let rate = DEFAULT_SAMPLE_RATE;
let reporter: Reporter | null = null;

/**
 * Wired by the Worker entry (in every isolate, Durable Objects included, since they load the same entry). This
 * module stays free of `cloudflare:workers` and analytics imports because db.ts and every scoped store import it.
 */
export function setPerfReporter(report: Reporter): void {
  reporter = report;
}

/** PERF_SAMPLE_RATE: the share of cheap units reported (0 < rate <= 1). */
export function configurePerf(sampleRate: string | undefined): void {
  const raw = Number(sampleRate);
  rate = Number.isFinite(raw) && raw > 0 && raw <= 1 ? raw : DEFAULT_SAMPLE_RATE;
}

/** Records one D1 statement (or batch) of the current unit of work. */
export function countD1(ms: number, replica = false): void {
  const scope = storage.getStore();
  if (!scope) return;
  if (replica) scope.d1ReplicaQueries++;
  else scope.d1Queries++;
  scope.d1Ms += ms;
}

/** Records one call into a RepoState or AccountState. */
export function countDo(ms: number): void {
  const scope = storage.getStore();
  if (!scope) return;
  scope.doCalls++;
  scope.doMs += ms;
}

/** Rows a scoped-object statement read and wrote (Durable Object SQL is billed by rows). */
export function countDoRows(read: number, written: number): void {
  const scope = storage.getStore();
  if (!scope) return;
  scope.doRowsRead += read;
  scope.doRowsWritten += written;
}

/** Rows a D1 statement read and wrote, from its result meta. */
export function countD1Rows(read: number, written: number): void {
  const scope = storage.getStore();
  if (!scope) return;
  scope.d1RowsRead += read;
  scope.d1RowsWritten += written;
}

/** Times `work` into `count`, whether it resolves or throws. */
export async function timed<T>(count: (ms: number) => void, work: () => Promise<T>): Promise<T> {
  const start = performance.now();
  try {
    return await work();
  } finally {
    count(performance.now() - start);
  }
}

type Outcome = {
  event: "server_request" | "server_job";
  properties: Record<string, string | number | boolean | null>;
  /** The work completed but failed (a request answered 5xx). */
  failed?: boolean;
};

function report(scope: PerfScope, outcome: Outcome, threw: boolean): void {
  const failed = threw || outcome.failed === true;
  const duration = Date.now() - scope.startedAt;
  const expensive =
    failed || duration >= SLOW_MS || scope.d1Queries + scope.d1ReplicaQueries >= HEAVY_D1 || scope.doCalls >= HEAVY_DO;
  if (!expensive && Math.random() >= rate) return;
  const properties = {
    ...outcome.properties,
    duration_ms: duration,
    d1_queries: scope.d1Queries,
    d1_ms: Math.round(scope.d1Ms),
    d1_replica_queries: scope.d1ReplicaQueries,
    d1_rows_read: scope.d1RowsRead,
    d1_rows_written: scope.d1RowsWritten,
    do_calls: scope.doCalls,
    do_ms: Math.round(scope.doMs),
    do_rows_read: scope.doRowsRead,
    do_rows_written: scope.doRowsWritten,
    failed,
    expensive,
    cold_start: scope.coldStart,
    // Expensive units are always sent (weight 1); sampled ones stand for 1 / rate units.
    sample_weight: expensive ? 1 : Math.round(1 / rate),
  };
  try {
    reporter?.(outcome.event, properties);
  } catch {
    // Metrics never fail the work they measure.
  }
}

/**
 * Runs one unit of work with cost accounting. `describe` turns the result into the event (request: status and
 * route). A nested call (a cron task inside the scheduled handler) gets its own counters.
 */
export async function withPerf<T>(describe: (result: T | null) => Outcome, work: () => Promise<T>): Promise<T> {
  const scope: PerfScope = {
    d1Queries: 0,
    d1Ms: 0,
    d1ReplicaQueries: 0,
    d1RowsRead: 0,
    d1RowsWritten: 0,
    doCalls: 0,
    doMs: 0,
    doRowsRead: 0,
    doRowsWritten: 0,
    startedAt: Date.now(),
    coldStart,
  };
  coldStart = false;
  try {
    const result = await storage.run(scope, work);
    report(scope, describe(result), false);
    return result;
  } catch (error) {
    report(scope, describe(null), true);
    throw error;
  }
}

/** Instruments a background job (cron task, queue batch, alarm). */
export function withJobPerf<T>(job: string, work: () => Promise<T>): Promise<T> {
  return withPerf(() => ({ event: "server_job", properties: { job } }), work);
}