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

#include <zencore/logging/registry.h>

#include <zencore/logging/ansicolorsink.h>
#include <zencore/logging/messageonlyformatter.h>

#include <thread>

namespace zen::logging {

struct Registry::FlushThread
{
	std::thread Thread;
};

Registry&
Registry::Instance()
{
	static Registry s_Instance;
	return s_Instance;
}

Registry::Registry()
{
	// Create default logger with a stdout color sink
	auto DefaultSink = std::make_shared<AnsiColorStdoutSink>();
	m_DefaultLogger	 = std::make_shared<Logger>("", std::vector<SinkPtr>{DefaultSink});
	m_Loggers[""]	 = m_DefaultLogger;
}

Registry::~Registry()
{
	StopPeriodicFlush();
}

void
Registry::Register(std::shared_ptr<Logger> InLogger)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	m_Loggers[InLogger->Name()] = std::move(InLogger);
}

void
Registry::Drop(const std::string& Name)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	m_Loggers.erase(Name);
}

std::shared_ptr<Logger>
Registry::Get(const std::string& Name)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	auto						It = m_Loggers.find(Name);
	if (It != m_Loggers.end())
	{
		return It->second;
	}
	return {};
}

void
Registry::SetDefaultLogger(std::shared_ptr<Logger> InLogger)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	if (InLogger)
	{
		m_Loggers[InLogger->Name()] = InLogger;
	}
	m_DefaultLogger = std::move(InLogger);
}

Logger*
Registry::DefaultLoggerRaw()
{
	return m_DefaultLogger.get();
}

std::shared_ptr<Logger>
Registry::DefaultLogger()
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	return m_DefaultLogger;
}

void
Registry::SetGlobalLevel(LogLevel Level)
{
	m_GlobalLevel.store(Level, std::memory_order_relaxed);
	std::lock_guard<std::mutex> Lock(m_Mutex);
	for (auto& [Name, CurLogger] : m_Loggers)
	{
		CurLogger->SetLevel(Level);
	}
}

LogLevel
Registry::GetGlobalLevel() const
{
	return m_GlobalLevel.load(std::memory_order_relaxed);
}

void
Registry::SetLevels(const LogLevels& Levels, LogLevel* DefaultLevel)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);

	if (DefaultLevel)
	{
		m_GlobalLevel.store(*DefaultLevel, std::memory_order_relaxed);
		for (auto& [Name, CurLogger] : m_Loggers)
		{
			CurLogger->SetLevel(*DefaultLevel);
		}
	}

	for (auto& [LoggerName, Level] : Levels)
	{
		auto It = m_Loggers.find(LoggerName);
		if (It != m_Loggers.end())
		{
			It->second->SetLevel(Level);
		}
	}
}

void
Registry::FlushAll()
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	for (auto& [Name, CurLogger] : m_Loggers)
	{
		try
		{
			CurLogger->Flush();
		}
		catch (const std::exception&)
		{
		}
	}
}

void
Registry::FlushOn(LogLevel Level)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	m_FlushLevel = Level;
	for (auto& [Name, CurLogger] : m_Loggers)
	{
		CurLogger->SetFlushLevel(Level);
	}
}

void
Registry::FlushEvery(std::chrono::seconds Interval)
{
	StartPeriodicFlush(Interval);
}

void
Registry::SetFormatter(std::unique_ptr<Formatter> InFormatter)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	for (auto& [Name, CurLogger] : m_Loggers)
	{
		CurLogger->SetFormatter(InFormatter->Clone());
	}
}

void
Registry::ApplyAll(std::function<void(std::shared_ptr<Logger>)> Func)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	for (auto& [Name, CurLogger] : m_Loggers)
	{
		Func(CurLogger);
	}
}

void
Registry::SetErrorHandler(std::function<void(const std::string&)> Handler)
{
	std::lock_guard<std::mutex> Lock(m_Mutex);
	m_ErrorHandler = std::move(Handler);
}

void
Registry::Shutdown()
{
	StopPeriodicFlush();
	FlushAll();

	std::lock_guard<std::mutex> Lock(m_Mutex);
	m_Loggers.clear();
	m_DefaultLogger.reset();
}

void
Registry::StartPeriodicFlush(std::chrono::seconds Interval)
{
	StopPeriodicFlush();

	m_PeriodicFlushRunning.store(true, std::memory_order_relaxed);
	m_FlushThread = std::make_unique<FlushThread>();

	m_FlushThread->Thread = std::thread([this, Interval] {
		while (m_PeriodicFlushRunning.load(std::memory_order_relaxed))
		{
			{
				std::unique_lock<std::mutex> Lock(m_PeriodicFlushMutex);
				m_PeriodicFlushCv.wait_for(Lock, Interval, [this] { return !m_PeriodicFlushRunning.load(std::memory_order_relaxed); });
			}

			if (m_PeriodicFlushRunning.load(std::memory_order_relaxed))
			{
				FlushAll();
			}
		}
	});
}

void
Registry::StopPeriodicFlush()
{
	if (m_FlushThread)
	{
		m_PeriodicFlushRunning.store(false, std::memory_order_relaxed);
		{
			std::lock_guard<std::mutex> Lock(m_PeriodicFlushMutex);
			m_PeriodicFlushCv.notify_one();
		}
		if (m_FlushThread->Thread.joinable())
		{
			m_FlushThread->Thread.join();
		}
		m_FlushThread.reset();
	}
}

}  // namespace zen::logging