/* SPDX-License-Identifier: GPL-2.0 */
/*
 * Copyright (c) 2024 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2024 David Vernet <dvernet@meta.com>
 * Copyright (c) 2024 Tejun Heo <tj@kernel.org>
 */
#include <stdio.h>
#include <unistd.h>
#include <signal.h>
#include <libgen.h>
#include <bpf/bpf.h>
#include "scx_test.h"

const char help_fmt[] =
"The runner for sched_ext tests.\n"
"\n"
"The runner is statically linked against all testcases, and runs them all serially.\n"
"It's required for the testcases to be serial, as only a single host-wide sched_ext\n"
"scheduler may be loaded at any given time."
"\n"
"Usage: %s [-t TEST] [-s] [-l] [-q]\n"
"\n"
"  -t TEST       Only run tests whose name includes this string\n"
"  -s            Include print output for skipped tests\n"
"  -l            List all available tests\n"
"  -q            Don't print the test descriptions during run\n"
"  -h            Display this help and exit\n";

static volatile int exit_req;
static bool quiet, print_skipped, list;

#define MAX_SCX_TESTS 2048

static struct scx_test __scx_tests[MAX_SCX_TESTS];
static unsigned __scx_num_tests = 0;

static void sigint_handler(int simple)
{
	exit_req = 1;
}

static void print_test_preamble(const struct scx_test *test, bool quiet)
{
	printf("===== START =====\n");
	printf("TEST: %s\n", test->name);
	if (!quiet)
		printf("DESCRIPTION: %s\n", test->description);
	printf("OUTPUT:\n");

	/*
	 * The tests may fork with the preamble buffered
	 * in the children's stdout. Flush before the test
	 * to avoid printing the message multiple times.
	 */
	fflush(stdout);
	fflush(stderr);
}

static const char *status_to_result(enum scx_test_status status)
{
	switch (status) {
	case SCX_TEST_PASS:
	case SCX_TEST_SKIP:
		return "ok";
	case SCX_TEST_FAIL:
		return "not ok";
	default:
		return "<UNKNOWN>";
	}
}

static void print_test_result(const struct scx_test *test,
			      enum scx_test_status status,
			      unsigned int testnum)
{
	const char *result = status_to_result(status);
	const char *directive = status == SCX_TEST_SKIP ? "SKIP " : "";

	printf("%s %u %s # %s\n", result, testnum, test->name, directive);
	printf("=====  END  =====\n");
}

static bool should_skip_test(const struct scx_test *test, const char * filter)
{
	return !strstr(test->name, filter);
}

static enum scx_test_status run_test(const struct scx_test *test)
{
	enum scx_test_status status;
	void *context = NULL;

	if (test->setup) {
		status = test->setup(&context);
		if (status != SCX_TEST_PASS)
			return status;
	}

	status = test->run(context);

	if (test->cleanup)
		test->cleanup(context);

	return status;
}

static bool test_valid(const struct scx_test *test)
{
	if (!test) {
		fprintf(stderr, "NULL test detected\n");
		return false;
	}

	if (!test->name) {
		fprintf(stderr,
			"Test with no name found. Must specify test name.\n");
		return false;
	}

	if (!test->description) {
		fprintf(stderr, "Test %s requires description.\n", test->name);
		return false;
	}

	if (!test->run) {
		fprintf(stderr, "Test %s has no run() callback\n", test->name);
		return false;
	}

	return true;
}

int main(int argc, char **argv)
{
	const char *filter = NULL;
	const char *failed_tests[MAX_SCX_TESTS];
	const char *skipped_tests[MAX_SCX_TESTS];
	unsigned testnum = 0, i;
	unsigned passed = 0, skipped = 0, failed = 0;
	int opt;

	signal(SIGINT, sigint_handler);
	signal(SIGTERM, sigint_handler);

	libbpf_set_strict_mode(LIBBPF_STRICT_ALL);

	while ((opt = getopt(argc, argv, "qslt:h")) != -1) {
		switch (opt) {
		case 'q':
			quiet = true;
			break;
		case 's':
			print_skipped = true;
			break;
		case 'l':
			list = true;
			break;
		case 't':
			filter = optarg;
			break;
		default:
			fprintf(stderr, help_fmt, basename(argv[0]));
			return opt != 'h';
		}
	}

	if (optind < argc) {
		fprintf(stderr, "Unexpected argument '%s'. Use -t to filter tests.\n",
			argv[optind]);
		return 1;
	}

	if (filter) {
		for (i = 0; i < __scx_num_tests; i++) {
			if (!should_skip_test(&__scx_tests[i], filter))
				break;
		}
		if (i == __scx_num_tests) {
			fprintf(stderr, "No tests matched filter '%s'\n", filter);
			fprintf(stderr, "Available tests (use -l to list):\n");
			for (i = 0; i < __scx_num_tests; i++)
				fprintf(stderr, "  %s\n", __scx_tests[i].name);
			return 1;
		}
	}

	for (i = 0; i < __scx_num_tests; i++) {
		enum scx_test_status status;
		struct scx_test *test = &__scx_tests[i];

		if (exit_req)
			break;

		if (list) {
			printf("%s\n", test->name);
			if (i == (__scx_num_tests - 1))
				return 0;
			continue;
		}

		if (filter && should_skip_test(test, filter)) {
			/*
			 * Printing the skipped tests and their preambles can
			 * add a lot of noise to the runner output. Printing
			 * this is only really useful for CI, so let's skip it
			 * by default.
			 */
			if (print_skipped) {
				print_test_preamble(test, quiet);
				print_test_result(test, SCX_TEST_SKIP, ++testnum);
			}
			continue;
		}

		print_test_preamble(test, quiet);
		status = run_test(test);
		print_test_result(test, status, ++testnum);
		switch (status) {
		case SCX_TEST_PASS:
			passed++;
			break;
		case SCX_TEST_SKIP:
			skipped_tests[skipped++] = test->name;
			break;
		case SCX_TEST_FAIL:
			failed_tests[failed++] = test->name;
			break;
		}
	}
	printf("\n\n=============================\n\n");
	printf("RESULTS:\n\n");
	printf("PASSED:  %u\n", passed);
	printf("SKIPPED: %u\n", skipped);
	printf("FAILED:  %u\n", failed);
	if (skipped > 0) {
		printf("\nSkipped tests:\n");
		for (i = 0; i < skipped; i++)
			printf("  - %s\n", skipped_tests[i]);
	}
	if (failed > 0) {
		printf("\nFailed tests:\n");
		for (i = 0; i < failed; i++)
			printf("  - %s\n", failed_tests[i]);
	}

	return failed > 0 ? 1 : 0;
}

void scx_test_register(struct scx_test *test)
{
	SCX_BUG_ON(!test_valid(test), "Invalid test found");
	SCX_BUG_ON(__scx_num_tests >= MAX_SCX_TESTS, "Maximum tests exceeded");

	__scx_tests[__scx_num_tests++] = *test;
}