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

const gate = new URL('../scripts/worker_ready.sh', import.meta.url).pathname;
const bootstrap = new URL('../scripts/bootstrap_selfhost_smolvlm.sh', import.meta.url);
const ALIAS = 'SmolVLM2-2.2B-Instruct';

const FAKE_SERVER = `#!/usr/bin/env python3
import json, os, sys
from http.server import BaseHTTPRequestHandler, HTTPServer

mode = os.environ['WORKER_MODE']
counter_path = os.environ['WORKER_COUNTER']
count = int(open(counter_path).read() or '0') if os.path.exists(counter_path) else 0

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        global count
        count += 1
        open(counter_path, 'w').write(str(count))
        loaded = mode == 'loaded' or (mode == 'flip' and count > 2)
        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 withWorker(mode, run) {
  const sandbox = await mkdtemp(join(tmpdir(), 'timmy-worker-ready-'));
  const counter = join(sandbox, 'count.txt');
  await writeFile(counter, '0');
  const server = join(sandbox, 'server.py');
  await writeFile(server, FAKE_SERVER);
  const port = String(20000 + (process.pid % 4000));
  const { spawn } = await import('node:child_process');
  const proc = spawn('python3', [server], {
    env: { ...process.env, WORKER_MODE: mode, WORKER_COUNTER: counter, WORKER_PORT: port, WORKER_ALIAS: ALIAS },
    stdio: 'ignore',
  });
  // Wait for the socket to accept connections before probing.
  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));
  }
  try {
    return await run(port);
  } finally {
    proc.kill('SIGKILL');
  }
}

function gateRun(port, extra = []) {
  return spawnSync('bash', [gate, '--port', port, '--alias', ALIAS, ...extra], { encoding: 'utf8' });
}

test('readiness passes once the pinned model is actually served', async () => {
  await withWorker('loaded', async port => {
    const result = gateRun(port);
    assert.equal(result.status, 0, result.stderr);
    assert.match(result.stdout, new RegExp(`${ALIAS} is loaded`));
  });
});

test('a reachable worker with the model still loading is not ready', async () => {
  await withWorker('unloaded', async port => {
    const result = gateRun(port);
    assert.equal(result.status, 3, 'reachable-but-unloaded must be its own failure, not a network error');
    assert.match(result.stderr, /not loaded yet/i);
  });
});

test('an unreachable worker is reported as unreachable', async () => {
  const result = spawnSync('bash', [gate, '--port', '9', '--alias', ALIAS], { encoding: 'utf8' });
  assert.equal(result.status, 2);
  assert.match(result.stderr, /not reachable/i);
});

test('readiness never probes anything but loopback', () => {
  const result = spawnSync('bash', [gate, '--host', '10.0.0.5', '--port', '8080'], { encoding: 'utf8' });
  assert.equal(result.status, 64);
  assert.match(result.stderr, /loopback/i);
});

test('readiness can wait for the model load it is gating', async () => {
  await withWorker('flip', async port => {
    const started = Date.now();
    const result = gateRun(port, ['--wait', '15']);
    assert.equal(result.status, 0, result.stderr);
    assert.ok(Date.now() - started >= 300, 'the gate must actually wait rather than pass on the first probe');
  });
});

test('the bootstrap exposes the same post-load gate instead of duplicating it', async () => {
  const script = await readFile(bootstrap, 'utf8');
  assert.match(script, /worker_ready\.sh/);
  assert.match(script, /^  ready\)$/m);
});