// Benchmark: net.http_post (pg_net) vs pg_notify + an external LISTEN relay.
//
//   bun scripts/bench_pg_net_vs_pg_notify.ts
//
// Needs Docker and `bun install`. The script starts three containers on one
// Docker network (no host NAT in the request path):
//   - db:    the Supabase Postgres image production runs, with pg_net (4 CPU, 4 GB)
//   - sink:  a Bun HTTP server standing in for /functions/v1/triggers/*
//   - relay: a Bun process that LISTENs and forwards each payload as an HTTP
//            POST (1 CPU, 256 MB). pg_notify cannot send HTTP itself, so a
//            production switch needs this always-on process.
//   - load:  pgbench in its own container (2 CPU), so the load generator's CPU
//            is not counted as database CPU
// For each variant it reports caller latency, TPS, WAL per transaction, database CPU (db container cgroup), pg_net
// worker CPU / memory, relay CPU / memory, net.* table growth and the time
// from the first enqueue until the sink received the whole burst (BENCH_SINK_DELAY_MS adds endpoint latency). Containers are removed at the end
// (BENCH_KEEP=1 keeps them).

import { spawn } from 'node:child_process'
import { resolve } from 'node:path'
import process from 'node:process'
import pg from 'pg'

const IMAGE = process.env.BENCH_PG_IMAGE ?? 'public.ecr.aws/supabase/postgres:17.6.1.143'
const BUN_IMAGE = process.env.BENCH_BUN_IMAGE ?? 'oven/bun:1.3-alpine'
const NETWORK = 'pgnet-bench'
const DB = 'pgnet-bench-db'
const SINK = 'pgnet-bench-sink'
const RELAY = 'pgnet-bench-relay'
const LOAD = 'pgnet-bench-load'
const DB_PORT = Number(process.env.BENCH_DB_PORT ?? 55439)
const SINK_PORT = Number(process.env.BENCH_SINK_PORT ?? 18080)
const HOST_PG_URL = `postgres://supabase_admin:postgres@127.0.0.1:${DB_PORT}/postgres`
const NET_PG_URL = `postgres://supabase_admin:postgres@${DB}:5432/postgres`
const SINK_URL = `http://${SINK}:8080/sink`
const CHANNEL = 'capgo_http_request'
const TOTAL_TX = Number(process.env.BENCH_TOTAL_TX ?? 20000)
const CLIENTS = (process.env.BENCH_CLIENTS ?? '1,16,64').split(',').map(Number)
// Simulated endpoint latency (an edge function answers in tens of ms).
const SINK_DELAY_MS = Number(process.env.BENCH_SINK_DELAY_MS ?? 0)
const REPO = resolve(import.meta.dir, '..')
const SCRIPT_IN_CONTAINER = '/app/scripts/bench_pg_net_vs_pg_notify.ts'

// Same headers and body shape process_function_queue sends today.
const HEADERS = { 'Content-Type': 'application/json', 'apisecret': 'bench-secret-0123456789abcdef0123456789abcdef' }
const BODY_SQL = `pg_catalog.jsonb_build_object('queue_name', 'on_version_update', 'batch_size', 950, 'wait_for_completion', false)`

function sh(cmd: string, args: string[], input?: string): Promise<{ code: number, out: string, err: string }> {
  return new Promise((done) => {
    const child = spawn(cmd, args, { stdio: ['pipe', 'pipe', 'pipe'] })
    let out = ''
    let err = ''
    child.stdout.on('data', d => out += d)
    child.stderr.on('data', d => err += d)
    child.on('close', code => done({ code: code ?? 1, out, err }))
    if (input)
      child.stdin.write(input)
    child.stdin.end()
  })
}

async function docker(args: string[], input?: string) {
  const res = await sh('docker', args, input)
  if (res.code !== 0)
    throw new Error(`docker ${args.join(' ')} failed: ${res.err}`)
  return res.out
}

const exec = (container: string, args: string[], input?: string) => docker(['exec', '-i', container, ...args], input)

/** CPU (ms) and memory (KB) of a whole container from its cgroup. */
async function cgroup(container: string) {
  const out = await exec(container, ['sh', '-c', 'cat /sys/fs/cgroup/cpu.stat /sys/fs/cgroup/memory.stat'])
  const num = (name: string) => Number(new RegExp(`^${name} (\\d+)`, 'm').exec(out)?.[1] ?? 0)
  return { cpuMs: num('usage_usec') / 1000, anonKb: num('anon') / 1024 }
}

/** CPU (ms) and private memory (KB) of one Postgres process. */
async function proc(pid: number) {
  const stat = await exec(DB, ['cat', `/proc/${pid}/stat`])
  const fields = stat.slice(stat.lastIndexOf(')') + 2).split(' ')
  const status = await exec(DB, ['cat', `/proc/${pid}/status`])
  return {
    cpuMs: (Number(fields[11]) + Number(fields[12])) * 10,
    anonKb: Number(/RssAnon:\s+(\d+)/.exec(status)?.[1] ?? 0),
  }
}

// ---------------------------------------------------------------------------
// sink (runs in its container): counts requests
// ---------------------------------------------------------------------------
function runSink() {
  let count = 0
  Bun.serve({
    port: 8080,
    hostname: '0.0.0.0',
    async fetch(req) {
      const path = new URL(req.url).pathname
      if (path === '/count')
        return Response.json({ count })
      if (path === '/reset') {
        count = 0
        return Response.json({ count })
      }
      await req.text()
      if (SINK_DELAY_MS)
        await Bun.sleep(SINK_DELAY_MS)
      count++
      return Response.json({ status: 'ok' })
    },
  })
}

// ---------------------------------------------------------------------------
// relay (runs in its container): LISTEN -> HTTP POST
// ---------------------------------------------------------------------------
async function runRelay() {
  const client = new pg.Client({ connectionString: NET_PG_URL, application_name: 'bench_notify_relay' })
  await client.connect()
  client.on('notification', (msg) => {
    fetch(SINK_URL, { method: 'POST', headers: HEADERS, body: msg.payload, signal: AbortSignal.timeout(8000) })
      .then(r => r.text())
      .catch(error => console.error('relay POST failed', error))
  })
  await client.query(`LISTEN ${CHANNEL}`)
  console.log('relay ready')
}

// ---------------------------------------------------------------------------
// orchestrator (host)
// ---------------------------------------------------------------------------
type Variant = 'baseline' | 'pg_net' | 'pg_notify'
type Workload = 'enqueue_only' | 'update_plus_enqueue'

function pgbenchScript(variant: Variant, workload: Workload) {
  const enqueue = {
    baseline: `SELECT 1;`,
    pg_net: `SELECT net.http_post(url := '${SINK_URL}', headers := '${JSON.stringify(HEADERS)}'::jsonb, body := ${BODY_SQL}, timeout_milliseconds := 8000);`,
    pg_notify: `SELECT pg_catalog.pg_notify('${CHANNEL}', ${BODY_SQL}::text);`,
  }[variant]
  if (workload === 'enqueue_only')
    return `${enqueue}\n`
  // A business write that fires the enqueue in the same transaction, like the
  // updates_cache_purge triggers on channels / apps / app_versions.
  return `\\set id random(1, 10000)\nBEGIN;\nUPDATE public.bench_rows SET n = n + 1 WHERE id = :id;\n${enqueue}\nCOMMIT;\n`
}

async function sinkCount() {
  const res = await fetch(`http://127.0.0.1:${SINK_PORT}/count`)
  return ((await res.json()) as { count: number }).count
}

/**
 * Resolves with the performance.now() at which the sink had every request.
 * The default allows pg_net's ~200 req/s at a quarter of that rate, plus 60 s.
 */
async function waitDelivered(target: number, timeoutMs = 60_000 + target * 20) {
  const start = performance.now()
  while (performance.now() - start < timeoutMs) {
    if (await sinkCount() >= target)
      return performance.now()
    await Bun.sleep(5)
  }
  throw new Error(`only ${await sinkCount()} of ${target} requests delivered after ${timeoutMs} ms`)
}

async function waitForDb() {
  for (let i = 0; i < 120; i++) {
    const res = await sh('docker', ['exec', DB, 'pg_isready', '-h', '127.0.0.1', '-U', 'supabase_admin'])
    if (res.code === 0) {
      const client = new pg.Client({ connectionString: HOST_PG_URL })
      try {
        await client.connect()
        await client.end()
        return
      }
      catch {}
    }
    await Bun.sleep(500)
  }
  throw new Error('database did not start')
}

async function removeContainers() {
  await sh('docker', ['rm', '-f', LOAD, RELAY, SINK, DB])
  await sh('docker', ['network', 'rm', NETWORK])
}

async function startContainers() {
  await removeContainers()
  await docker(['network', 'create', NETWORK])
  await docker(['run', '-d', '--name', DB, '--network', NETWORK, '--cpus=4', '--memory=4g', '-p', `${DB_PORT}:5432`, '-e', 'POSTGRES_PASSWORD=postgres', IMAGE, 'postgres', '-c', 'listen_addresses=*', '-c', 'shared_preload_libraries=pg_net', '-c', 'max_connections=200'])
  const bun = (name: string, mode: string, extra: string[]) => docker(['run', '-d', '--name', name, '--network', NETWORK, ...extra, '-v', `${REPO}:/app:ro`, '-w', '/app', BUN_IMAGE, 'bun', SCRIPT_IN_CONTAINER, mode])
  await bun(SINK, 'sink', ['--cpus=2', '-p', `${SINK_PORT}:8080`, '-e', `BENCH_SINK_DELAY_MS=${SINK_DELAY_MS}`])
  await waitForDb()
  await bun(RELAY, 'relay', ['--cpus=1', '--memory=256m'])
  await docker(['run', '-d', '--name', LOAD, '--network', NETWORK, '--cpus=2', '-e', 'PGPASSWORD=postgres', '--entrypoint', 'sleep', IMAGE, 'infinity'])
}

interface Result {
  variant: Variant
  workload: Workload
  clients: number
  tx: number
  tps: number
  latencyAvgMs: number
  walBytesPerTx: number
  dbCpuMsPerTx: number
  /** First enqueue -> sink received every request. */
  deliveredMs: number
  /** pgbench end -> sink received every request. */
  drainAfterEnqueueMs: number
  delivered: number
  pgNetWorkerCpuMs?: number
  pgNetWorkerAnonKb?: number
  netTablesKb?: number
  relayCpuMs?: number
  relayAnonKb?: number
  relayBackendCpuMs?: number
}

async function main() {
  await startContainers()
  const db = new pg.Client({ connectionString: HOST_PG_URL })
  await db.connect()
  await db.query(`CREATE EXTENSION IF NOT EXISTS pg_net`)
  await db.query(`CREATE TABLE public.bench_rows (id int PRIMARY KEY, n bigint NOT NULL DEFAULT 0)`)
  await db.query(`INSERT INTO public.bench_rows (id) SELECT g FROM generate_series(1, 10000) g`)

  // The relay connects once the database accepts connections.
  let relayBackendPid = 0
  for (let i = 0; i < 100 && !relayBackendPid; i++) {
    relayBackendPid = Number((await db.query(`SELECT pid FROM pg_stat_activity WHERE application_name = 'bench_notify_relay' AND query ILIKE 'LISTEN%'`)).rows[0]?.pid ?? 0)
    if (!relayBackendPid)
      await Bun.sleep(200)
  }
  const workerPid = Number((await db.query(`SELECT pid FROM pg_stat_activity WHERE backend_type ILIKE 'pg_net%'`)).rows[0]?.pid ?? 0)
  if (!relayBackendPid || !workerPid)
    throw new Error(`missing relay backend (${relayBackendPid}) or pg_net worker (${workerPid})`)

  const results: Result[] = []
  for (const workload of ['enqueue_only', 'update_plus_enqueue'] as Workload[]) {
    for (const clients of CLIENTS) {
      for (const variant of ['baseline', 'pg_net', 'pg_notify'] as Variant[]) {
        const perClient = Math.ceil(TOTAL_TX / clients)
        const tx = perClient * clients
        // Same state for every run: empty pg_net tables (prod truncates
        // net._http_response hourly), clean checkpoint, zeroed sink.
        await db.query(`TRUNCATE net._http_response`)
        await db.query(`VACUUM FULL net.http_request_queue, net._http_response, public.bench_rows`)
        await db.query(`CHECKPOINT`)
        await fetch(`http://127.0.0.1:${SINK_PORT}/reset`)
        await Bun.sleep(500)

        await exec(LOAD, ['sh', '-c', 'cat > /tmp/bench.sql'], pgbenchScript(variant, workload))
        const walStart = (await db.query(`SELECT pg_current_wal_insert_lsn() AS lsn`)).rows[0].lsn
        const dbStart = await cgroup(DB)
        const relayStart = await cgroup(RELAY)
        const workerStart = await proc(workerPid)
        const relayBackendStart = await proc(relayBackendPid)

        const enqueueStart = performance.now()
        const out = await exec(LOAD, ['pgbench', '-n', '-U', 'supabase_admin', '-h', DB, '-d', 'postgres', '-c', String(clients), '-j', String(Math.min(clients, 4)), '-t', String(perClient), '-M', 'prepared', '-f', '/tmp/bench.sql'])
        const enqueueEnd = performance.now()
        const deliveredAt = variant === 'baseline' ? enqueueEnd : await waitDelivered(tx)
        // pg_net writes each response row after the request; let it settle.
        await Bun.sleep(1000)

        const walBytes = Number((await db.query(`SELECT pg_wal_lsn_diff(pg_current_wal_insert_lsn(), $1) AS b`, [walStart])).rows[0].b)
        const dbEnd = await cgroup(DB)
        const relayEnd = await cgroup(RELAY)
        const workerEnd = await proc(workerPid)
        const relayBackendEnd = await proc(relayBackendPid)
        const netTablesKb = Number((await db.query(`SELECT (pg_total_relation_size('net.http_request_queue') + pg_total_relation_size('net._http_response')) / 1024 AS kb`)).rows[0].kb)

        const result: Result = {
          variant,
          workload,
          clients,
          tx,
          tps: Math.round(Number(/tps = ([\d.]+)/.exec(out)?.[1] ?? 0)),
          latencyAvgMs: Number(/latency average = ([\d.]+) ms/.exec(out)?.[1] ?? 0),
          walBytesPerTx: Math.round(walBytes / tx),
          // Includes the pg_net worker and the relay's backend: all of it is
          // database CPU.
          dbCpuMsPerTx: Number(((dbEnd.cpuMs - dbStart.cpuMs) / tx).toFixed(4)),
          deliveredMs: Math.round(deliveredAt - enqueueStart),
          drainAfterEnqueueMs: Math.round(deliveredAt - enqueueEnd),
          delivered: await sinkCount(),
        }
        if (variant === 'pg_net') {
          result.pgNetWorkerCpuMs = workerEnd.cpuMs - workerStart.cpuMs
          result.pgNetWorkerAnonKb = workerEnd.anonKb
          result.netTablesKb = netTablesKb
        }
        if (variant === 'pg_notify') {
          result.relayCpuMs = Math.round(relayEnd.cpuMs - relayStart.cpuMs)
          result.relayAnonKb = Math.round(relayEnd.anonKb)
          result.relayBackendCpuMs = relayBackendEnd.cpuMs - relayBackendStart.cpuMs
        }
        results.push(result)
        console.error(JSON.stringify(result))
      }
    }
  }
  await db.end()
  console.log(JSON.stringify(results, null, 2))
}

const mode = process.argv[2] ?? 'run'
if (mode === 'sink') {
  runSink()
}
else if (mode === 'relay') {
  await runRelay()
}
else {
  try {
    await main()
  }
  finally {
    if (!process.env.BENCH_KEEP)
      await removeContainers()
  }
}