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

#include "wipe_cmd.h"

#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iohash.h>
#include <zencore/logging.h>
#include <zencore/parallelwork.h>
#include <zencore/string.h>
#include <zencore/timer.h>
#include <zencore/trace.h>
#include <zenutil/workerpools.h>

#include "consoleprogress.h"

#include <signal.h>

#include <iostream>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
#include <tsl/robin_set.h>
ZEN_THIRD_PARTY_INCLUDES_END

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#else
#	include <fcntl.h>
#	include <sys/file.h>
#	include <sys/stat.h>
#	include <unistd.h>
#endif

namespace zen {

namespace wipe_impl {
	static std::atomic<bool>   AbortFlag		  = false;
	static std::atomic<bool>   PauseFlag		  = false;
	static bool				   IsVerbose		  = false;
	static bool				   Quiet			  = false;
	static ConsoleProgressMode ProgressMode		  = ConsoleProgressMode::Pretty;
	const bool				   SingleThreaded	  = false;
	bool					   BoostWorkerThreads = true;

	WorkerThreadPool& GetIOWorkerPool()
	{
		return SingleThreaded		? GetSyncWorkerPool()
			   : BoostWorkerThreads ? GetLargeWorkerPool(EWorkloadType::Burst)
									: GetMediumWorkerPool(EWorkloadType::Burst);
	}

#undef ZEN_CONSOLE_VERBOSE
#define ZEN_CONSOLE_VERBOSE(fmtstr, ...)                            \
	if (IsVerbose)                                                  \
	{                                                               \
		ZEN_CONSOLE_LOG(zen::logging::Info, fmtstr, ##__VA_ARGS__); \
	}

	static void SignalCallbackHandler(int SigNum)
	{
		if (SigNum == SIGINT)
		{
			PauseFlag = false;
			AbortFlag = true;
		}
#if ZEN_PLATFORM_WINDOWS
		if (SigNum == SIGBREAK)
		{
			PauseFlag = false;
			AbortFlag = true;
		}
#endif	// ZEN_PLATFORM_WINDOWS
	}

	bool IsReadOnly(uint32_t Attributes)
	{
#if ZEN_PLATFORM_WINDOWS
		return IsFileAttributeReadOnly(Attributes);
#else
		return IsFileModeReadOnly(Attributes);
#endif
	}

	bool IsFileWithRetry(const std::filesystem::path& Path)
	{
		std::error_code Ec;
		bool			Result = IsFile(Path, Ec);
		for (size_t Retries = 0; Ec && Retries < 3; Retries++)
		{
			Sleep(100 + int(Retries * 50));
			Ec.clear();
			Result = IsFile(Path, Ec);
		}
		if (Ec)
		{
			zen::ThrowSystemError(Ec.value(), Ec.message());
		}
		return Result;
	}

	bool SetFileReadOnlyWithRetry(const std::filesystem::path& Path, bool ReadOnly)
	{
		std::error_code Ec;
		bool			Result = SetFileReadOnly(Path, ReadOnly, Ec);
		for (size_t Retries = 0; Ec && Retries < 3; Retries++)
		{
			Sleep(100 + int(Retries * 50));
			if (!IsFileWithRetry(Path))
			{
				return false;
			}
			Ec.clear();
			Result = SetFileReadOnly(Path, ReadOnly, Ec);
		}
		if (Ec)
		{
			zen::ThrowSystemError(Ec.value(), Ec.message());
		}
		return Result;
	}

	void RemoveFileWithRetry(const std::filesystem::path& Path)
	{
		std::error_code Ec;
		RemoveFile(Path, Ec);
		for (size_t Retries = 0; Ec && Retries < 3; Retries++)
		{
			Sleep(100 + int(Retries * 50));
			if (!IsFileWithRetry(Path))
			{
				return;
			}
			Ec.clear();
			RemoveFile(Path, Ec);
		}
		if (Ec)
		{
			zen::ThrowSystemError(Ec.value(), Ec.message());
		}
	}

	void RemoveDirWithRetry(const std::filesystem::path& Path)
	{
		std::error_code Ec;
		RemoveDir(Path, Ec);
		for (size_t Retries = 0; Ec && Retries < 3; Retries++)
		{
			Sleep(100 + int(Retries * 50));
			if (!IsDir(Path))
			{
				return;
			}
			Ec.clear();
			RemoveDir(Path, Ec);
		}
		if (Ec)
		{
			zen::ThrowSystemError(Ec.value(), Ec.message());
		}
	}

	bool CleanDirectory(const std::filesystem::path&	  Path,
						std::span<const std::string_view> ExcludeDirectories,
						bool							  RemoveReadonly,
						bool							  Dryrun)
	{
		ZEN_TRACE_CPU("CleanDirectory");
		Stopwatch Timer;

		std::unique_ptr<ProgressBase>			   ProgressOwner(CreateConsoleProgress(ProgressMode));
		std::unique_ptr<ProgressBase::ProgressBar> Progress = ProgressOwner->CreateProgressBar("Clean Folder");

		std::atomic<bool>	  CleanWipe			  = true;
		std::atomic<uint64_t> DiscoveredItemCount = 0;
		std::atomic<uint64_t> DeletedItemCount	  = 0;
		std::atomic<uint64_t> DeletedByteCount	  = 0;
		std::atomic<uint64_t> FailedDeleteCount	  = 0;

		std::vector<std::filesystem::path>			   SubdirectoriesToDelete;
		tsl::robin_map<IoHash, size_t, IoHash::Hasher> SubdirectoriesToDeleteLookup;
		tsl::robin_set<IoHash, IoHash::Hasher>		   SubdirectoriesToKeep;
		RwLock										   SubdirectoriesLock;

		auto AddFoundDirectory = [&](std::filesystem::path Directory, bool Keep) -> bool {
			bool Added = false;
			if (Keep)
			{
				bool IsLeaf = true;
				while (Directory != Path)
				{
					const std::string		   DirectoryString	 = Directory.generic_string();
					IoHash					   DirectoryNameHash = IoHash::HashBuffer(DirectoryString.data(), DirectoryString.length());
					RwLock::ExclusiveLockScope _(SubdirectoriesLock);
					if (auto It = SubdirectoriesToKeep.find(DirectoryNameHash); It == SubdirectoriesToKeep.end())
					{
						SubdirectoriesToKeep.insert(DirectoryNameHash);
						if (IsLeaf)
						{
							Added = true;
						}
					}
					else
					{
						break;
					}
					Directory = Directory.parent_path();
					IsLeaf	  = false;
				}
			}
			else
			{
				bool IsLeaf = true;
				while (Directory != Path)
				{
					const std::string		   DirectoryString	 = Directory.generic_string();
					IoHash					   DirectoryNameHash = IoHash::HashBuffer(DirectoryString.data(), DirectoryString.length());
					RwLock::ExclusiveLockScope _(SubdirectoriesLock);
					if (SubdirectoriesToKeep.contains(DirectoryNameHash))
					{
						break;
					}
					if (auto It = SubdirectoriesToDeleteLookup.find(DirectoryNameHash); It == SubdirectoriesToDeleteLookup.end())
					{
						SubdirectoriesToDeleteLookup.insert({DirectoryNameHash, SubdirectoriesToDelete.size()});
						SubdirectoriesToDelete.push_back(Directory);
						if (IsLeaf)
						{
							Added = true;
						}
					}
					else
					{
						break;
					}
					Directory = Directory.parent_path();
					IsLeaf	  = false;
				}
			}
			return Added;
		};

		ParallelWork Work(AbortFlag, PauseFlag, WorkerThreadPool::EMode::EnableBacklog);

		struct AsyncVisitor : public GetDirectoryContentVisitor
		{
			AsyncVisitor(const std::filesystem::path&							 InPath,
						 std::atomic<bool>&										 InCleanWipe,
						 std::atomic<uint64_t>&									 InDiscoveredItemCount,
						 std::atomic<uint64_t>&									 InDeletedItemCount,
						 std::atomic<uint64_t>&									 InDeletedByteCount,
						 std::atomic<uint64_t>&									 InFailedDeleteCount,
						 std::span<const std::string_view>						 InExcludeDirectories,
						 bool													 InRemoveReadonly,
						 bool													 InDryrun,
						 const std::function<bool(std::filesystem::path, bool)>& InAddFoundDirectoryFunc)
			: Path(InPath)
			, CleanWipe(InCleanWipe)
			, DiscoveredItemCount(InDiscoveredItemCount)
			, DeletedItemCount(InDeletedItemCount)
			, DeletedByteCount(InDeletedByteCount)
			, FailedDeleteCount(InFailedDeleteCount)
			, ExcludeDirectories(InExcludeDirectories)
			, RemoveReadonly(InRemoveReadonly)
			, Dryrun(InDryrun)
			, AddFoundDirectoryFunc(InAddFoundDirectoryFunc)
			{
			}
			virtual void AsyncVisitDirectory(const std::filesystem::path& RelativeRoot, DirectoryContent&& Content) override
			{
				ZEN_TRACE_CPU("CleanDirectory_AsyncVisitDirectory");
				if (!AbortFlag)
				{
					if (!RelativeRoot.empty())
					{
						DiscoveredItemCount++;
					}
					if (Content.FileNames.empty())
					{
						const std::filesystem::path ParentPath	  = Path / RelativeRoot;
						bool						KeepDirectory = RelativeRoot.empty();

						bool Added = AddFoundDirectoryFunc(ParentPath, KeepDirectory);
						if (Added)
						{
							ZEN_CONSOLE_VERBOSE("{} directory {}", KeepDirectory ? "Keeping" : "Removing", ParentPath);
						}
					}
					else
					{
						DiscoveredItemCount += Content.FileNames.size();

						const std::string RelativeRootString = RelativeRoot.generic_string();
						bool			  RemoveContent		 = true;
						for (const std::string_view ExcludeDirectory : ExcludeDirectories)
						{
							if (RelativeRootString.starts_with(ExcludeDirectory))
							{
								if (RelativeRootString.length() > ExcludeDirectory.length())
								{
									const char MaybePathDelimiter = RelativeRootString[ExcludeDirectory.length()];
									if (MaybePathDelimiter == '/' || MaybePathDelimiter == '\\' ||
										MaybePathDelimiter == std::filesystem::path::preferred_separator)
									{
										RemoveContent = false;
										break;
									}
								}
								else
								{
									RemoveContent = false;
									break;
								}
							}
						}

						const std::filesystem::path ParentPath	  = Path / RelativeRoot;
						bool						KeepDirectory = RelativeRoot.empty();

						if (RemoveContent)
						{
							ZEN_TRACE_CPU("DeleteFiles");
							uint64_t RemovedCount = 0;
							for (size_t FileIndex = 0; FileIndex < Content.FileNames.size(); FileIndex++)
							{
								const std::filesystem::path& FileName = Content.FileNames[FileIndex];
								const std::filesystem::path	 FilePath = (ParentPath / FileName).make_preferred();
								try
								{
									const uint32_t Attributes = Content.FileAttributes[FileIndex];
									const bool	   IsReadonly = IsReadOnly(Attributes);
									bool		   RemoveFile = false;
									if (IsReadonly)
									{
										if (RemoveReadonly)
										{
											if (!Dryrun)
											{
												SetFileReadOnlyWithRetry(FilePath, false);
											}
											RemoveFile = true;
										}
									}
									else
									{
										RemoveFile = true;
									}

									if (RemoveFile)
									{
										if (!Dryrun)
										{
											RemoveFileWithRetry(FilePath);
										}
										DeletedItemCount++;
										DeletedByteCount += Content.FileSizes[FileIndex];
										RemovedCount++;
										ZEN_CONSOLE_VERBOSE("Removed file {}", FilePath);
									}
									else
									{
										ZEN_CONSOLE_VERBOSE("Skipped readonly file {}", FilePath);
										KeepDirectory = true;
									}
								}
								catch (const std::exception& Ex)
								{
									ZEN_WARN("Failed removing file {}. Reason: {}", FilePath, Ex.what());
									FailedDeleteCount++;
									CleanWipe	  = false;
									KeepDirectory = true;
								}
							}
							ZEN_CONSOLE_VERBOSE("Removed {} files in {}", RemovedCount, ParentPath);
						}
						else
						{
							ZEN_CONSOLE_VERBOSE("Skipped removal of {} files in {}", Content.FileNames.size(), ParentPath);
						}
						bool Added = AddFoundDirectoryFunc(ParentPath, KeepDirectory);
						if (Added)
						{
							ZEN_CONSOLE_VERBOSE("{} directory {}", KeepDirectory ? "Keeping" : "Removing", ParentPath);
						}
					}
				}
			}
			const std::filesystem::path&					 Path;
			std::atomic<bool>&								 CleanWipe;
			std::atomic<uint64_t>&							 DiscoveredItemCount;
			std::atomic<uint64_t>&							 DeletedItemCount;
			std::atomic<uint64_t>&							 DeletedByteCount;
			std::atomic<uint64_t>&							 FailedDeleteCount;
			std::span<const std::string_view>				 ExcludeDirectories;
			const bool										 RemoveReadonly;
			const bool										 Dryrun;
			std::function<bool(std::filesystem::path, bool)> AddFoundDirectoryFunc;
		} Visitor(Path,
				  CleanWipe,
				  DiscoveredItemCount,
				  DeletedItemCount,
				  DeletedByteCount,
				  FailedDeleteCount,
				  ExcludeDirectories,
				  RemoveReadonly,
				  Dryrun,
				  AddFoundDirectory);

		uint64_t LastUpdateTimeMs = Timer.GetElapsedTimeMs();

		GetDirectoryContent(Path,
							DirectoryContentFlags::IncludeFiles | DirectoryContentFlags::Recursive |
								DirectoryContentFlags::IncludeFileSizes | DirectoryContentFlags::IncludeAttributes,
							Visitor,
							GetIOWorkerPool(),
							Work.PendingWork());

		Work.Wait(ProgressOwner->GetProgressUpdateDelayMS(), [&](bool IsAborted, bool IsPaused, ptrdiff_t PendingWork) {
			ZEN_UNUSED(PendingWork);
			if (Quiet)
			{
				return;
			}
			LastUpdateTimeMs = Timer.GetElapsedTimeMs();

			uint64_t Deleted	  = DeletedItemCount.load();
			uint64_t DeletedBytes = DeletedByteCount.load();
			uint64_t Discovered	  = DiscoveredItemCount.load();
			Progress->UpdateState({.Task		   = "Removing files  ",
								   .Details		   = fmt::format("Found {}, Deleted {} ({})", Discovered, Deleted, NiceBytes(DeletedBytes)),
								   .TotalCount	   = Discovered,
								   .RemainingCount = Discovered - Deleted,
								   .Status		   = ProgressBase::ProgressBar::State::CalculateStatus(IsAborted, IsPaused)},
								  false);
		});

		std::vector<std::filesystem::path> DirectoriesToDelete;
		DirectoriesToDelete.reserve(SubdirectoriesToDelete.size());
		for (auto It : SubdirectoriesToDeleteLookup)
		{
			const IoHash& DirHash = It.first;
			if (auto KeepIt = SubdirectoriesToKeep.find(DirHash); KeepIt == SubdirectoriesToKeep.end())
			{
				DirectoriesToDelete.emplace_back(std::move(SubdirectoriesToDelete[It.second]));
			}
		}

		std::sort(DirectoriesToDelete.begin(),
				  DirectoriesToDelete.end(),
				  [](const std::filesystem::path& Lhs, const std::filesystem::path& Rhs) {
					  return Lhs.string().length() > Rhs.string().length();
				  });

		for (size_t SubDirectoryIndex = 0; SubDirectoryIndex < DirectoriesToDelete.size(); SubDirectoryIndex++)
		{
			ZEN_TRACE_CPU("DeleteDirs");
			const std::filesystem::path& DirectoryToDelete = DirectoriesToDelete[SubDirectoryIndex];
			try
			{
				if (!Dryrun)
				{
					RemoveDirWithRetry(DirectoryToDelete);
				}
				ZEN_CONSOLE_VERBOSE("Removed directory {}", DirectoryToDelete);
				DeletedItemCount++;
			}
			catch (const std::exception& Ex)
			{
				if (!Quiet)
				{
					ZEN_CONSOLE_WARN("Failed removing directory {}. Reason: {}", DirectoryToDelete, Ex.what());
				}
				CleanWipe = false;
				FailedDeleteCount++;
			}

			uint64_t NowMs = Timer.GetElapsedTimeMs();
			if ((NowMs - LastUpdateTimeMs) >= ProgressOwner->GetProgressUpdateDelayMS())
			{
				LastUpdateTimeMs = NowMs;

				uint64_t Deleted	  = DeletedItemCount.load();
				uint64_t DeletedBytes = DeletedByteCount.load();
				uint64_t Discovered	  = DiscoveredItemCount.load();
				Progress->UpdateState({.Task	   = "Removing folders",
									   .Details	   = fmt::format("Found {}, Deleted {} ({})", Discovered, Deleted, NiceBytes(DeletedBytes)),
									   .TotalCount = DirectoriesToDelete.size(),
									   .RemainingCount = DirectoriesToDelete.size() - SubDirectoryIndex},
									  false);
			}
		}

		Progress->Finish();

		uint64_t ElapsedTimeMs = Timer.GetElapsedTimeMs();
		if (!Quiet)
		{
			ZEN_CONSOLE("Wiped folder '{}' {} ({}) ({} failed) in {}",
						Path,
						DeletedItemCount.load(),
						NiceBytes(DeletedByteCount.load()),
						FailedDeleteCount.load(),
						NiceTimeSpanMs(ElapsedTimeMs));
		}
		if (FailedDeleteCount.load() > 0)
		{
			throw std::runtime_error(fmt::format("Failed to delete {} files/directories in '{}'", FailedDeleteCount.load(), Path));
		}
		return CleanWipe;
	}
}  // namespace wipe_impl

WipeCommand::WipeCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "d", "directory", "Directory to wipe", cxxopts::value(m_Directory), "<directory>");
	m_Options.add_option("", "r", "keep-readonly", "Leave read-only files", cxxopts::value(m_KeepReadOnlyFiles), "<keepreadonly>");
	m_Options.add_option("", "q", "quiet", "Reduce output to console", cxxopts::value(m_Quiet), "<quiet>");
	m_Options.add_option("", "y", "yes", "Don't query for confirmation", cxxopts::value(m_Yes), "<yes>");
	m_Options.add_option("", "", "dryrun", "Do a dry run without deleting anything", cxxopts::value(m_Dryrun), "<dryrun>");
	m_Options.add_option("output", "", "plain-progress", "Show progress using plain output", cxxopts::value(m_PlainProgress), "<progress>");
	m_Options.add_option("output", "", "verbose", "Enable verbose console output", cxxopts::value(m_Verbose), "<verbose>");
	m_Options.add_option("",
						 "",
						 "boost-workers",
						 "Increase the number of worker threads - may cause computer to be less responsive",
						 cxxopts::value(m_BoostWorkerThreads),
						 "<boostworkers>");

	m_Options.parse_positional({"directory"});
}

WipeCommand::~WipeCommand() = default;

void
WipeCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	using namespace wipe_impl;
	ZEN_UNUSED(GlobalOptions);

	ScopedSignalHandler SigIntGuard(SIGINT, SignalCallbackHandler);
#if ZEN_PLATFORM_WINDOWS
	ScopedSignalHandler SigBreakGuard(SIGBREAK, SignalCallbackHandler);
#endif

	if (!ParseOptions(argc, argv))
	{
		return;
	}

	Quiet			   = m_Quiet;
	IsVerbose		   = m_Verbose;
	ProgressMode	   = m_PlainProgress ? ConsoleProgressMode::Plain : ConsoleProgressMode::Pretty;
	BoostWorkerThreads = m_BoostWorkerThreads;

	MakeSafeAbsolutePathInPlace(m_Directory);

	if (!IsDir(m_Directory))
	{
		return;
	}

	while (!m_Yes)
	{
		const std::string Prompt = fmt::format("Do you want to wipe directory '{}'? (yes/no) ", m_Directory);
		printf("%s", Prompt.c_str());
		std::string Reponse;
		std::getline(std::cin, Reponse);
		Reponse = ToLower(Reponse);
		if (Reponse == "y" || Reponse == "yes")
		{
			m_Yes = true;
		}
		else if (Reponse == "n" || Reponse == "no")
		{
			return;
		}
	}

	CleanDirectory(m_Directory, {}, !m_KeepReadOnlyFiles, m_Dryrun);
}

}  // namespace zen