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

#pragma once

#include <zencore/trace.h>

#if ZEN_PLATFORM_WINDOWS
#	include <intrin.h>

#	define PLATFORM_RETURN_ADDRESS()					 _ReturnAddress()
#	define PLATFORM_RETURN_ADDRESS_POINTER()			 _AddressOfReturnAddress()
#	define PLATFORM_RETURN_ADDRESS_FOR_CALLSTACKTRACING PLATFORM_RETURN_ADDRESS_POINTER
#endif

////////////////////////////////////////////////////////////////////////////////
#if !defined(UE_CALLSTACK_TRACE_ENABLED)
#	if ZEN_WITH_TRACE
#		if ZEN_PLATFORM_WINDOWS
#			define UE_CALLSTACK_TRACE_ENABLED 1
#		endif
#	endif
#endif

#if !defined(UE_CALLSTACK_TRACE_ENABLED)
#	define UE_CALLSTACK_TRACE_ENABLED 0
#endif

////////////////////////////////////////////////////////////////////////////////
#if UE_CALLSTACK_TRACE_ENABLED

#	include "platformtls.h"

namespace zen {

/**
 * Creates callstack tracing.
 * @param Malloc Allocator instance to use.
 */
void CallstackTrace_Create(class FMalloc* Malloc);

/**
 * Initializes callstack tracing. On some platforms this has to be delayed due to initialization order.
 */
void CallstackTrace_Initialize();

/**
 * Capture the current callstack, and trace the definition if it has not already been encountered. The returned value
 * can be used in trace events and be resolved in analysis.
 * @return Unique id identifying the current callstack.
 */
uint32_t CallstackTrace_GetCurrentId();

/**
 * Callstack Trace Scoped Macro to avoid resolving the full callstack
 * can be used when some external libraries are not compiled with frame pointers
 * preventing us to resolve it without crashing. Instead the callstack will be
 * only the caller address.
 */
#	define CALLSTACK_TRACE_LIMIT_CALLSTACKRESOLVE_SCOPE() FCallStackTraceLimitResolveScope PREPROCESSOR_JOIN(FCTLMScope, __LINE__)

extern uint32_t GCallStackTracingTlsSlotIndex;

/**
 * @return the fallback callstack address
 */
inline void*
CallstackTrace_GetFallbackPlatformReturnAddressData()
{
	if (FPlatformTLS::IsValidTlsSlot(GCallStackTracingTlsSlotIndex))
		return FPlatformTLS::GetTlsValue(GCallStackTracingTlsSlotIndex);
	else
		return nullptr;
}

/**
 * @return Needs full callstack resolve
 */
inline bool
CallstackTrace_ResolveFullCallStack()
{
	return CallstackTrace_GetFallbackPlatformReturnAddressData() == nullptr;
}

/*
 * Callstack Trace scope for override CallStack
 */
class FCallStackTraceLimitResolveScope
{
public:
	ZEN_FORCENOINLINE FCallStackTraceLimitResolveScope()
	{
		if (FPlatformTLS::IsValidTlsSlot(GCallStackTracingTlsSlotIndex))
		{
			FPlatformTLS::SetTlsValue(GCallStackTracingTlsSlotIndex, PLATFORM_RETURN_ADDRESS_FOR_CALLSTACKTRACING());
		}
	}

	ZEN_FORCENOINLINE ~FCallStackTraceLimitResolveScope()
	{
		if (FPlatformTLS::IsValidTlsSlot(GCallStackTracingTlsSlotIndex))
		{
			FPlatformTLS::SetTlsValue(GCallStackTracingTlsSlotIndex, nullptr);
		}
	}
};

}  // namespace zen

#else  // UE_CALLSTACK_TRACE_ENABLED

namespace zen {

inline void
CallstackTrace_Create(class FMalloc* /*Malloc*/)
{
}

inline void
CallstackTrace_Initialize()
{
}

inline uint32_t
CallstackTrace_GetCurrentId()
{
	return 0;
}

inline void*
CallstackTrace_GetCurrentReturnAddressData()
{
	return nullptr;
}

inline void*
CallstackTrace_GetFallbackPlatformReturnAddressData()
{
	return nullptr;
}

inline bool
CallstackTrace_ResolveFullCallStack()
{
	return true;
}

#	define CALLSTACK_TRACE_LIMIT_CALLSTACKRESOLVE_SCOPE()

}  // namespace zen

#endif	// UE_CALLSTACK_TRACE_ENABLED