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

#include "deferreddeleter.h"

#if ZEN_WITH_COMPUTE_SERVICES

#	include <zencore/filesystem.h>
#	include <zencore/fmtutils.h>
#	include <zencore/logging.h>
#	include <zencore/thread.h>

#	include <algorithm>
#	include <chrono>

namespace zen::compute {

using namespace std::chrono_literals;

using Clock = std::chrono::steady_clock;

// Default deferral: how long to wait before attempting deletion.
// This gives memory-mapped file handles time to close naturally.
static constexpr auto DeferralPeriod = 60s;

// Shortened deferral after MarkReady(): the client has collected results
// so handles should be released soon, but we still wait briefly.
static constexpr auto ReadyGracePeriod = 5s;

// Interval between retry attempts for directories that failed deletion.
static constexpr auto RetryInterval = 5s;

static constexpr int MaxRetries = 10;

DeferredDirectoryDeleter::DeferredDirectoryDeleter() : m_Thread(&DeferredDirectoryDeleter::ThreadFunction, this)
{
}

DeferredDirectoryDeleter::~DeferredDirectoryDeleter()
{
	Shutdown();
}

void
DeferredDirectoryDeleter::Enqueue(int ActionLsn, std::filesystem::path Path)
{
	{
		std::lock_guard Lock(m_Mutex);
		m_Queue.push_back({ActionLsn, std::move(Path)});
	}
	m_Cv.notify_one();
}

void
DeferredDirectoryDeleter::MarkReady(int ActionLsn)
{
	{
		std::lock_guard Lock(m_Mutex);
		m_ReadyLsns.push_back(ActionLsn);
	}
	m_Cv.notify_one();
}

void
DeferredDirectoryDeleter::Shutdown()
{
	{
		std::lock_guard Lock(m_Mutex);
		m_Done = true;
	}
	m_Cv.notify_one();

	if (m_Thread.joinable())
	{
		m_Thread.join();
	}
}

void
DeferredDirectoryDeleter::ThreadFunction()
{
	SetCurrentThreadName("ZenDirCleanup");

	struct PendingEntry
	{
		int					  ActionLsn;
		std::filesystem::path Path;
		Clock::time_point	  ReadyTime;
		int					  Attempts = 0;
	};

	std::vector<PendingEntry> PendingList;

	auto TryDelete = [](PendingEntry& Entry) -> bool {
		std::error_code Ec;
		std::filesystem::remove_all(Entry.Path, Ec);
		return !Ec;
	};

	for (;;)
	{
		bool Shutting = false;

		// Drain the incoming queue and process MarkReady signals

		{
			std::unique_lock Lock(m_Mutex);

			if (m_Queue.empty() && m_ReadyLsns.empty() && !m_Done)
			{
				if (PendingList.empty())
				{
					m_Cv.wait(Lock, [this] { return !m_Queue.empty() || !m_ReadyLsns.empty() || m_Done; });
				}
				else
				{
					auto NextReady = PendingList.front().ReadyTime;
					for (const auto& Entry : PendingList)
					{
						if (Entry.ReadyTime < NextReady)
						{
							NextReady = Entry.ReadyTime;
						}
					}

					m_Cv.wait_until(Lock, NextReady, [this] { return !m_Queue.empty() || !m_ReadyLsns.empty() || m_Done; });
				}
			}

			// Move new items into PendingList with the full deferral deadline
			auto Now = Clock::now();
			for (auto& Entry : m_Queue)
			{
				PendingList.push_back({Entry.ActionLsn, std::move(Entry.Path), Now + DeferralPeriod, 0});
			}
			m_Queue.clear();

			// Apply MarkReady: shorten ReadyTime for matching entries
			for (int Lsn : m_ReadyLsns)
			{
				for (auto& Entry : PendingList)
				{
					if (Entry.ActionLsn == Lsn)
					{
						auto NewReady = Now + ReadyGracePeriod;
						if (NewReady < Entry.ReadyTime)
						{
							Entry.ReadyTime = NewReady;
						}
					}
				}
			}
			m_ReadyLsns.clear();

			Shutting = m_Done;
		}

		// Process items whose deferral period has elapsed (or all items on shutdown)

		auto Now = Clock::now();

		for (size_t i = 0; i < PendingList.size();)
		{
			auto& Entry = PendingList[i];

			if (!Shutting && Now < Entry.ReadyTime)
			{
				++i;
				continue;
			}

			if (TryDelete(Entry))
			{
				if (Entry.Attempts > 0)
				{
					ZEN_INFO("Retry succeeded for directory '{}'", Entry.Path);
				}

				PendingList[i] = std::move(PendingList.back());
				PendingList.pop_back();
			}
			else
			{
				++Entry.Attempts;

				if (Entry.Attempts >= MaxRetries)
				{
					ZEN_WARN("Giving up on deleting '{}' after {} attempts", Entry.Path, Entry.Attempts);
					PendingList[i] = std::move(PendingList.back());
					PendingList.pop_back();
				}
				else
				{
					ZEN_WARN("Unable to delete directory '{}' (attempt {}), will retry", Entry.Path, Entry.Attempts);
					Entry.ReadyTime = Now + RetryInterval;
					++i;
				}
			}
		}

		// Exit once shutdown is requested and nothing remains

		if (Shutting && PendingList.empty())
		{
			return;
		}
	}
}

}  // namespace zen::compute

#endif

#if ZEN_WITH_TESTS

#	include <zencore/testing.h>

namespace zen::compute {

void
deferreddeleter_forcelink()
{
}

}  // namespace zen::compute

#endif

#if ZEN_WITH_TESTS && ZEN_WITH_COMPUTE_SERVICES

#	include <zencore/testutils.h>

namespace zen::compute {

TEST_CASE("DeferredDirectoryDeleter.DeletesSingleDirectory")
{
	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 DirToDelete = TempDir.Path() / "subdir";
	CreateDirectories(DirToDelete / "nested");

	CHECK(std::filesystem::exists(DirToDelete));

	{
		DeferredDirectoryDeleter Deleter;
		Deleter.Enqueue(1, DirToDelete);
	}

	CHECK(!std::filesystem::exists(DirToDelete));
}

TEST_CASE("DeferredDirectoryDeleter.DeletesMultipleDirectories")
{
	ScopedTemporaryDirectory TempDir;

	constexpr int					   NumDirs = 10;
	std::vector<std::filesystem::path> Dirs;

	for (int i = 0; i < NumDirs; ++i)
	{
		auto Dir = TempDir.Path() / std::to_string(i);
		CreateDirectories(Dir / "child");
		Dirs.push_back(std::move(Dir));
	}

	{
		DeferredDirectoryDeleter Deleter;
		for (int i = 0; i < NumDirs; ++i)
		{
			CHECK(std::filesystem::exists(Dirs[i]));
			Deleter.Enqueue(100 + i, Dirs[i]);
		}
	}

	for (const auto& Dir : Dirs)
	{
		CHECK(!std::filesystem::exists(Dir));
	}
}

TEST_CASE("DeferredDirectoryDeleter.ShutdownIsIdempotent")
{
	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 Dir = TempDir.Path() / "idempotent";
	CreateDirectories(Dir);

	DeferredDirectoryDeleter Deleter;
	Deleter.Enqueue(42, Dir);
	Deleter.Shutdown();
	Deleter.Shutdown();

	CHECK(!std::filesystem::exists(Dir));
}

TEST_CASE("DeferredDirectoryDeleter.HandlesNonExistentPath")
{
	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 NoSuchDir = TempDir.Path() / "does_not_exist";

	{
		DeferredDirectoryDeleter Deleter;
		Deleter.Enqueue(99, NoSuchDir);
	}
}

TEST_CASE("DeferredDirectoryDeleter.ExplicitShutdownBeforeDestruction")
{
	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 Dir = TempDir.Path() / "explicit";
	CreateDirectories(Dir / "inner");

	DeferredDirectoryDeleter Deleter;
	Deleter.Enqueue(7, Dir);
	Deleter.Shutdown();

	CHECK(!std::filesystem::exists(Dir));
}

TEST_CASE("DeferredDirectoryDeleter.MarkReadyShortensDeferral")
{
	ScopedTemporaryDirectory TempDir;
	std::filesystem::path	 Dir = TempDir.Path() / "markready";
	CreateDirectories(Dir / "child");

	DeferredDirectoryDeleter Deleter;
	Deleter.Enqueue(50, Dir);

	// Without MarkReady the full deferral (60s) would apply.
	// MarkReady shortens it to 5s, and shutdown bypasses even that.
	Deleter.MarkReady(50);
	Deleter.Shutdown();

	CHECK(!std::filesystem::exists(Dir));
}

}  // namespace zen::compute

#endif	// ZEN_WITH_TESTS && ZEN_WITH_COMPUTE_SERVICES