// Copyright Epic Games, Inc. All Rights Reserved.

#define ZEN_TEST_WITH_RUNNER 1

#include "zencore/testing.h"

#include "zencore/filesystem.h"
#include "zencore/logging.h"
#include "zencore/process.h"
#include "zencore/trace.h"

#if ZEN_WITH_TESTS

#	include <zencore/callstack.h>

#	include <chrono>
#	include <clocale>
#	include <csignal>
#	include <cstdlib>
#	include <cstdio>
#	include <string>
#	include <vector>

#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
#		include <execinfo.h>
#		include <unistd.h>
#	endif

namespace zen::testing {

using namespace std::literals;

static void
PrintCrashCallstack([[maybe_unused]] const char* SignalName)
{
#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	// Use write() + backtrace_symbols_fd() which are async-signal-safe
	write(STDERR_FILENO, "\n*** Caught ", 12);
	write(STDERR_FILENO, SignalName, strlen(SignalName));
	write(STDERR_FILENO, " — callstack:\n", 15);

	void* Frames[64];
	int	  FrameCount = backtrace(Frames, 64);
	backtrace_symbols_fd(Frames, FrameCount, STDERR_FILENO);
#	elif ZEN_PLATFORM_WINDOWS
	// On Windows we're called from SEH, not a signal handler, so heap/locks are safe
	void*	 Addresses[64];
	uint32_t FrameCount = GetCallstack(2, 64, Addresses);
	if (FrameCount > 0)
	{
		std::vector<std::string> Symbols = GetFrameSymbols(FrameCount, Addresses);
		fprintf(stderr, "\n*** Caught %s - callstack:\n", SignalName);
		for (uint32_t i = 0; i < FrameCount; ++i)
		{
			fprintf(stderr, "  %2u: %s\n", i, Symbols[i].c_str());
		}
	}
#	endif
}

#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC

static void
CrashSignalHandler(int Signal)
{
	const char* SignalName = "Unknown signal";
	switch (Signal)
	{
		case SIGSEGV:
			SignalName = "SIGSEGV";
			break;
		case SIGABRT:
			SignalName = "SIGABRT";
			break;
		case SIGFPE:
			SignalName = "SIGFPE";
			break;
		case SIGBUS:
			SignalName = "SIGBUS";
			break;
		case SIGILL:
			SignalName = "SIGILL";
			break;
	}

	PrintCrashCallstack(SignalName);

	// Re-raise with default handler so the process terminates normally
	signal(Signal, SIG_DFL);
	raise(Signal);
}

#	endif	// ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC

#	if ZEN_PLATFORM_WINDOWS

static LONG CALLBACK
CrashVectoredHandler(PEXCEPTION_POINTERS ExceptionInfo)
{
	// Only handle fatal exceptions, not first-chance exceptions used for normal control flow
	switch (ExceptionInfo->ExceptionRecord->ExceptionCode)
	{
		case EXCEPTION_ACCESS_VIOLATION:
			PrintCrashCallstack("EXCEPTION_ACCESS_VIOLATION");
			break;
		case EXCEPTION_STACK_OVERFLOW:
			PrintCrashCallstack("EXCEPTION_STACK_OVERFLOW");
			break;
		case EXCEPTION_ILLEGAL_INSTRUCTION:
			PrintCrashCallstack("EXCEPTION_ILLEGAL_INSTRUCTION");
			break;
		case EXCEPTION_INT_DIVIDE_BY_ZERO:
			PrintCrashCallstack("EXCEPTION_INT_DIVIDE_BY_ZERO");
			break;
		default:
			break;
	}

	// Continue search so doctest's handler can report the test case context
	return EXCEPTION_CONTINUE_SEARCH;
}

#	endif	// ZEN_PLATFORM_WINDOWS

static void
InstallCrashSignalHandlers()
{
#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	signal(SIGSEGV, CrashSignalHandler);
	signal(SIGABRT, CrashSignalHandler);
	signal(SIGFPE, CrashSignalHandler);
	signal(SIGBUS, CrashSignalHandler);
	signal(SIGILL, CrashSignalHandler);
#	elif ZEN_PLATFORM_WINDOWS
	AddVectoredExceptionHandler(0 /*called last among vectored handlers*/, CrashVectoredHandler);
#	endif
}

struct TestListener : public doctest::IReporter
{
	const std::string_view ColorYellow = "\033[0;33m"sv;
	const std::string_view ColorNone   = "\033[0m"sv;

	// constructor has to accept the ContextOptions by ref as a single argument
	TestListener(const doctest::ContextOptions&) {}

	void report_query(const doctest::QueryData& /*in*/) override {}

	void test_run_start() override { RunStart = std::chrono::steady_clock::now(); }

	void test_run_end(const doctest::TestRunStats& in) override
	{
		auto   elapsed		  = std::chrono::steady_clock::now() - RunStart;
		double elapsedSeconds = std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count() / 1000.0;

		// Write machine-readable summary to file if requested (used by xmake test summary table)
		const char* summaryFile = std::getenv("ZEN_TEST_SUMMARY_FILE");
		if (summaryFile && summaryFile[0] != '\0')
		{
			if (FILE* f = std::fopen(summaryFile, "w"))
			{
				std::fprintf(f,
							 "cases_total=%u\ncases_passed=%u\nassertions_total=%d\nassertions_passed=%d\n"
							 "elapsed_seconds=%.3f\n",
							 in.numTestCasesPassingFilters,
							 in.numTestCasesPassingFilters - in.numTestCasesFailed,
							 in.numAsserts,
							 in.numAsserts - in.numAssertsFailed,
							 elapsedSeconds);
				for (const auto& failure : FailedTests)
				{
					std::fprintf(f, "failed=%s|%s|%u\n", failure.Name.c_str(), failure.File.c_str(), failure.Line);
				}
				std::fclose(f);
			}
		}
	}

	void test_case_start(const doctest::TestCaseData& in) override
	{
		Current = &in;
		ZEN_CONSOLE("{}======== TEST_CASE: {:<50} ========{}", ColorYellow, Current->m_name, ColorNone);
	}

	// called when a test case is reentered because of unfinished subcases
	void test_case_reenter(const doctest::TestCaseData& /*in*/) override
	{
		ZEN_CONSOLE("{}-------------------------------------------------------------------------------{}", ColorYellow, ColorNone);
	}

	void test_case_end(const doctest::CurrentTestCaseStats& in) override
	{
		if (!in.testCaseSuccess && Current)
		{
			FailedTests.push_back({Current->m_name, Current->m_file.c_str(), Current->m_line});
		}
		Current = nullptr;
	}

	void test_case_exception(const doctest::TestCaseException& /*in*/) override {}

	void subcase_start(const doctest::SubcaseSignature& in) override
	{
		ZEN_CONSOLE("{}-------- SUBCASE:   {:<50} --------{}",
					ColorYellow,
					fmt::format("{}/{}", Current->m_name, in.m_name.c_str()),
					ColorNone);
	}

	void subcase_end() override {}

	void log_assert(const doctest::AssertData& /*in*/) override {}

	void log_message(const doctest::MessageData& /*in*/) override {}

	void test_case_skipped(const doctest::TestCaseData& /*in*/) override {}

	const doctest::TestCaseData*		  Current  = nullptr;
	std::chrono::steady_clock::time_point RunStart = {};

	struct FailedTestInfo
	{
		std::string Name;
		std::string File;
		unsigned	Line;
	};
	std::vector<FailedTestInfo> FailedTests;
};

struct TestRunner::Impl
{
	Impl() { REGISTER_LISTENER("ZenTestListener", 1, TestListener); }
	doctest::Context Session;
};

TestRunner::TestRunner()
{
	m_Impl = std::make_unique<Impl>();
}

TestRunner::~TestRunner()
{
}

void
TestRunner::SetDefaultSuiteFilter(const char* Pattern)
{
	m_Impl->Session.setOption("test-suite", Pattern);
}

int
TestRunner::ApplyCommandLine(int Argc, char const* const* Argv)
{
	m_Impl->Session.applyCommandLine(Argc, Argv);

	for (int i = 1; i < Argc; ++i)
	{
		if (Argv[i] == "--debug"sv)
		{
			zen::logging::SetLogLevel(zen::logging::Debug);
		}
		else if (Argv[i] == "--verbose"sv)
		{
			zen::logging::SetLogLevel(zen::logging::Trace);
		}
	}

	return 0;
}

int
TestRunner::Run()
{
	return m_Impl->Session.run();
}

int
RunTestMain(int Argc, char* Argv[], const char* ExecutableName, void (*ForceLink)())
{
#	if ZEN_PLATFORM_WINDOWS
	setlocale(LC_ALL, "en_us.UTF8");
#	endif

	ForceLink();

#	if ZEN_PLATFORM_LINUX
	zen::IgnoreChildSignals();
#	endif

#	if ZEN_WITH_TRACE
	zen::TraceInit(ExecutableName);
	zen::TraceOptions TraceCommandlineOptions;
	if (GetTraceOptionsFromCommandline(TraceCommandlineOptions))
	{
		TraceConfigure(TraceCommandlineOptions);
	}
#	endif

	zen::logging::InitializeLogging();
	zen::MaximizeOpenFileCount();
	InstallCrashSignalHandlers();

	TestRunner Runner;

	// Derive default suite filter from ExecutableName: "zencore-test" -> "core.*"
	if (ExecutableName)
	{
		std::string_view Name = ExecutableName;
		if (Name.starts_with("zen"))
		{
			Name.remove_prefix(3);
		}
		if (Name.ends_with("-test"))
		{
			Name.remove_suffix(5);
		}
		if (!Name.empty())
		{
			std::string Filter(Name);
			Filter += ".*";
			Runner.SetDefaultSuiteFilter(Filter.c_str());
		}
	}

	Runner.ApplyCommandLine(Argc, Argv);
	return Runner.Run();
}

}  // namespace zen::testing

#endif	// ZEN_WITH_TESTS