#!/usr/bin/env node
// Record the target-GPU cost of the pinned worker (issue #15).
//
// This script measures; it never invents. If no GPU can be probed, if the pinned model is not
// loaded, or if the cold start cannot be observed, it exits non-zero and writes no receipt at all,
// because a receipt with zeroed or assumed numbers is worse than no receipt.
import { spawn } from 'node:child_process';
import { readFile, writeFile } from 'node:fs/promises';
import { fileURLToPath } from 'node:url';
import path from 'node:path';

const LOOPBACK = /^(127\.0\.0\.1|localhost|::1)$/;

function parseArgs(argv) {
  const options = {
    out: '',
    host: '127.0.0.1',
    port: '8080',
    alias: 'SmolVLM2-2.2B-Instruct',
    gpuProbe: 'nvidia-smi',
    warmSamples: 20,
    startCommand: '',
    readyTimeoutSeconds: 900,
  };
  for (let index = 0; index < argv.length; index += 1) {
    const [flag, inline] = argv[index].split('=');
    const value = inline ?? argv[index + 1];
    const consume = inline === undefined ? 2 : 1;
    if (flag === '--out') options.out = value;
    else if (flag === '--host') options.host = value;
    else if (flag === '--port') options.port = value;
    else if (flag === '--alias') options.alias = value;
    else if (flag === '--gpu-probe') options.gpuProbe = value;
    else if (flag === '--warm-samples') options.warmSamples = Number(value);
    else if (flag === '--start-command') options.startCommand = value;
    else if (flag === '--ready-timeout') options.readyTimeoutSeconds = Number(value);
    else fail(`Unknown argument: ${flag}`);
    index += consume - 1;
  }
  return options;
}

function fail(message, code = 64) {
  process.stderr.write(`${message}\n`);
  process.exit(code);
}

const reject = message => {
  process.stderr.write(`${message}\n`);
  process.exit(1);
};

const sleep = ms => new Promise(resolve => setTimeout(resolve, ms));

async function modelsServed({ host, port, alias }) {
  try {
    const response = await fetch(`http://${host}:${port}/v1/models`, { signal: AbortSignal.timeout(3000) });
    if (!response.ok) return false;
    const body = await response.json();
    const ids = Array.isArray(body?.data) ? body.data.map(item => item?.id).filter(Boolean) : [];
    return ids.some(id => id === alias || id.includes(alias));
  } catch {
    return false;
  }
}

function runCommand(command, args) {
  return new Promise(resolve => {
    const child = spawn(command, args, { stdio: ['ignore', 'pipe', 'pipe'] });
    let stdout = '';
    let stderr = '';
    child.stdout.on('data', chunk => { stdout += chunk; });
    child.stderr.on('data', chunk => { stderr += chunk; });
    child.on('error', error => resolve({ code: 127, stdout, stderr: String(error.message) }));
    child.on('close', code => resolve({ code, stdout, stderr }));
  });
}

async function probeGpu(probeCommand) {
  const result = await runCommand(probeCommand, [
    '--query-gpu=name,memory.total,memory.used',
    '--format=csv,noheader,nounits',
  ]);
  if (result.code !== 0) {
    reject(`No GPU could be probed with ${probeCommand}: ${result.stderr.trim() || `exit ${result.code}`}. A benchmark receipt needs a real device, so nothing was written.`);
  }
  const [name, total, used] = result.stdout.trim().split('\n')[0].split(',').map(part => part.trim());
  if (!name || !/^\d+$/.test(total) || !/^\d+$/.test(used)) {
    reject(`The GPU probe returned unreadable output (${JSON.stringify(result.stdout.trim())}); refusing to record VRAM it cannot read.`);
  }
  return { name, vramTotalMiB: Number(total), vramUsedMiB: Number(used) };
}

function percentile(sorted, fraction) {
  const index = Math.min(sorted.length - 1, Math.max(0, Math.ceil(fraction * sorted.length) - 1));
  return sorted[index];
}

const options = parseArgs(process.argv.slice(2));
if (!options.out) fail('--out is required so the receipt has a destination');
if (!LOOPBACK.test(options.host)) fail(`The worker is a loopback service; refusing to benchmark ${options.host}.`);
if (!Number.isInteger(options.warmSamples) || options.warmSamples < 1) fail('--warm-samples must be a positive integer');

const receipts = JSON.parse(await readFile(fileURLToPath(new URL('../deploy/worker/receipts.json', import.meta.url)), 'utf8'));

// The device is probed first: nothing is timed until we know the numbers have somewhere to be real.
const gpu = await probeGpu(options.gpuProbe);

if (!(await modelsServed(options))) {
  reject(`${options.alias} is not loaded on ${options.host}:${options.port} yet, so there is no ready worker to time. Start the worker and wait for the post-load gate.`);
}

let coldStartMs = null;
if (options.startCommand) {
  const started = Date.now();
  const started_process = await runCommand('bash', ['-c', options.startCommand]);
  if (started_process.code !== 0) {
    reject(`The cold-start command failed (exit ${started_process.code}): ${started_process.stderr.trim()}`);
  }
  const deadline = Date.now() + options.readyTimeoutSeconds * 1000;
  let ready = false;
  while (Date.now() < deadline) {
    if (await modelsServed(options)) { ready = true; break; }
    await sleep(250);
  }
  if (!ready) reject(`The worker did not become ready within ${options.readyTimeoutSeconds}s, so the cold start could not be measured.`);
  coldStartMs = Date.now() - started;
}

const samples = [];
for (let index = 0; index < options.warmSamples; index += 1) {
  const started = process.hrtime.bigint();
  const ok = await modelsServed(options);
  const elapsedMs = Number(process.hrtime.bigint() - started) / 1e6;
  if (!ok) reject('The worker stopped answering during the warm pass; the receipt would be partial, so nothing was written.');
  samples.push(Math.round(elapsedMs * 1000) / 1000);
}
const sorted = [...samples].sort((left, right) => left - right);

const receipt = {
  schemaVersion: 1,
  status: 'measured_on_target_gpu',
  measuredAt: new Date().toISOString(),
  gpu,
  coldStartMs,
  warmLatencyMs: {
    samples: sorted.length,
    min: sorted[0],
    p50: percentile(sorted, 0.5),
    p95: percentile(sorted, 0.95),
    max: sorted[sorted.length - 1],
  },
  worker: { baseUrl: `http://${options.host}:${options.port}/v1`, alias: options.alias },
  pins: {
    llamaCppCommit: receipts.llamaCpp.commit,
    modelSha256: receipts.models.map(model => model.sha256),
  },
};

await writeFile(path.resolve(options.out), `${JSON.stringify(receipt, null, 2)}\n`, 'utf8');
process.stdout.write(`${path.resolve(options.out)}\n`);