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

#include <zencore/crashhandler.h>

#include <stdio.h>
#include <string.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#	include <DbgHelp.h>
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
#	include <execinfo.h>
#	include <signal.h>
#	include <unistd.h>
#endif

namespace zen {

#if ZEN_PLATFORM_WINDOWS

static LONG WINAPI
CrashExceptionFilter(PEXCEPTION_POINTERS ExceptionInfo)
{
	const char* ExceptionName = nullptr;

	switch (ExceptionInfo->ExceptionRecord->ExceptionCode)
	{
		case EXCEPTION_ACCESS_VIOLATION:
			ExceptionName = "EXCEPTION_ACCESS_VIOLATION";
			break;
		case EXCEPTION_ILLEGAL_INSTRUCTION:
			ExceptionName = "EXCEPTION_ILLEGAL_INSTRUCTION";
			break;
		case EXCEPTION_INT_DIVIDE_BY_ZERO:
			ExceptionName = "EXCEPTION_INT_DIVIDE_BY_ZERO";
			break;
		case EXCEPTION_ARRAY_BOUNDS_EXCEEDED:
			ExceptionName = "EXCEPTION_ARRAY_BOUNDS_EXCEEDED";
			break;
		case EXCEPTION_FLT_DIVIDE_BY_ZERO:
			ExceptionName = "EXCEPTION_FLT_DIVIDE_BY_ZERO";
			break;
		default:
			return EXCEPTION_CONTINUE_SEARCH;
	}

	fprintf(stderr,
			"\n*** FATAL: %s at address 0x%p (pid %lu)\n",
			ExceptionName,
			ExceptionInfo->ExceptionRecord->ExceptionAddress,
			GetCurrentProcessId());

	// Capture backtrace from the exception context using StackWalk64 so we get
	// the actual crash site rather than the exception dispatch frames

	HANDLE Process = GetCurrentProcess();
	HANDLE Thread  = GetCurrentThread();

	// SymInitialize is safe to call if already initialized — it returns FALSE
	// but existing state remains valid for SymFromAddr calls
	SymInitialize(Process, NULL, TRUE);

	CONTEXT Context = *ExceptionInfo->ContextRecord;

	STACKFRAME64 StackFrame;
	memset(&StackFrame, 0, sizeof(StackFrame));

#	if ZEN_ARCH_X64
	DWORD MachineType			= IMAGE_FILE_MACHINE_AMD64;
	StackFrame.AddrPC.Offset	= Context.Rip;
	StackFrame.AddrPC.Mode		= AddrModeFlat;
	StackFrame.AddrFrame.Offset = Context.Rbp;
	StackFrame.AddrFrame.Mode	= AddrModeFlat;
	StackFrame.AddrStack.Offset = Context.Rsp;
	StackFrame.AddrStack.Mode	= AddrModeFlat;
#	elif ZEN_ARCH_ARM64
	DWORD MachineType			= IMAGE_FILE_MACHINE_ARM64;
	StackFrame.AddrPC.Offset	= Context.Pc;
	StackFrame.AddrPC.Mode		= AddrModeFlat;
	StackFrame.AddrFrame.Offset = Context.Fp;
	StackFrame.AddrFrame.Mode	= AddrModeFlat;
	StackFrame.AddrStack.Offset = Context.Sp;
	StackFrame.AddrStack.Mode	= AddrModeFlat;
#	endif

	char		 SymbolBuffer[sizeof(SYMBOL_INFO) + 1024];
	SYMBOL_INFO* Symbol	 = (SYMBOL_INFO*)SymbolBuffer;
	Symbol->SizeOfStruct = sizeof(SYMBOL_INFO);
	Symbol->MaxNameLen	 = 1023;

	fprintf(stderr, "Backtrace:\n");

	for (int FrameIndex = 0; FrameIndex < 64; ++FrameIndex)
	{
		if (!StackWalk64(MachineType, Process, Thread, &StackFrame, &Context, NULL, SymFunctionTableAccess64, SymGetModuleBase64, NULL))
		{
			break;
		}

		DWORD64 Address = StackFrame.AddrPC.Offset;
		if (Address == 0)
		{
			break;
		}

		DWORD64 Displacement = 0;
		if (SymFromAddr(Process, Address, &Displacement, Symbol))
		{
			fprintf(stderr,
					"  #%-2d %s+0x%llx [0x%llx]\n",
					FrameIndex,
					Symbol->Name,
					(unsigned long long)Displacement,
					(unsigned long long)Address);
		}
		else
		{
			fprintf(stderr, "  #%-2d [0x%llx]\n", FrameIndex, (unsigned long long)Address);
		}
	}

	fprintf(stderr, "\n");
	fflush(stderr);

	return EXCEPTION_CONTINUE_SEARCH;
}

void
InstallCrashHandler()
{
	SetUnhandledExceptionFilter(CrashExceptionFilter);
}

#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC

// Async-signal-safe helper: write a decimal number to a file descriptor
static void
WriteDecimal(int Fd, unsigned long Value)
{
	char Buffer[20];
	int	 Pos = sizeof(Buffer);

	if (Value == 0)
	{
		Buffer[--Pos] = '0';
	}
	else
	{
		while (Value > 0)
		{
			Buffer[--Pos] = '0' + (char)(Value % 10);
			Value /= 10;
		}
	}

	write(Fd, Buffer + Pos, sizeof(Buffer) - Pos);
}

static void
CrashSignalHandler(int Signal)
{
	// Note: backtrace/backtrace_symbols_fd are not strictly async-signal-safe
	// (glibc's backtrace may call malloc internally), but they are widely used
	// in signal handlers in practice. This is best-effort.

	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;
	}

	write(STDERR_FILENO, "\n*** FATAL: Caught ", 19);
	write(STDERR_FILENO, SignalName, strlen(SignalName));
	write(STDERR_FILENO, " (pid ", 6);
	WriteDecimal(STDERR_FILENO, (unsigned long)getpid());
	write(STDERR_FILENO, ")\nBacktrace:\n", 13);

	void* Frames[64];
	int	  FrameCount = backtrace(Frames, 64);
	backtrace_symbols_fd(Frames, FrameCount, STDERR_FILENO);

	write(STDERR_FILENO, "\n", 1);

	// Re-raise with default handler so the process generates a core dump
	// and terminates with the correct exit status
	signal(Signal, SIG_DFL);
	raise(Signal);
}

void
InstallCrashHandler()
{
	struct sigaction Action;
	memset(&Action, 0, sizeof(Action));
	Action.sa_handler = CrashSignalHandler;
	Action.sa_flags	  = SA_RESETHAND;  // one-shot: auto-reset to default after firing
	sigemptyset(&Action.sa_mask);

	sigaction(SIGSEGV, &Action, nullptr);
	sigaction(SIGABRT, &Action, nullptr);
	sigaction(SIGFPE, &Action, nullptr);
	sigaction(SIGBUS, &Action, nullptr);
	sigaction(SIGILL, &Action, nullptr);
}

#endif

}  // namespace zen