package test;

// WW uses one fork per test because it does not yet have Hare's
// setjmp/onabort machinery (ref/hare/test/+test.ha:250-322).  The pipe
// makes a normal return distinguishable from a test calling os.exit(0).

import fnmatch;
import os;
import time;

def TST_PASS: i32 = 0;
def TST_FAIL: i32 = 1;
def TST_SKIP: i32 = 2;
def TST_HARNESS: i32 = 3;

def TST_COMPLETE: u8 = 80u8;	// P
def TST_SKIPPED: u8 = 83u8;	// S
def TST_EXPECT_ABORT: u8 = 65u8;	// A
def TST_GROUP_READY: u8 = 71u8;	// G
def TST_CONTROL_MAX: i32 = 4096;
def TST_FRAME_HEADER: i32 = 3;
def TST_EINTR: i64 = -4i64;
def TST_EAGAIN: i64 = -11i64;
def TST_ESRCH: i32 = -3;
def TST_DEFAULT_TIMEOUT_MS: i64 = 30000i64;
def TST_MAX_TIMEOUT_MS: i64 = 3600000i64;
def TST_POLL_MS: i64 = 1i64;
def TST_TERM_GRACE_MS: i64 = 100i64;
def TST_TIMEOUT_PREFIX_LEN: i32 = 12;
def TST_PACKAGE_PREFIX_LEN: i32 = 9;

type tstresult = struct {
	kind: i32,
	completed: bool,
	timedout: bool,
	exited: bool,
	detail: i32,
	control: [TST_CONTROL_MAX]u8,
	ncontrol: i32,
};

let tstactive: bool = false;
let tstcontrolfd: i32 = -1;
let tstname: str;
let tstwantabort: bool = false;

// run preserves the compiler-generated [](str, *fn() void) descriptor.
// Package assembly can therefore aggregate tests without another registry
// format or a compiler invocation per assertion.
export fn run(tests: [](str, *fn() void)) i32 = {
	let av: []str = os.args();
	let list: bool = tstlistmode(av);
	let timeoutms: i64 = tsttimeout(av);
	let pkgprefix: str = "";
	if (timeoutms <= 0) {
		tstputs("test: invalid or duplicate -timeout-ms option\n");
		return 1;
	};
	if (!tstpackage(av, &pkgprefix)) {
		tstputs("test: invalid or duplicate -package option\n");
		return 1;
	};
	let selected: i32 = 0;
	let started: i32 = 0;
	let completed: i32 = 0;
	let passed: i32 = 0;
	let failed: i32 = 0;
	let skipped: i32 = 0;
	let harness: i32 = 0;
	let i: i32 = 0;
	for (i < tests.len) {
		let row: (str, *fn() void) = tests[i];
		let name: str = tstqualified(pkgprefix, row.0);
		if (!tstenabled(name, av)) { i += 1; continue; };
		selected += 1;
		if (list) {
			tstputs(name);
			tstputs("\n");
			i += 1;
			continue;
		};

		tstputs(name);
		tstputs(" ... ");
		let result: tstresult;
		if (!tstrunone(row.1, name, timeoutms, &result)) {
			harness += 1;
			tstputs("HARNESS (could not start)\n");
			i += 1;
			continue;
		};
		started += 1;
		if (result.completed) { completed += 1; };
		switch (result.kind) {
		case TST_PASS:
			passed += 1;
			tstputs("ok\n");
		case TST_SKIP:
			skipped += 1;
			tstputs("SKIP: ");
			tstfdwriteall(os.STDOUT_FILENO,
			    &result.control[TST_FRAME_HEADER],
			    (result.ncontrol - TST_FRAME_HEADER): u64);
			tstputs("\n");
		case TST_FAIL:
			failed += 1;
			tstputfailure(&result);
		case TST_HARNESS:
			harness += 1;
			tstputharness(&result);
		};
		i += 1;
	};

	if (list) {
		return 0;
	};
	if (selected == 0) {
		tsterrputs("testing: warning: no tests to run\n");
		tstputaccounting(tests.len, 0, 0, 0);
		return 0;
	};
	tstputuint(passed);
	tstputs(" passed, ");
	tstputuint(failed);
	tstputs(" failed, ");
	tstputuint(skipped);
	tstputs(" skipped, ");
	tstputuint(harness);
	tstputs(" harness errors\n");
	tstputaccounting(tests.len, selected, started, completed);
	if (failed != 0 || harness != 0 || started != selected ||
	    completed != started) {
		return 1;
	};
	return 0;
};

// skip and expectabort are records rather than exit-code conventions: an
// exit status alone cannot say whether user code returned, skipped, or died.
export fn skip(reason: str) never = {
	if (!tstactive) { abort("test.skip called outside a test"); };
	if (reason.len > TST_CONTROL_MAX - TST_FRAME_HEADER) {
		tstwriteframe(0u8, "");
		os.exit(0);
	};
	tstwriteframe(TST_SKIPPED, reason);
	os.close(tstcontrolfd);
	os.exit(0);
};

export fn expectabort() void = {
	if (!tstactive) { abort("test.expectabort called outside a test"); };
	if (!tstwantabort) {
		tstwriteframe(TST_EXPECT_ABORT, "");
		tstwantabort = true;
	};
};

export fn current() str = {
	if (!tstactive) { abort("test.current called outside a test"); };
	return tstname;
};

fn tstlistarg(s: str) bool = {
	return s == "-list" || s == "--list";
};

fn tsttimeoutarg(s: str) bool = {
	let prefix: str = "-timeout-ms=";
	if (s.len < prefix.len) { return false; };
	let i: i32 = 0;
	for (i < prefix.len) {
		if (s[i] != prefix[i]) { return false; };
		i += 1;
	};
	return true;
};

fn tstpackagearg(s: str) bool = {
	let prefix: str = "-package=";
	if (s.len < prefix.len) { return false; };
	let i: i32 = 0;
	for (i < prefix.len) {
		if (s[i] != prefix[i]) { return false; };
		i += 1;
	};
	return true;
};

fn tstpackage(av: []str, out: *str) bool = {
	let seen: bool = false;
	let i: i32 = 1;
	for (i < av.len) {
		if (tstpackagearg(av[i])) {
			if (seen || av[i].len == TST_PACKAGE_PREFIX_LEN) {
				return false;
			};
			seen = true;
			out.ptr = av[i].ptr + (TST_PACKAGE_PREFIX_LEN: u64);
			out.len = av[i].len - TST_PACKAGE_PREFIX_LEN;
		};
		i += 1;
	};
	return true;
};

fn tstqualified(pkgprefix: str, name: str) str = {
	if (pkgprefix.len == 0) { return name; };
	let b: []u8 = alloc([], (pkgprefix.len + name.len + 1): u64)!;
	b.len = pkgprefix.len + name.len + 1;
	let i: i32 = 0;
	for (i < pkgprefix.len) { b[i] = pkgprefix[i]; i += 1; };
	b[pkgprefix.len] = '.';
	i = 0;
	for (i < name.len) {
		b[pkgprefix.len + 1 + i] = name[i];
		i += 1;
	};
	let out: str;
	out.ptr = b.ptr;
	out.len = b.len;
	return out;
};

fn tsttimeout(av: []str) i64 = {
	let timeoutms: i64 = TST_DEFAULT_TIMEOUT_MS;
	let seen: bool = false;
	let i: i32 = 1;
	for (i < av.len) {
		if (tsttimeoutarg(av[i])) {
			if (seen || av[i].len == TST_TIMEOUT_PREFIX_LEN) {
				return -1i64;
			};
			seen = true;
			let n: i64 = 0i64;
			let j: i32 = TST_TIMEOUT_PREFIX_LEN;
			for (j < av[i].len) {
				let c: u8 = av[i][j];
				if (c < 48u8 || c > 57u8) { return -1i64; };
				let digit: i64 = (c - 48u8): i64;
				if (n > (TST_MAX_TIMEOUT_MS - digit) / 10i64) {
					return -1i64;
				};
				n = n * 10i64 + digit;
				j += 1;
			};
			if (n == 0) { return -1i64; };
			timeoutms = n;
		};
		i += 1;
	};
	return timeoutms;
};

fn tstlistmode(av: []str) bool = {
	let i: i32 = 1;
	for (i < av.len) {
		if (tstlistarg(av[i])) { return true; };
		i += 1;
	};
	return false;
};

fn tstenabled(name: str, av: []str) bool = {
	let havefilter: bool = false;
	let leaf: str = tstleaf(name);
	let i: i32 = 1;
	for (i < av.len) {
		if (!tstlistarg(av[i]) && !tsttimeoutarg(av[i])
			&& !tstpackagearg(av[i])) {
			havefilter = true;
			if (fnmatch.fnmatch(av[i], name, fnmatch.flag.NONE) ||
			    fnmatch.fnmatch(av[i], leaf, fnmatch.flag.NONE)) {
				return true;
			};
		};
		i += 1;
	};
	return !havefilter;
};

fn tstleaf(name: str) str = {
	let off: i32 = 0;
	let i: i32 = 0;
	for (i < name.len) {
		if (name[i] == 46u8) { off = i + 1; };
		i += 1;
	};
	let leaf: str;
	leaf.ptr = name.ptr + (off: u64);
	leaf.len = name.len - off;
	return leaf;
};

// One atomic bounded frame is enough after the leader is reaped. Descendants
// may inherit the writer, so record collection never waits for EOF.
fn tstrunone(f: *fn() void, name: str, timeoutms: i64,
    result: *tstresult) bool = {
	result.kind = TST_HARNESS;
	result.completed = false;
	result.timedout = false;
	result.exited = false;
	result.detail = 0;
	result.ncontrol = 0;
	let fds: [2]i32;
	if (os.pipe2(&fds, os.O_CLOEXEC | os.O_NONBLOCK) < 0) {
		return false;
	};
	let pid: i32 = os.fork();
	if (pid < 0) {
		os.close(fds[0]);
		os.close(fds[1]);
		return false;
	};
	if (pid == 0) {
		os.close(fds[0]);
		tstactive = true;
		tstcontrolfd = fds[1];
		tstname = name;
		tstwantabort = false;
		if (os.setpgid(0, 0) != 0) {
			tstwritebyte(0u8);
			os.exit(1);
		};
		tstwritebyte(TST_GROUP_READY);
		(*f)();
		if (!tstwantabort) { tstwriteframe(TST_COMPLETE, ""); };
		os.close(tstcontrolfd);
		os.exit(0);
	};
	if (pid <= 1) {
		os.close(fds[0]);
		os.close(fds[1]);
		return false;
	};
	os.close(fds[1]);
	let parentgroup: i32 = os.setpgid(pid, pid);
	let ready: u8 = 0u8;
	let began: time.instant = time.now(time.clock.monotonic);
	let readied: i32 = tstreadready(fds[0], &ready, began, timeoutms);
	if (readied == 0) {
		result.timedout = true;
		if (parentgroup == 0) { os.kill(-pid, os.SIGKILL); }
		else { os.kill(pid, os.SIGKILL); };
		let status: i32 = 0;
		if (tstwait(pid, &status, 0) == pid) { result.completed = true; };
		result.kind = TST_FAIL;
		os.close(fds[0]);
		return true;
	};
	if (readied < 0 || ready != TST_GROUP_READY) {
		if (pid > 1) { os.kill(pid, os.SIGKILL); };
		let ignored: i32 = 0;
		tstwait(pid, &ignored, 0);
		os.close(fds[0]);
		return true;
	};
	// The child byte validates the group even when it wins the setpgid race.
	let status: i32 = 0;
	let waited: i32 = 0;
	for (true) {
		waited = tstwait(pid, &status, os.WNOHANG);
		if (waited == pid || waited < 0) { break; };
		if (tstelapsedms(began) >= timeoutms) {
			result.timedout = true;
			os.kill(-pid, os.SIGTERM);
			let termbegan: time.instant = time.now(time.clock.monotonic);
			for (tstelapsedms(termbegan) < TST_TERM_GRACE_MS) {
				waited = tstwait(pid, &status, os.WNOHANG);
				if (waited == pid || waited < 0) { break; };
				tstsleep();
			};
			if (waited == 0) {
				os.kill(-pid, os.SIGKILL);
				waited = tstwait(pid, &status, 0);
			};
			break;
		};
		tstsleep();
	};
	if (waited != pid) {
		os.kill(-pid, os.SIGKILL);
		waited = tstwait(pid, &status, 0);
		if (waited == pid) {
			result.completed = true;
			tstcleargroup(pid);
		};
		os.close(fds[0]);
		return true;
	};
	result.completed = true;
	if (!tstcleargroup(pid)) {
		os.close(fds[0]);
		return true;
	};
	let full: bool = false;
	for (result.ncontrol < TST_CONTROL_MAX) {
		let n: i64 = os.read(fds[0], &result.control[result.ncontrol],
		    (TST_CONTROL_MAX - result.ncontrol): u64);
		if (n == TST_EINTR) { continue; };
		if (n == TST_EAGAIN || n == 0) { break; };
		if (n < 0) {
			os.close(fds[0]);
			return true;
		};
		result.ncontrol += n: i32;
	};
	if (result.ncontrol == TST_CONTROL_MAX) {
		let extra: u8 = 0u8;
		let n: i64 = os.read(fds[0], &extra, 1u64);
		full = n != TST_EAGAIN && n != 0;
	};
	os.close(fds[0]);
	if (full) { return true; };

	result.exited = os.wifexited(status);
	if (result.exited) {
		result.detail = os.wexitstatus(status);
	} else if (os.wifsignaled(status)) {
		result.detail = os.wtermsig(status);
	} else {
		return true;
	};
	tstinterpret(result);
	return true;
};

fn tstinterpret(result: *tstresult) void = {
	if (result.timedout) {
		result.kind = TST_FAIL;
		return;
	};
	if (result.ncontrol == 0) {
		if (result.exited && result.detail == 0) {
			result.kind = TST_HARNESS;
		} else {
			result.kind = TST_FAIL;
		};
		return;
	};
	if (result.ncontrol < TST_FRAME_HEADER) {
		result.kind = TST_HARNESS;
		return;
	};
	let code: u8 = result.control[0];
	let payload: i32 = (result.control[1]: i32) |
	    ((result.control[2]: i32) << 8i32);
	if (result.ncontrol != TST_FRAME_HEADER + payload) {
		result.kind = TST_HARNESS;
		return;
	};
	if (code == TST_COMPLETE && payload == 0) {
		if (result.exited && result.detail == 0) {
			result.kind = TST_PASS;
		} else {
			result.kind = TST_HARNESS;
		};
		return;
	};
	if (code == TST_SKIPPED) {
		if (result.exited && result.detail == 0) {
			result.kind = TST_SKIP;
		} else {
			result.kind = TST_HARNESS;
		};
		return;
	};
	if (code == TST_EXPECT_ABORT && payload == 0) {
		if (result.exited && result.detail != 0) {
			result.kind = TST_PASS;
		} else {
			result.kind = TST_FAIL;
		};
		return;
	};
	result.kind = TST_HARNESS;
};

fn tstwritebyte(b: u8) void = {
	if (!tstwriteall(&b, 1u64)) {
		os.exit(1);
	};
};

fn tstwriteframe(code: u8, payload: str) void = {
	let frame: [TST_CONTROL_MAX]u8;
	frame[0] = code;
	frame[1] = (payload.len & 255): u8;
	frame[2] = ((payload.len >> 8) & 255): u8;
	let i: i32 = 0;
	for (i < payload.len) {
		frame[TST_FRAME_HEADER + i] = payload[i];
		i += 1;
	};
	if (!tstwriteall(&frame[0], (TST_FRAME_HEADER + payload.len): u64)) {
		os.exit(1);
	};
};

fn tstwriteall(buf: *u8, n: u64) bool = {
	return tstfdwriteall(tstcontrolfd, buf, n);
};

fn tstfdwriteall(fd: i32, buf: *u8, n: u64) bool = {
	let off: u64 = 0u64;
	for (off < n) {
		let wrote: i64 = os.write(fd, buf + off, n - off);
		if (wrote == TST_EINTR) { continue; };
		if (wrote <= 0) { return false; };
		off += wrote: u64;
	};
	return true;
};

fn tstreadready(fd: i32, b: *u8, began: time.instant,
    timeoutms: i64) i32 = {
	for (true) {
		let n: i64 = os.read(fd, b, 1u64);
		if (n == TST_EINTR) { continue; };
		if (n == 1i64) { return 1; };
		if (n == TST_EAGAIN) {
			if (tstelapsedms(began) >= timeoutms) { return 0; };
			tstsleep();
			continue;
		};
		return -1;
	};
};

fn tstwait(pid: i32, status: *i32, options: i32) i32 = {
	for (true) {
		let waited: i32 = os.wait4(pid, status, options, nil: *void);
		if (waited: i64 != TST_EINTR) { return waited; };
	};
};

fn tstelapsedms(began: time.instant) i64 = {
	let now: time.instant = time.now(time.clock.monotonic);
	return (time.diff(began, now): i64) / (time.millisecond: i64);
};

fn tstsleep() void = {
	time.sleep((TST_POLL_MS * (time.millisecond: i64)): time.duration,
	    time.clock.monotonic);
};

fn tstgroupexists(pgid: i32) i32 = {
	if (pgid <= 1) { return -1; };
	let r: i32 = os.kill(-pgid, 0);
	if (r == 0) { return 1; };
	if (r == TST_ESRCH) { return 0; };
	return -1;
};

fn tstcleargroup(pgid: i32) bool = {
	let state: i32 = tstgroupexists(pgid);
	if (state == 0) { return true; };
	if (state < 0) { return false; };
	os.kill(-pgid, os.SIGTERM);
	let began: time.instant = time.now(time.clock.monotonic);
	for (tstelapsedms(began) < TST_TERM_GRACE_MS) {
		state = tstgroupexists(pgid);
		if (state == 0) { return true; };
		if (state < 0) { return false; };
		tstsleep();
	};
	os.kill(-pgid, os.SIGKILL);
	began = time.now(time.clock.monotonic);
	for (tstelapsedms(began) < TST_TERM_GRACE_MS) {
		state = tstgroupexists(pgid);
		if (state == 0) { return true; };
		if (state < 0) { return false; };
		tstsleep();
	};
	return tstgroupexists(pgid) == 0;
};

fn tstputfailure(result: *tstresult) void = {
	if (result.timedout) {
		tstputs("FAIL(timeout)\n");
		return;
	};
	if (result.ncontrol == TST_FRAME_HEADER &&
	    result.control[0] == TST_EXPECT_ABORT &&
	    result.exited && result.detail == 0) {
		tstputs("FAIL (expected abort)\n");
		return;
	};
	if (result.exited) {
		tstputs("FAIL (exit ");
		tstputuint(result.detail);
		tstputs(")\n");
	} else {
		tstputs("FAIL (signal ");
		tstputuint(result.detail);
		tstputs(")\n");
	};
};

fn tstputharness(result: *tstresult) void = {
	if (result.completed && result.ncontrol == 0 && result.exited &&
	    result.detail == 0) {
		tstputs("HARNESS (incomplete result)\n");
	} else if (result.ncontrol == 0) {
		tstputs("HARNESS (wait or read error)\n");
	} else {
		tstputs("HARNESS (malformed or contradictory result)\n");
	};
};

fn tstputaccounting(discovered: i32, selected: i32, started: i32,
    completed: i32) void = {
	tstputuint(discovered);
	tstputs(" discovered, ");
	tstputuint(selected);
	tstputs(" selected, ");
	tstputuint(started);
	tstputs(" started, ");
	tstputuint(completed);
	tstputs(" completed\n");
};

fn tstputs(s: str) void = {
	tstfdwriteall(os.STDOUT_FILENO, s.ptr, s.len: u64);
};

fn tsterrputs(s: str) void = {
	tstfdwriteall(os.STDERR_FILENO, s.ptr, s.len: u64);
};

fn tstputuint(n: i32) void = {
	let buf: [16]u8;
	let i: i32 = 16;
	let v: i32 = n;
	if (v == 0) {
		i -= 1;
		buf[i] = 48u8;
	} else {
		for (v > 0) {
			i -= 1;
			buf[i] = (48 + v % 10): u8;
			v = v / 10;
		};
	};
	tstfdwriteall(os.STDOUT_FILENO, &buf[i], (16 - i): u64);
};