#ifndef RT_TEST_H
#define RT_TEST_H

/*
 * rttest.h - Testing library
 * Copyright (C) 2026 Kevin Trogant
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

#include "rtcore.h"

#ifndef RTT_API
#define RTT_API RTC_API
#endif

/* Opaque testing harness. A test driver is the component keeping track of registered tests and suits */
struct test_driver;
typedef struct test_driver test_driver;

/* Helper macro for defining a test function. The suite needs to point to the suite object the test belongs to,
 * or NULL if the test is freestanding.
 */
#define TEST(_Name) void _Name(test_driver *testing, void *suite)
typedef TEST(test_fn);

/* Helper macro for defining a test suite.
 * The suite pointer points to the suite object, which needs to remain valid throughout test execution */
#define SUITE_SETUP(_Name) void _Name(test_driver *testing, void *suite)
/* Helper macro for defining a test suite shutdown. */
#define SUITE_TEARDOWN(_Name) void _Name(test_driver *testing, void *suite)
typedef SUITE_SETUP(suite_setup_fn);
typedef SUITE_TEARDOWN(suite_teardown_fn);

/* Mark the suite parameter of a test as unused.
 * Mostly for documentation purposes, or if you compile with -Wunused-parameter
 */
#define NoSuite Unused(suite)

/* Creates a new test driver. Allocates internal storage with the given size.
 * All internal allocations are made from that storage */
RTT_API test_driver *CreateTestDriver(isize mem_size);

/* Cleanup of the test driver */
RTT_API void DestroyTestDriver(test_driver *testing);

/* Runs the tests.
 * Takes the argc and argv passed to the program. You can strip some prefix of argv. For example, if you only want
 * to pass arguments after some --testing flag, you could do:
 *
 * for (int i = 0; i < argc; ++i) {
 *	 if (strcmp(argv[i], "--testing") == 0) {
 *	   ExecuteTests(testing, argc - i - 1, &argv[i+1]);
 *	 }
 * }
 *
 * The supported arguments are:
 * -t | --test <name>  -  executes only the test with the given name.
 *                        For tests inside suites, the syntax "MySuite.MyTest" is supported.
 * -s | --suite <name> -  executes all tests inside the given test suite.
 *
 * Returns true if all executed tests pass, false if any test failed.
 */
RTT_API b32 ExecuteTests(test_driver *testing, int argc, char **argv);

/* Writes a CMake file with CTest commands for all registered tests to the given output path.
 *
 * The executable is the CMake executable that needs to be executed (you can add arguments to it)
 */
RTT_API b32 WriteCTestFile(test_driver *testing, s8 executable, s8 output_path);

/* Internal utility, you don't need to call this yourself. */
RTT_API void ReportTestFailure(test_driver *testing, s8 file, int line, s8 reason);

RTT_API void RegisterSuite_(test_driver *testing, void *suite, suite_setup_fn setup, suite_teardown_fn teardown, s8 name);
RTT_API void RegisterTest_(test_driver *testing, void *suite, test_fn fn, s8 name);

/* Registers a new test suite.
 * Usage:
 *    SUITE_SETUP(SetupMySuite) { ... }
 *    SUITE_TEARDOWN(TeardownMySuite) { ... }
 *
 *    static suite_type my_suite;
 *    RegisterSuite(testing, my_suite, SetupMySuite, TeardownMySuite)
 * This registers a suite with name "my_suite"
 */
#define RegisterSuite(_Testing, _Suite, _Setup, _Teardown) RegisterSuite_(_Testing, &_Suite, _Setup, _Teardown, S8( # _Suite))

/* Registers a new test inside a suite
 * Usage:
 *    TEST(MyTest) { ... }
 *
 *    RegisterTest(testing, my_suite, MyTest)
 * This registers a test with name "MyTest" inside the given suite.
 */
#define RegisterTest(_Testing, _Suite, _Fn) RegisterTest_(_Testing, &_Suite, _Fn, S8( # _Fn))

/* Registers a new test without a suite
 * Usage:
 *    TEST(MyTest) { ... }
 *
 *    RegisterTest(testing, MyTest)
 * This registers a test with name "MyTest"
 */
#define RegisterStandaloneTest(_Testing, _Fn) RegisterTest_(_Testing, NULL, _Fn, S8( # _Fn))

/* Assertion macros.
 * These report a failure if the condition is not true and abort the test.
 */

#define AssertTrue(_Testing, _Expr) { if (!(_Expr)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AssertTrue(" # _Expr ")" ) ); return; } }
#define AssertFalse(_Testing, _Expr) { if ((_Expr)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AssertFalse(" # _Expr ")" ) ); return; } }
#define AssertEqual(_Testing, _A, _B) { if ((_A) != (_B)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AssertEqual(" # _A ", " # _B ")" ) ); return; } }
#define AssertNotEqual(_Testing, _A, _B) { if ((_A) == (_B)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AsserNotEqual(" # _A ", " # _B ")" ) ); return; } }
#define AssertNotNull(_Testing, _Ptr) { if ((_Ptr) == NULL) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AssertNotNull(" # _Ptr ")" ) ); return; } }
#define AssertNull(_Testing, _Ptr) { if ((_Ptr) != NULL) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AssertNull(" # _Ptr ")" ) ); return; } }

/* Expect macros.
 * These report a failure if the condition is not true, but do not abort the test.
 */

#define ExpectTrue(_Testing, _Expr) { if (!(_Expr)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "ExpectTrue(" # _Expr ")" ) ); } }
#define ExpectFalse(_Testing, _Expr) { if ((_Expr)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "ExpectFalse(" # _Expr ")" ) ); } }
#define ExpectEqual(_Testing, _A, _B) { if ((_A) != (_B)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "ExpectEqual(" # _A ", " # _B ")" ) ); } }
#define ExpectNotEqual(_Testing, _A, _B) { if ((_A) == (_B)) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "AsserNotEqual(" # _A ", " # _B ")" ) ); } }
#define ExpectNotNull(_Testing, _Ptr) { if ((_Ptr) == NULL) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "ExpectNotNull(" # _Ptr ")" ) ); } }
#define ExpectNull(_Testing, _Ptr) { if ((_Ptr) != NULL) { ReportTestFailure((_Testing), S8(__FILE__), __LINE__, S8( "ExpectNull(" # _Ptr ")" ) ); } }

#endif

#ifdef RT_TEST_IMPLEMENTATION
#undef RT_TEST_IMPLEMENTATION

#include <stdlib.h>
#include <stdio.h>

typedef struct
{
	void *ptr;
	suite_setup_fn *setup;
	suite_teardown_fn *teardown;
	s8 name;
} suite;

typedef struct
{
	void *suite_ptr;
	test_fn *fn;
	s8 name;
} test;

struct test_driver
{
	suite *suites;
	test *tests;
	arena allocator;
	void *mem;
	int suite_capacity;
	int num_suites;
	int test_capacity;
	int num_tests;
	b32 failed;
};

test_driver *
CreateTestDriver(isize mem_size)
{
	void *mem = malloc((size_t)mem_size);
	arena allocator = MakeArena(mem, mem_size);
	test_driver *driver = Alloc(&allocator, test_driver);
	driver->allocator = allocator;
	driver->mem = mem;
	return driver;
}

void
DestroyTestDriver(test_driver *testing)
{
	free(testing->mem);
}

internal int
SortSuites(const void *a, const void *b)
{
	const suite *sa = a, *sb = b;
	if ((uintptr_t)sa->ptr < (uintptr_t)sb->ptr)
		return -1;
	if ((uintptr_t)sa->ptr > (uintptr_t)sb->ptr)
		return 1;
	return 0;
}

internal int
SearchSuite(const void *key, const void *candidate)
{
	const suite *c = candidate;
	if ((uintptr_t)key < (uintptr_t)c->ptr)
		return -1;
	if ((uintptr_t)key > (uintptr_t)c->ptr)
		return 1;
	return 0;
}

internal int
SortTests(const void *a, const void *b)
{
	const test *sa = a, *sb = b;
	if ((uintptr_t)sa->suite_ptr < (uintptr_t)sb->suite_ptr)
		return -1;
	if ((uintptr_t)sa->suite_ptr > (uintptr_t)sb->suite_ptr)
		return 1;
	return 0;
}

RTT_API b32
ExecuteTests(test_driver *testing, int argc, char **argv)
{
	s8 test_name = {0}, suite_name = {0};
	if (argc > 0)
	{
		for (int i = 0; i < argc; ++i)
		{
			if ((strcmp(argv[i], "--test") == 0) || (strcmp(argv[i], "-t") == 0))
			{
				if (i == argc - 1)
				{
					fprintf(stderr, "Expected test name after argument --%s\n", argv[i]);
					break;
				}
				test_name = (s8){.data = (u8 *)argv[i + 1], strlen(argv[i + 1])};
				break;
			}
			else if ((strcmp(argv[i], "--suite") == 0) || (strcmp(argv[i], "-s") == 0))
			{
				if (i == argc - 1)
				{
					fprintf(stderr, "Expected suite name after argument --%s\n", argv[i]);
					break;
				}
				suite_name = (s8){.data = (u8 *)argv[i + 1], strlen(argv[i + 1])};
				break;
			}
		}
	}

	/* Sort suites and tests by suite ptr */
	qsort(testing->suites, testing->num_suites, sizeof(suite), SortSuites);
	qsort(testing->tests, testing->num_tests, sizeof(test), SortTests);

	int success = 1;

	test *tests = testing->tests;

	/* Run all */
	if (!test_name.data && !suite_name.data)
	{
		/* Execute tests not attached to a suite */
		int test_idx = 0;
		for (; test_idx < testing->num_tests; ++test_idx)
		{
			if (tests[test_idx].suite_ptr != NULL)
			{
				break;
			}
			testing->failed = 0;
			printf("[<NONE>.%.*s]\t...\n", (int)tests[test_idx].name.length, (const char *)tests[test_idx].name.data);
			tests[test_idx].fn(testing, tests[test_idx].suite_ptr);
			printf("[<NONE>.%.*s]\t%s\n",
			       (int)tests[test_idx].name.length,
			       (const char *)tests[test_idx].name.data,
			       testing->failed ? "FAIL" : "OK");
			if (testing->failed)
				success = 0;
		}

		suite *cur_suite = NULL;
		if (test_idx < testing->num_tests)
		{
			cur_suite = &testing->suites[0];
			printf("[%.*s] SETUP\t...\n", (int)cur_suite->name.length, (const char *)cur_suite->name.data);
			cur_suite->setup(testing, cur_suite->ptr);
			printf("[%.*s] SETUP\t%s\n",
			       (int)cur_suite->name.length,
			       (const char *)cur_suite->name.data,
			       testing->failed ? "FAIL" : "OK");
			if (testing->failed)
			{
				return 0;
			}
		}
		for (; test_idx < testing->num_tests; ++test_idx)
		{
			suite *old_suite = cur_suite;
			while (cur_suite && cur_suite->ptr != tests[test_idx].suite_ptr)
			{
				++cur_suite;
			}
			if (!cur_suite)
			{
				fprintf(stderr,
				        "Failed to find suite for test %.*s. Make sure that the suite was registered.\n",
				        (int)tests[test_idx].name.length,
				        (char *)tests[test_idx].name.data);

				return 0;
			}
			if (old_suite && old_suite != cur_suite)
			{
				old_suite->teardown(testing, old_suite->ptr);

				testing->failed = 0;
				printf("[%.*s] SETUP\t...\n", (int)cur_suite->name.length, (const char *)cur_suite->name.data);
				cur_suite->setup(testing, cur_suite->ptr);
				printf("[%.*s] SETUP\t%s\n",
				       (int)cur_suite->name.length,
				       (const char *)cur_suite->name.data,
				       testing->failed ? "FAIL" : "OK");
				if (testing->failed)
				{
					success = 0;
					break;
				}
			}
			testing->failed = 0;
			printf("[%.*s.%.*s]\t...\n",
			       (int)cur_suite->name.length,
			       (const char *)cur_suite->name.data,
			       (int)tests[test_idx].name.length,
			       (const char *)tests[test_idx].name.data);
			tests[test_idx].fn(testing, tests[test_idx].suite_ptr);
			printf("[%.*s.%.*s]\t%s\n",
			       (int)cur_suite->name.length,
			       (const char *)cur_suite->name.data,
			       (int)tests[test_idx].name.length,
			       (const char *)tests[test_idx].name.data,
			       testing->failed ? "FAIL" : "OK");
			if (testing->failed)
				success = 0;
		}

		if (cur_suite)
			cur_suite->teardown(testing, cur_suite->ptr);
	}
	else if (suite_name.data)
	{
		/* Run only one suite */
		suite *s = NULL;
		for (int i = 0; i < testing->num_suites; ++i)
		{
			if (S8Equals(testing->suites[i].name, suite_name))
			{
				s = &testing->suites[i];
				break;
			}
		}
		if (!s)
		{
			fprintf(stderr,
			        "Failed to find suite %.*s. Make sure that the suite was registered.\n",
			        (int)suite_name.length,
			        (char *)suite_name.data);

			return 0;
		}

		printf("[%.*s] SETUP\t...\n", (int)s->name.length, (const char *)s->name.data);
		s->setup(testing, s->ptr);
		printf("[%.*s] SETUP\t%s\n", (int)s->name.length, (const char *)s->name.data, testing->failed ? "FAIL" : "OK");
		if (testing->failed)
		{
			return 0;
		}

		for (int test_idx = 0; test_idx < testing->num_tests; ++test_idx)
		{
			if (tests[test_idx].suite_ptr != s->ptr)
				continue;

			testing->failed = 0;
			printf("[%.*s.%.*s]\t...\n",
			       (int)s->name.length,
			       (const char *)s->name.data,
			       (int)tests[test_idx].name.length,
			       (const char *)tests[test_idx].name.data);
			tests[test_idx].fn(testing, tests[test_idx].suite_ptr);
			printf("[%.*s.%.*s]\t%s\n",
			       (int)s->name.length,
			       (const char *)s->name.data,
			       (int)tests[test_idx].name.length,
			       (const char *)tests[test_idx].name.data,
			       testing->failed ? "FAIL" : "OK");
			if (testing->failed)
				success = 0;
		}

		s->teardown(testing, s->ptr);
	}
	else if (test_name.data)
	{
		/* Run one specific test.
		 *
		 * The syntax is either "Suite.Test" or just "Test"
		 */
		u8 *dot = S8Chr(test_name, '.');
		if (dot)
		{
			suite_name = S8Span(test_name.data, dot);
			test_name = S8Span(dot+1, S8End(test_name)+1);
		}

		suite *s = NULL;
		if (suite_name.data)
		{
			for (int i = 0; i < testing->num_suites; ++i)
			{
				if (S8Equals(testing->suites[i].name, suite_name))
				{
					s = &testing->suites[i];
					break;
				}
			}
		}

		for (int test_idx = 0; test_idx < testing->num_tests; ++test_idx)
		{
			if (s && tests[test_idx].suite_ptr != s->ptr)
				continue;
			if (!S8Equals(tests[test_idx].name, test_name))
				continue;

			if (tests[test_idx].suite_ptr)
			{
				if (!s)
				{
					s = bsearch(tests[test_idx].suite_ptr,
					            testing->suites,
					            testing->num_suites,
					            sizeof(suite),
					            SearchSuite);
					if (!s)
					{
						fprintf(stderr,
						        "Failed to find suite for test %.*s. Make sure that the suite was registered.\n",
						        (int)tests[test_idx].name.length,
						        (char *)tests[test_idx].name.data);

						return 0;
					}
				}

				testing->failed = 0;
				printf("[%.*s] SETUP\t...\n", (int)s->name.length, (const char *)s->name.data);
				s->setup(testing, s->ptr);
				printf("[%.*s] SETUP\t%s\n",
				       (int)s->name.length,
				       (const char *)s->name.data,
				       testing->failed ? "FAIL" : "OK");
				if (testing->failed)
				{
					success = 0;
					break;
				}

				testing->failed = 0;
				printf("[%.*s.%.*s]\t...\n",
				       (int)s->name.length,
				       (const char *)s->name.data,
				       (int)tests[test_idx].name.length,
				       (const char *)tests[test_idx].name.data);
				tests[test_idx].fn(testing, tests[test_idx].suite_ptr);
				printf("[%.*s.%.*s]\t%s\n",
				       (int)s->name.length,
				       (const char *)s->name.data,
				       (int)tests[test_idx].name.length,
				       (const char *)tests[test_idx].name.data,
				       testing->failed ? "FAIL" : "OK");
				if (testing->failed)
					success = 0;
			}
			else
			{
				testing->failed = 0;
				printf("[<NONE>.%.*s]\t...\n",
				       (int)tests[test_idx].name.length,
				       (const char *)tests[test_idx].name.data);
				tests[test_idx].fn(testing, tests[test_idx].suite_ptr);
				printf("[<NONE>.%.*s]\t%s\n",
				       (int)tests[test_idx].name.length,
				       (const char *)tests[test_idx].name.data,
				       testing->failed ? "FAIL" : "OK");
				if (testing->failed)
					success = 0;
			}
			break;
		}
	}

	return success;
}

RTT_API b32
WriteCTestFile(test_driver *testing, s8 executable, s8 output_path)
{
	char p[260];
	memcpy(p, output_path.data, Min(CountOf(p) - 1, output_path.length));
	p[Min(CountOf(p) - 1, output_path.length)] = 0;
	FILE *f = fopen(p, "w");
	if (!f)
	{
		fprintf(stderr, "Failed to open %s for writing.\n", p);
		return 0;
	}

	b32 success = 1;

	for (int test_idx = 0; test_idx < testing->num_tests; ++test_idx)
	{
		test *t = &testing->tests[test_idx];
		suite *s = NULL;
		if (t->suite_ptr)
		{
			s = bsearch(testing->tests[test_idx].suite_ptr,
						testing->suites,
						testing->num_suites,
						sizeof(suite),
						SearchSuite);
			if (!s)
			{
				fprintf(stderr,
						"Failed to find suite for test %.*s. Make sure that the suite was registered.\n",
						(int)testing->tests[test_idx].name.length,
						(char *)testing->tests[test_idx].name.data);
				success = 0;
				goto out;
			}
		}

		if (s)
		{
			fprintf(f,
					"add_test(NAME %.*s_%.*s COMMAND %.*s --test %.*s.%.*s WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})\n",
					(int)s->name.length,
					(const char *)s->name.data,
					(int)t->name.length,
					(const char *)t->name.data,
					(int)executable.length,
					(const char *)executable.data,
					(int)s->name.length,
					(const char *)s->name.data,
					(int)t->name.length,
					(const char *)t->name.data);
		}
		else
		{
			fprintf(f,
					"add_test(NAME %.*s COMMAND %.*s --test %.*s WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})\n",
					(int)t->name.length,
					(const char *)t->name.data,
					(int)executable.length,
					(const char *)executable.data,
					(int)t->name.length,
					(const char *)t->name.data);
		}
	}
out:
	fclose(f);
	return success;
}

RTT_API void
RegisterSuite_(test_driver *testing, void *suite_ptr, suite_setup_fn setup, suite_teardown_fn teardown, s8 name)
{
	if (testing->num_suites == testing->suite_capacity)
	{
		suite *tmp = AllocArray(&testing->allocator, suite, Max(2 * testing->suite_capacity, 64));
		if (testing->num_suites > 0)
			memcpy(tmp, testing->suites, sizeof(suite) * testing->num_suites);
		testing->suites = tmp;
	}
	testing->suites[testing->num_suites].name = name;
	testing->suites[testing->num_suites].ptr = suite_ptr;
	testing->suites[testing->num_suites].setup = setup;
	testing->suites[testing->num_suites].teardown = teardown;
	++testing->num_suites;
}

RTT_API void
RegisterTest_(test_driver *testing, void *suite, test_fn fn, s8 name)
{
	if (testing->num_tests == testing->test_capacity)
	{
		test *tmp = AllocArray(&testing->allocator, test, Max(2 * testing->test_capacity, 64));
		if (testing->num_tests > 0)
			memcpy(tmp, testing->tests, sizeof(test) * testing->num_tests);
		testing->tests = tmp;
	}
	testing->tests[testing->num_tests].name = name;
	testing->tests[testing->num_tests].suite_ptr = suite;
	testing->tests[testing->num_tests].fn = fn;
	++testing->num_tests;
}

RTT_API void
ReportTestFailure(test_driver *testing, s8 file, int line, s8 reason)
{
	Assert(testing);
	testing->failed = 1;
	printf("%.*s:%d failed: %.*s\n",
	       (int)file.length,
	       (char *)file.data,
	       line,
	       (int)reason.length,
	       (char *)reason.data);
}

#endif