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

#include "zencore/logging.h"

#include <zencore/logging/ansicolorsink.h>
#include <zencore/logging/logger.h>
#include <zencore/logging/messageonlyformatter.h>
#include <zencore/logging/nullsink.h>
#include <zencore/logging/registry.h>
#include <zencore/string.h>
#include <zencore/testing.h>
#include <zencore/thread.h>

#include <zencore/memory/llm.h>

#if ZEN_PLATFORM_WINDOWS
#	pragma section(".zlog$a", read)
#	pragma section(".zlog$f", read)
#	pragma section(".zlog$m", read)
#	pragma section(".zlog$s", read)
#	pragma section(".zlog$z", read)
#endif

namespace zen {

LoggerRef TheDefaultLogger;

}  // namespace zen

namespace zen::logging {

using MemoryBuffer_t = fmt::basic_memory_buffer<char, 250>;

struct LoggingContext
{
	inline LoggingContext();
	inline ~LoggingContext();

	zen::logging::MemoryBuffer_t MessageBuffer;

	inline std::string_view Message() const { return std::string_view(MessageBuffer.data(), MessageBuffer.size()); }
};

LoggingContext::LoggingContext()
{
}

LoggingContext::~LoggingContext()
{
}

//////////////////////////////////////////////////////////////////////////

static inline bool
IsErrorLevel(LogLevel InLevel)
{
	return (InLevel == Err || InLevel == Critical);
};

void
EmitLogMessage(LoggerRef& Logger, LogLevel InLevel, const std::string_view Message)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	Logger.m_Logger->Log(InLevel, Message);
	if (IsErrorLevel(InLevel))
	{
		if (LoggerRef ErrLogger = zen::logging::ErrorLog())
		{
			ErrLogger.m_Logger->Log(InLevel, Message);
		}
	}
}

void
EmitLogMessage(LoggerRef& Logger, LogLevel InLevel, std::string_view Format, fmt::format_args Args)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	zen::logging::LoggingContext LogCtx;
	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
	zen::logging::EmitLogMessage(Logger, InLevel, LogCtx.Message());
}

void
EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, LogLevel InLevel, const std::string_view Message)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	Logger.m_Logger->Log(InLocation, InLevel, Message);
	if (IsErrorLevel(InLevel))
	{
		if (LoggerRef ErrLogger = zen::logging::ErrorLog())
		{
			ErrLogger.m_Logger->Log(InLocation, InLevel, Message);
		}
	}
}

void
EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, LogLevel InLevel, std::string_view Format, fmt::format_args Args)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	zen::logging::LoggingContext LogCtx;
	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
	zen::logging::EmitLogMessage(Logger, InLocation, InLevel, LogCtx.Message());
}

void
EmitConsoleLogMessage(LogLevel InLevel, const std::string_view Message)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	ConsoleLog().m_Logger->Log(InLevel, Message);
}

#define ZEN_COLOR_YELLOW	 "\033[0;33m"
#define ZEN_COLOR_RED		 "\033[0;31m"
#define ZEN_BRIGHT_COLOR_RED "\033[1;31m"
#define ZEN_COLOR_RESET		 "\033[0m"

void
EmitConsoleLogMessage(LogLevel InLevel, std::string_view Format, fmt::format_args Args)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);
	zen::logging::LoggingContext LogCtx;

	// We are not using a format option for console which include log level since it would interfere with normal console output

	switch (InLevel)
	{
		case Warn:
			fmt::format_to(fmt::appender(LogCtx.MessageBuffer), ZEN_COLOR_YELLOW "Warning: " ZEN_COLOR_RESET);
			break;
		case Err:
			fmt::format_to(fmt::appender(LogCtx.MessageBuffer), ZEN_BRIGHT_COLOR_RED "Error: " ZEN_COLOR_RESET);
			break;
		case Critical:
			fmt::format_to(fmt::appender(LogCtx.MessageBuffer), ZEN_COLOR_RED "Critical: " ZEN_COLOR_RESET);
			break;
		default:
			break;
	}
	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
	zen::logging::EmitConsoleLogMessage(InLevel, LogCtx.Message());
}

}  // namespace zen::logging

namespace zen::logging {

constinit std::string_view LevelNames[] = {ZEN_LEVEL_NAME_TRACE,
										   ZEN_LEVEL_NAME_DEBUG,
										   ZEN_LEVEL_NAME_INFO,
										   ZEN_LEVEL_NAME_WARNING,
										   ZEN_LEVEL_NAME_ERROR,
										   ZEN_LEVEL_NAME_CRITICAL,
										   ZEN_LEVEL_NAME_OFF};

LogLevel
ParseLogLevelString(std::string_view Name)
{
	for (int Level = 0; Level < LogLevelCount; ++Level)
	{
		if (LevelNames[Level] == Name)
			return static_cast<LogLevel>(Level);
	}

	if (Name == "warn")
	{
		return Warn;
	}

	if (Name == "err")
	{
		return Err;
	}

	return Off;
}

std::string_view
ToStringView(LogLevel Level)
{
	if (int(Level) < LogLevelCount)
	{
		return LevelNames[int(Level)];
	}

	return "None";
}

}  // namespace zen::logging

//////////////////////////////////////////////////////////////////////////

namespace zen::logging {

RwLock		LogLevelsLock;
std::string LogLevels[LogLevelCount];

void
ConfigureLogLevels(LogLevel Level, std::string_view Loggers)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	RwLock::ExclusiveLockScope _(LogLevelsLock);
	LogLevels[Level] = Loggers;
}

void
RefreshLogLevels(LogLevel* DefaultLevel)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	Registry::LogLevels Levels;

	{
		RwLock::SharedLockScope _(LogLevelsLock);

		for (int i = 0; i < LogLevelCount; ++i)
		{
			LogLevel CurrentLevel{i};

			std::string_view Spec = LogLevels[i];

			while (!Spec.empty())
			{
				std::string LoggerName;

				if (auto CommaPos = Spec.find_first_of(','); CommaPos != std::string_view::npos)
				{
					LoggerName = Spec.substr(0, CommaPos);
					Spec.remove_prefix(CommaPos + 1);
				}
				else
				{
					LoggerName = Spec;
					Spec	   = {};
				}

				Levels[LoggerName] = CurrentLevel;
			}
		}
	}

	Registry::Instance().SetLevels(Levels, DefaultLevel);
}

void
RefreshLogLevels(LogLevel DefaultLevel)
{
	RefreshLogLevels(&DefaultLevel);
}

void
RefreshLogLevels()
{
	RefreshLogLevels(nullptr);
}

void
SetLogLevel(LogLevel NewLogLevel)
{
	Registry::Instance().SetGlobalLevel(NewLogLevel);
}

LogLevel
GetLogLevel()
{
	return Registry::Instance().GetGlobalLevel();
}

LoggerRef
Default()
{
	ZEN_ASSERT(TheDefaultLogger);
	return TheDefaultLogger;
}

void
SetDefault(std::string_view NewDefaultLoggerId)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	auto NewDefaultLogger = Registry::Instance().Get(std::string(NewDefaultLoggerId));
	ZEN_ASSERT(NewDefaultLogger);

	Registry::Instance().SetDefaultLogger(NewDefaultLogger);
	TheDefaultLogger = LoggerRef(*NewDefaultLogger);
}

LoggerRef TheErrorLogger;

LoggerRef
ErrorLog()
{
	return TheErrorLogger;
}

void
SetErrorLog(std::string_view NewErrorLoggerId)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	if (NewErrorLoggerId.empty())
	{
		TheErrorLogger = {};
	}
	else
	{
		auto NewErrorLogger = Registry::Instance().Get(std::string(NewErrorLoggerId));

		ZEN_ASSERT(NewErrorLogger);

		TheErrorLogger = LoggerRef(*NewErrorLogger.get());
	}
}

RwLock g_LoggerMutex;

LoggerRef
Get(std::string_view Name)
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	std::shared_ptr<Logger> FoundLogger = Registry::Instance().Get(std::string(Name));

	if (!FoundLogger)
	{
		g_LoggerMutex.WithExclusiveLock([&] {
			FoundLogger = Registry::Instance().Get(std::string(Name));

			if (!FoundLogger)
			{
				FoundLogger = Default().m_Logger->Clone(std::string(Name));
				Registry::Instance().Register(FoundLogger);
			}
		});
	}

	return *FoundLogger;
}

std::once_flag			ConsoleInitFlag;
std::shared_ptr<Logger> ConLogger;

void
SuppressConsoleLog()
{
	if (ConLogger)
	{
		Registry::Instance().Drop("console");
		ConLogger = {};
	}

	auto NullSinkPtr = std::make_shared<NullSink>();
	ConLogger		 = std::make_shared<Logger>("console", std::vector<SinkPtr>{NullSinkPtr});
	Registry::Instance().Register(ConLogger);
}

LoggerRef
ConsoleLog()
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	std::call_once(ConsoleInitFlag, [&] {
		if (!ConLogger)
		{
			auto ConsoleSink = std::make_shared<AnsiColorStdoutSink>();
			ConsoleSink->SetFormatter(std::make_unique<MessageOnlyFormatter>());
			ConLogger = std::make_shared<Logger>("console", std::vector<SinkPtr>{ConsoleSink});
			Registry::Instance().Register(ConLogger);
		}
	});

	return *ConLogger;
}

void
ResetConsoleLog()
{
	LoggerRef ConLog = ConsoleLog();

	for (auto& sink : ConLog.m_Logger->Sinks())
	{
		sink->SetFormatter(std::make_unique<MessageOnlyFormatter>());
	}
}

void
InitializeLogging()
{
	ZEN_MEMSCOPE(ELLMTag::Logging);

	TheDefaultLogger = *Registry::Instance().DefaultLoggerRaw();
}

void
ShutdownLogging()
{
	Registry::Instance().Shutdown();
	TheDefaultLogger = {};
}

bool
EnableVTMode()
{
#if ZEN_PLATFORM_WINDOWS
	// Set output mode to handle virtual terminal sequences
	HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
	if (hOut == INVALID_HANDLE_VALUE)
	{
		return false;
	}

	DWORD dwMode = 0;
	if (!GetConsoleMode(hOut, &dwMode))
	{
		return false;
	}

	dwMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
	if (!SetConsoleMode(hOut, dwMode))
	{
		return false;
	}
#endif

	return true;
}

void
FlushLogging()
{
	Registry::Instance().FlushAll();
}

}  // namespace zen::logging

namespace zen {

bool
LoggerRef::ShouldLog(logging::LogLevel Level) const
{
	return m_Logger->ShouldLog(Level);
}

void
LoggerRef::SetLogLevel(logging::LogLevel NewLogLevel)
{
	m_Logger->SetLevel(NewLogLevel);
}

logging::LogLevel
LoggerRef::GetLogLevel()
{
	return m_Logger->GetLevel();
}

void
LoggerRef::Flush()
{
	m_Logger->Flush();
}

thread_local ScopedActivityBase* t_ScopeStack = nullptr;

ScopedActivityBase*
GetThreadActivity()
{
	return t_ScopeStack;
}

ScopedActivityBase::ScopedActivityBase() : m_NextScope{t_ScopeStack}
{
	t_ScopeStack = this;
}

ScopedActivityBase::~ScopedActivityBase()
{
	if (t_ScopeStack != this)
	{
		ZEN_ERROR("invalid t_ScopeStack in ~ScopedActivityBase(). Expected {:#x}, found {:#x}", (uintptr_t)this, (uintptr_t)t_ScopeStack);
		return;
	}
	t_ScopeStack = m_NextScope;
}

std::string_view
EmitActivitiesForLogging(StringBuilderBase& OutString)
{
	OutString.Reset();

	for (auto Bread = GetThreadActivity(); Bread; Bread = Bread->GetNext())
	{
		OutString.Append("\n|>|>|> ");
		Bread->Emit(OutString);
	}

	return OutString.ToView();
}

//////////////////////////////////////////////////////////////////////////

ScopedActivityString::~ScopedActivityString()
{
}

void
ScopedActivityString::Emit(StringBuilderBase& Target)
{
	Target.Append(m_Text);
}

#if ZEN_WITH_TESTS

void
logging_forcelink()
{
}

using namespace std::literals;

TEST_CASE("simple.bread")
{
	ExtendableStringBuilder<256> Crumbs;

	auto EmitBreadcrumbs = [&] {
		Crumbs.Reset();

		for (auto Bread = GetThreadActivity(); Bread; Bread = Bread->GetNext())
		{
			Bread->Emit(Crumbs);
		}

		return Crumbs.ToView();
	};

	SUBCASE("single")
	{
		ScopedActivityString _{"hello"};
		EmitBreadcrumbs();

		CHECK_EQ(Crumbs.ToView(), "hello"sv);
	}

	SUBCASE("multi")
	{
		ScopedActivityString $1{"hello"};
		ScopedActivityString $2{"world"};
		ScopedActivityString $3{"amaze"};

		CHECK_EQ(EmitBreadcrumbs(), "amazeworldhello"sv);
	}

	SUBCASE("multi_defer")
	{
		int					 n = 42;
		ScopedActivityString $1{"hello"};
		ScopedActivityString $2{"world"};
		ScopedActivityString $3{"amaze"};
		ScopedLazyActivity	 $4{[&](StringBuilderBase& Out) { Out << "plant"; }};
		{
			ScopedLazyActivity $5{[&](StringBuilderBase& Out) { Out << n; }};
			ZEN_LOG_SCOPE("{}:{}", "abc", "def");
			CHECK_EQ(EmitBreadcrumbs(), "abc:def42plantamazeworldhello"sv);
		}
		CHECK_EQ(EmitBreadcrumbs(), "plantamazeworldhello"sv);
	}
}

#endif

}  // namespace zen