import test from 'node:test';
import assert from 'node:assert/strict';
import { spawnSync, spawn } from 'node:child_process';
import { chmod, mkdtemp, readFile, writeFile } from 'node:fs/promises';
import { existsSync } from 'node:fs';
import { join } from 'node:path';
import { tmpdir } from 'node:os';

const benchmark = new URL('../scripts/record_worker_benchmark.mjs', import.meta.url).pathname;
const receiptsPath = new URL('../deploy/worker/receipts.json', import.meta.url);
const ALIAS = 'SmolVLM2-2.2B-Instruct';

const FAKE_WORKER = `#!/usr/bin/env python3
import json, os
from http.server import BaseHTTPRequestHandler, HTTPServer
loaded = os.environ['WORKER_MODE'] == 'loaded'
class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        body = json.dumps({'data': [{'id': os.environ['WORKER_ALIAS']}] if loaded else []}).encode()
        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        self.end_headers()
        self.wfile.write(body)
    def log_message(self, *_):
        pass
HTTPServer(('127.0.0.1', int(os.environ['WORKER_PORT'])), Handler).serve_forever()
`;

async function sandbox() {
  const dir = await mkdtemp(join(tmpdir(), 'timmy-bench-'));
  const gpuProbe = join(dir, 'nvidia-smi-stub');
  await writeFile(gpuProbe, '#!/bin/sh\nprintf "NVIDIA A40, 46068, 1234\\n"\n');
  await chmod(gpuProbe, 0o755);
  const worker = join(dir, 'worker.py');
  await writeFile(worker, FAKE_WORKER);
  return { dir, gpuProbe, worker, out: join(dir, 'worker-benchmark.json') };
}

async function startWorker(box, mode) {
  const port = String(30000 + (process.pid % 2000));
  const proc = spawn('python3', [box.worker], {
    env: { ...process.env, WORKER_MODE: mode, WORKER_PORT: port, WORKER_ALIAS: ALIAS },
    stdio: 'ignore',
  });
  for (let attempt = 0; attempt < 50; attempt += 1) {
    const probe = spawnSync('bash', ['-c', `curl --silent --output /dev/null --max-time 1 http://127.0.0.1:${port}/v1/models`], { encoding: 'utf8' });
    if (probe.status === 0) break;
    await new Promise(resolve => setTimeout(resolve, 100));
  }
  return { proc, port };
}

function run(box, port, extra = []) {
  return spawnSync('node', [benchmark, '--out', box.out, '--port', port, '--alias', ALIAS, '--gpu-probe', box.gpuProbe, ...extra], { encoding: 'utf8' });
}

test('a benchmark refuses to record numbers it did not measure on a GPU', async () => {
  const box = await sandbox();
  const { proc, port } = await startWorker(box, 'loaded');
  try {
    const result = spawnSync('node', [benchmark, '--out', box.out, '--port', port, '--alias', ALIAS, '--gpu-probe', join(box.dir, 'missing-smi')], { encoding: 'utf8' });
    assert.notEqual(result.status, 0, 'measuring without a GPU probe must fail, not record zeros');
    assert.match(result.stderr, /GPU|probe|VRAM/i);
    assert.equal(existsSync(box.out), false, 'no receipt may be written from an unmeasured run');
  } finally {
    proc.kill('SIGKILL');
  }
});

test('a benchmark refuses to time a model that is not loaded', async () => {
  const box = await sandbox();
  const { proc, port } = await startWorker(box, 'unloaded');
  try {
    const result = run(box, port);
    assert.notEqual(result.status, 0);
    assert.match(result.stderr, /not loaded|ready/i);
    assert.equal(existsSync(box.out), false, 'an unloaded model must not produce a latency receipt');
  } finally {
    proc.kill('SIGKILL');
  }
});

test('a benchmark records cold start, warm latency, VRAM, and the pins it measured', async () => {
  const box = await sandbox();
  const { proc, port } = await startWorker(box, 'loaded');
  try {
    const result = run(box, port, ['--warm-samples', '5', '--start-command', 'bash -c "sleep 0.3"']);
    assert.equal(result.status, 0, result.stderr);
    const receipt = JSON.parse(await readFile(box.out, 'utf8'));
    assert.equal(receipt.status, 'measured_on_target_gpu');
    assert.equal(receipt.worker.alias, ALIAS);
    assert.match(receipt.worker.baseUrl, new RegExp(`^http://127\\.0\\.0\\.1:${port}/v1$`));
    assert.match(receipt.gpu.name, /NVIDIA A40/);
    assert.equal(receipt.gpu.vramUsedMiB, 1234);
    assert.equal(receipt.gpu.vramTotalMiB, 46068);
    assert.ok(receipt.coldStartMs >= 300, 'cold start is measured, not asserted');
    assert.equal(receipt.warmLatencyMs.samples, 5);
    assert.ok(receipt.warmLatencyMs.p50 >= 0 && receipt.warmLatencyMs.p95 >= receipt.warmLatencyMs.p50);

    const receipts = JSON.parse(await readFile(receiptsPath, 'utf8'));
    assert.equal(receipt.pins.llamaCppCommit, receipts.llamaCpp.commit);
    assert.deepEqual(receipt.pins.modelSha256, receipts.models.map(model => model.sha256));
  } finally {
    proc.kill('SIGKILL');
  }
});

test('a benchmark receipt carries no notes, image bytes, or user data', async () => {
  const box = await sandbox();
  const { proc, port } = await startWorker(box, 'loaded');
  try {
    const result = run(box, port, ['--warm-samples', '3']);
    assert.equal(result.status, 0, result.stderr);
    const raw = await readFile(box.out, 'utf8');
    const receipt = JSON.parse(raw);
    assert.deepEqual(Object.keys(receipt).sort(), [
      'coldStartMs',
      'gpu',
      'measuredAt',
      'pins',
      'schemaVersion',
      'status',
      'warmLatencyMs',
      'worker',
    ]);
    assert.doesNotMatch(raw, /data:image|base64|note|observations|symptom/i);
    // Cold start that was not measured is recorded as unmeasured, never as zero.
    assert.equal(receipt.coldStartMs, null);
  } finally {
    proc.kill('SIGKILL');
  }
});