#include "PackProfile.h"

#include <algorithm>
#include <array>
#include <cctype>
#include <chrono>
#include <condition_variable>
#include <filesystem>
#include <fstream>
#include <map>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

#include <windows.h>

#include "EterBase/Utils.h"

namespace
{

using Clock = std::chrono::steady_clock;

constexpr const char* kPackProfileEnv = "M2PACK_PROFILE";
constexpr const char* kPackProfileOutput = "log/pack_profile.txt";
constexpr std::size_t kTopListLimit = 12;
constexpr auto kSnapshotFlushInterval = std::chrono::seconds(1);

struct LoadStats
{
	std::uint64_t count = 0;
	std::uint64_t failures = 0;
	std::uint64_t output_bytes = 0;
	std::uint64_t elapsed_us = 0;
};

struct MountStats
{
	std::uint64_t count = 0;
	std::uint64_t failures = 0;
	std::uint64_t entry_count = 0;
	std::uint64_t elapsed_us = 0;
};

struct StageStats
{
	std::uint64_t count = 0;
	std::uint64_t input_bytes = 0;
	std::uint64_t output_bytes = 0;
	std::uint64_t elapsed_us = 0;
};

struct PhaseMarker
{
	std::string phase;
	std::uint64_t elapsed_us = 0;
};

struct PackProfileState
{
	bool initialized = false;
	bool enabled = false;
	std::string current_phase = "startup";
	Clock::time_point start_time {};
	std::vector<PhaseMarker> markers;
	std::map<std::string, LoadStats> load_by_format;
	std::map<std::string, LoadStats> load_by_phase_format;
	std::map<std::string, LoadStats> load_by_phase_pack;
	std::map<std::string, LoadStats> load_by_extension;
	std::map<std::string, MountStats> mount_by_format;
	std::map<std::string, MountStats> mount_by_pack;
	std::map<std::string, StageStats> stage_by_key;
	bool snapshot_dirty = false;
	bool stop_snapshot_thread = false;
	std::condition_variable snapshot_cv;
	std::thread snapshot_thread;
	std::mutex mutex;

	~PackProfileState()
	{
		StopSnapshotThread();
		if (enabled)
		{
			std::lock_guard<std::mutex> lock(mutex);
			WriteReportLocked("shutdown");
		}
	}

	void StartSnapshotThread()
	{
		snapshot_thread = std::thread([this]() {
			SnapshotLoop();
		});
	}

	void StopSnapshotThread()
	{
		{
			std::lock_guard<std::mutex> lock(mutex);
			stop_snapshot_thread = true;
		}
		snapshot_cv.notify_all();
		if (snapshot_thread.joinable())
		{
			snapshot_thread.join();
		}
	}

	void SnapshotLoop()
	{
		std::unique_lock<std::mutex> lock(mutex);
		for (;;)
		{
			const bool shouldStop = snapshot_cv.wait_for(
				lock,
				kSnapshotFlushInterval,
				[this]() { return stop_snapshot_thread; });
			if (shouldStop)
			{
				return;
			}

			if (!enabled || !snapshot_dirty)
			{
				continue;
			}

			WriteReportLocked("periodic");
			snapshot_dirty = false;
		}
	}

	void WriteReportLocked(std::string_view reason) const
	{
		std::error_code ec;
		std::filesystem::create_directories("log", ec);

		std::ofstream output(kPackProfileOutput, std::ios::binary | std::ios::trunc);
		if (!output)
		{
			return;
		}

		output << "PACK PROFILE REPORT\n";
		output << "reason=" << reason << "\n";
		output << "current_phase=" << current_phase << "\n";

		const auto uptime_us = start_time == Clock::time_point {}
			? 0
			: static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
				Clock::now() - start_time).count());
		output << "uptime_ms=" << FormatMs(uptime_us) << "\n\n";

		output << "[phase_markers]\n";
		for (const auto& marker : markers)
		{
			output << marker.phase << "\t" << FormatMs(marker.elapsed_us) << " ms\n";
		}
		output << "\n";

		WriteMountSection(output, "[mounts_by_format]", mount_by_format);
		WriteMountSection(output, "[mounts_by_pack]", mount_by_pack, kTopListLimit);
		WriteLoadSection(output, "[loads_by_format]", load_by_format);
		WriteLoadSection(output, "[loads_by_phase_format]", load_by_phase_format);
		WriteLoadSection(output, "[top_phase_pack_loads]", load_by_phase_pack, kTopListLimit);
		WriteLoadSection(output, "[top_extension_loads]", load_by_extension, kTopListLimit);
		WriteStageSection(output, "[loader_stages]", stage_by_key);
	}

	static std::string FormatMs(std::uint64_t elapsed_us)
	{
		std::ostringstream out;
		out.setf(std::ios::fixed, std::ios::floatfield);
		out.precision(3);
		out << (static_cast<double>(elapsed_us) / 1000.0);
		return out.str();
	}

	static std::string ToLowerAscii(std::string value)
	{
		std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) {
			return static_cast<char>(std::tolower(ch));
		});
		return value;
	}

	static std::string ParsePackLabel(std::string_view pack_path)
	{
		const std::filesystem::path path { std::string(pack_path) };
		const auto filename = path.filename().string();
		if (!filename.empty())
		{
			return ToLowerAscii(filename);
		}

		return ToLowerAscii(std::string(pack_path));
	}

	static std::string ParseExtension(std::string_view request_path)
	{
		const std::filesystem::path path { std::string(request_path) };
		auto extension = path.extension().string();
		if (extension.empty())
		{
			return "<none>";
		}

		return ToLowerAscii(std::move(extension));
	}

	static std::string MakePhaseKey(std::string_view phase, std::string_view suffix)
	{
		std::string key;
		key.reserve(phase.size() + suffix.size() + 3);
		key.append(phase);
		key.append(" | ");
		key.append(suffix);
		return key;
	}

	template <typename MapType>
	static std::vector<std::pair<std::string, typename MapType::mapped_type>> SortByTimeDesc(const MapType& values)
	{
		using StatsType = typename MapType::mapped_type;
		std::vector<std::pair<std::string, StatsType>> sorted;
		sorted.reserve(values.size());
		for (const auto& [key, stats] : values)
		{
			sorted.emplace_back(key, stats);
		}
		std::sort(sorted.begin(), sorted.end(), [](const auto& left, const auto& right) {
			if (left.second.elapsed_us != right.second.elapsed_us)
			{
				return left.second.elapsed_us > right.second.elapsed_us;
			}

			return left.first < right.first;
		});
		return sorted;
	}

	static void WriteLoadSection(std::ofstream& output, std::string_view header, const std::map<std::string, LoadStats>& values, std::size_t limit = 0)
	{
		output << header << "\n";
		const auto sorted = SortByTimeDesc(values);
		const std::size_t count = limit == 0 ? sorted.size() : std::min(limit, sorted.size());
		for (std::size_t i = 0; i < count; ++i)
		{
			const auto& [key, stats] = sorted[i];
			output
				<< key
				<< "\tcount=" << stats.count
				<< "\tfail=" << stats.failures
				<< "\tbytes=" << stats.output_bytes
				<< "\ttime_ms=" << FormatMs(stats.elapsed_us)
				<< "\n";
		}
		output << "\n";
	}

	static void WriteMountSection(std::ofstream& output, std::string_view header, const std::map<std::string, MountStats>& values, std::size_t limit = 0)
	{
		output << header << "\n";
		std::vector<std::pair<std::string, MountStats>> sorted(values.begin(), values.end());
		std::sort(sorted.begin(), sorted.end(), [](const auto& left, const auto& right) {
			if (left.second.elapsed_us != right.second.elapsed_us)
			{
				return left.second.elapsed_us > right.second.elapsed_us;
			}

			return left.first < right.first;
		});

		const std::size_t count = limit == 0 ? sorted.size() : std::min(limit, sorted.size());
		for (std::size_t i = 0; i < count; ++i)
		{
			const auto& [key, stats] = sorted[i];
			output
				<< key
				<< "\tcount=" << stats.count
				<< "\tfail=" << stats.failures
				<< "\tentries=" << stats.entry_count
				<< "\ttime_ms=" << FormatMs(stats.elapsed_us)
				<< "\n";
		}
		output << "\n";
	}

	static void WriteStageSection(std::ofstream& output, std::string_view header, const std::map<std::string, StageStats>& values)
	{
		output << header << "\n";
		std::vector<std::pair<std::string, StageStats>> sorted(values.begin(), values.end());
		std::sort(sorted.begin(), sorted.end(), [](const auto& left, const auto& right) {
			if (left.second.elapsed_us != right.second.elapsed_us)
			{
				return left.second.elapsed_us > right.second.elapsed_us;
			}

			return left.first < right.first;
		});

		for (const auto& [key, stats] : sorted)
		{
			output
				<< key
				<< "\tcount=" << stats.count
				<< "\tin=" << stats.input_bytes
				<< "\tout=" << stats.output_bytes
				<< "\ttime_ms=" << FormatMs(stats.elapsed_us)
				<< "\n";
		}
		output << "\n";
	}
};

PackProfileState g_pack_profile;

bool ParseBoolFlag(std::string value, bool& out)
{
	value.erase(std::remove_if(value.begin(), value.end(), [](unsigned char ch) {
		return std::isspace(ch) != 0;
	}), value.end());

	std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) {
		return static_cast<char>(std::tolower(ch));
	});

	if (value == "1" || value == "true" || value == "yes" || value == "on")
	{
		out = true;
		return true;
	}

	if (value == "0" || value == "false" || value == "no" || value == "off")
	{
		out = false;
		return true;
	}

	return false;
}

std::string GetEnvString(const char* name)
{
	char buffer[256];
	const DWORD len = GetEnvironmentVariableA(name, buffer, sizeof(buffer));
	if (len == 0 || len >= sizeof(buffer))
	{
		return {};
	}

	return std::string(buffer, buffer + len);
}

std::uint64_t CurrentElapsedUs()
{
	if (!g_pack_profile.enabled)
	{
		return 0;
	}

	return static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
		Clock::now() - g_pack_profile.start_time).count());
}

} // namespace

void InitializePackProfile(const char* commandLine)
{
	if (g_pack_profile.initialized)
	{
		return;
	}

	g_pack_profile.initialized = true;
	g_pack_profile.enabled = false;

	const std::string envProfile = GetEnvString(kPackProfileEnv);
	if (!envProfile.empty())
	{
		bool enabled = false;
		if (ParseBoolFlag(envProfile, enabled))
		{
			g_pack_profile.enabled = enabled;
		}
	}

	int argc = 0;
	PCHAR* argv = CommandLineToArgv(const_cast<PCHAR>(commandLine ? commandLine : ""), &argc);
	if (argv)
	{
		for (int i = 0; i < argc; ++i)
		{
			const std::string arg = argv[i];
			if (arg != "--m2pack-profile")
			{
				continue;
			}

			bool enabled = true;
			if (i + 1 < argc)
			{
				const std::string next = argv[i + 1];
				if (!next.starts_with("--") && ParseBoolFlag(next, enabled))
				{
					++i;
				}
			}

			g_pack_profile.enabled = enabled;
		}

		SAFE_FREE_GLOBAL(argv);
	}

	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);
	g_pack_profile.current_phase = "startup";
	g_pack_profile.start_time = Clock::now();
	g_pack_profile.markers.push_back(PhaseMarker {
		.phase = g_pack_profile.current_phase,
		.elapsed_us = 0,
	});
	g_pack_profile.WriteReportLocked("initialized");
	g_pack_profile.snapshot_dirty = false;
	g_pack_profile.stop_snapshot_thread = false;
	g_pack_profile.StartSnapshotThread();
}

bool IsPackProfileEnabled()
{
	return g_pack_profile.enabled;
}

void MarkPackProfilePhase(std::string_view phase)
{
	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);
	if (g_pack_profile.current_phase == phase)
	{
		return;
	}

	g_pack_profile.current_phase = std::string(phase);
	g_pack_profile.markers.push_back(PhaseMarker {
		.phase = g_pack_profile.current_phase,
		.elapsed_us = CurrentElapsedUs(),
	});
	g_pack_profile.WriteReportLocked(PackProfileState::MakePhaseKey("phase", phase));
	g_pack_profile.snapshot_dirty = false;
}

void RecordPackProfileMount(
	std::string_view format,
	std::string_view packPath,
	bool ok,
	std::size_t entryCount,
	std::uint64_t elapsedUs)
{
	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);

	auto& formatStats = g_pack_profile.mount_by_format[std::string(format)];
	formatStats.count += 1;
	formatStats.failures += ok ? 0 : 1;
	formatStats.entry_count += entryCount;
	formatStats.elapsed_us += elapsedUs;

	auto& packStats = g_pack_profile.mount_by_pack[PackProfileState::ParsePackLabel(packPath)];
	packStats.count += 1;
	packStats.failures += ok ? 0 : 1;
	packStats.entry_count += entryCount;
	packStats.elapsed_us += elapsedUs;
	g_pack_profile.snapshot_dirty = true;
}

void RecordPackProfileLoad(
	std::string_view format,
	std::string_view packPath,
	std::string_view requestPath,
	bool ok,
	std::uint64_t outputBytes,
	std::uint64_t elapsedUs)
{
	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);

	auto update = [ok, outputBytes, elapsedUs](LoadStats& stats) {
		stats.count += 1;
		stats.failures += ok ? 0 : 1;
		stats.output_bytes += outputBytes;
		stats.elapsed_us += elapsedUs;
	};

	update(g_pack_profile.load_by_format[std::string(format)]);
	update(g_pack_profile.load_by_phase_format[
		PackProfileState::MakePhaseKey(g_pack_profile.current_phase, format)]);
	update(g_pack_profile.load_by_phase_pack[
		PackProfileState::MakePhaseKey(g_pack_profile.current_phase, PackProfileState::ParsePackLabel(packPath))]);
	update(g_pack_profile.load_by_extension[PackProfileState::ParseExtension(requestPath)]);
	g_pack_profile.snapshot_dirty = true;
}

void RecordPackProfileStage(
	std::string_view format,
	std::string_view stage,
	std::uint64_t inputBytes,
	std::uint64_t outputBytes,
	std::uint64_t elapsedUs)
{
	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);
	auto& stats = g_pack_profile.stage_by_key[
		PackProfileState::MakePhaseKey(format, stage)];
	stats.count += 1;
	stats.input_bytes += inputBytes;
	stats.output_bytes += outputBytes;
	stats.elapsed_us += elapsedUs;
	g_pack_profile.snapshot_dirty = true;
}

void FlushPackProfileSnapshot(std::string_view reason)
{
	if (!g_pack_profile.enabled)
	{
		return;
	}

	std::lock_guard<std::mutex> lock(g_pack_profile.mutex);
	g_pack_profile.WriteReportLocked(reason);
	g_pack_profile.snapshot_dirty = false;
}