#include "M2PackRuntimeKeyProvider.h"

#include <algorithm>
#include <cstring>
#include <string>

#include <windows.h>

#include "EterBase/Debug.h"
#include "EterBase/Utils.h"

namespace
{

constexpr char M2PACK_SHARED_KEY_MAGIC[8] = {'M', '2', 'K', 'E', 'Y', 'S', '1', '\0'};
constexpr const char* M2PACK_DEFAULT_MAP_NAME = "Local\\M2PackSharedKeys";
constexpr const char* M2PACK_ENV_MASTER_KEY = "M2PACK_MASTER_KEY_HEX";
constexpr const char* M2PACK_ENV_PUBLIC_KEY = "M2PACK_SIGN_PUBKEY_HEX";
constexpr const char* M2PACK_ENV_MAP_NAME = "M2PACK_KEY_MAP";
constexpr const char* M2PACK_ENV_KEY_ID = "M2PACK_KEY_ID";

#pragma pack(push, 1)
struct M2PackSharedKeys
{
	char magic[8];
	uint32_t version;
	uint32_t flags;
	uint32_t key_id;
	uint8_t master_key[M2PACK_KEY_SIZE];
	uint8_t sign_public_key[M2PACK_PUBLIC_KEY_SIZE];
};
#pragma pack(pop)

struct RuntimeKeyState
{
	uint32_t master_key_id = 0;
	std::array<uint8_t, M2PACK_KEY_SIZE> master_key {};
	std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE> public_key {};
	bool runtime_master_key = false;
	bool runtime_public_key = false;
	bool initialized = false;
};

RuntimeKeyState g_state;

uint8_t HexNibble(char ch)
{
	if (ch >= '0' && ch <= '9')
		return static_cast<uint8_t>(ch - '0');
	if (ch >= 'a' && ch <= 'f')
		return static_cast<uint8_t>(10 + ch - 'a');
	if (ch >= 'A' && ch <= 'F')
		return static_cast<uint8_t>(10 + ch - 'A');
	return 0xff;
}

template <std::size_t N>
bool ParseHexInto(std::string value, std::array<uint8_t, N>& out)
{
	value.erase(std::remove_if(value.begin(), value.end(), [](unsigned char ch) {
		return std::isspace(ch) != 0;
	}), value.end());

	if (value.size() != N * 2)
		return false;

	for (std::size_t i = 0; i < N; ++i)
	{
		const uint8_t hi = HexNibble(value[i * 2]);
		const uint8_t lo = HexNibble(value[i * 2 + 1]);
		if (hi == 0xff || lo == 0xff)
			return false;
		out[i] = static_cast<uint8_t>((hi << 4) | lo);
	}

	return true;
}

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

uint32_t ParseUInt32(const std::string& value)
{
	if (value.empty())
		return 0;

	try
	{
		return static_cast<uint32_t>(std::stoul(value));
	}
	catch (...)
	{
		return 0;
	}
}

const std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE>* FindCompiledPublicKey(uint32_t keyId)
{
	for (std::size_t i = 0; i < M2PACK_SIGN_KEY_IDS.size(); ++i)
	{
		if (M2PACK_SIGN_KEY_IDS[i] == keyId)
			return &M2PACK_SIGN_PUBLIC_KEYS[i];
	}

	return nullptr;
}

bool LoadFromSharedMapping(const std::string& mappingName)
{
	const HANDLE mapping = OpenFileMappingA(FILE_MAP_READ, FALSE, mappingName.c_str());
	if (!mapping)
		return false;

	void* view = MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, sizeof(M2PackSharedKeys));
	if (!view)
	{
		CloseHandle(mapping);
		return false;
	}

	M2PackSharedKeys blob {};
	memcpy(&blob, view, sizeof(blob));

	UnmapViewOfFile(view);
	CloseHandle(mapping);

	if (memcmp(blob.magic, M2PACK_SHARED_KEY_MAGIC, sizeof(blob.magic)) != 0 || blob.version != 1)
	{
		TraceError("M2Pack key mapping exists but has invalid header");
		return false;
	}

	memcpy(g_state.master_key.data(), blob.master_key, g_state.master_key.size());
	g_state.master_key_id = blob.key_id;
	memcpy(g_state.public_key.data(), blob.sign_public_key, g_state.public_key.size());
	g_state.runtime_master_key = true;
	g_state.runtime_public_key = true;
	return true;
}

void ApplyCommandLineOption(const std::string& key, const std::string& value)
{
	if (key == "--m2pack-key-hex")
	{
		if (ParseHexInto(value, g_state.master_key))
		{
			g_state.runtime_master_key = true;
		}
		else
		{
			TraceError("Invalid value for --m2pack-key-hex");
		}
		return;
	}

	if (key == "--m2pack-key-id")
	{
		const auto parsed = ParseUInt32(value);
		if (parsed != 0)
		{
			g_state.master_key_id = parsed;
		}
		else
		{
			TraceError("Invalid value for --m2pack-key-id");
		}
		return;
	}

	if (key == "--m2pack-pubkey-hex")
	{
		if (ParseHexInto(value, g_state.public_key))
		{
			g_state.runtime_public_key = true;
		}
		else
		{
			TraceError("Invalid value for --m2pack-pubkey-hex");
		}
		return;
	}

	if (key == "--m2pack-key-map")
	{
		LoadFromSharedMapping(value);
		return;
	}
}

} // namespace

bool InitializeM2PackRuntimeKeyProvider(const char* commandLine)
{
	if (g_state.initialized)
		return true;

	g_state = RuntimeKeyState {};

	const std::string mapName = GetEnvString(M2PACK_ENV_MAP_NAME);
	if (!mapName.empty())
	{
		LoadFromSharedMapping(mapName);
	}
	else
	{
		LoadFromSharedMapping(M2PACK_DEFAULT_MAP_NAME);
	}

	const std::string envMaster = GetEnvString(M2PACK_ENV_MASTER_KEY);
	if (!envMaster.empty())
	{
		if (ParseHexInto(envMaster, g_state.master_key))
			g_state.runtime_master_key = true;
		else
			TraceError("Invalid M2PACK_MASTER_KEY_HEX value");
	}

	const std::string envKeyId = GetEnvString(M2PACK_ENV_KEY_ID);
	if (!envKeyId.empty())
	{
		const auto parsed = ParseUInt32(envKeyId);
		if (parsed != 0)
			g_state.master_key_id = parsed;
		else
			TraceError("Invalid M2PACK_KEY_ID value");
	}

	const std::string envPublic = GetEnvString(M2PACK_ENV_PUBLIC_KEY);
	if (!envPublic.empty())
	{
		if (ParseHexInto(envPublic, g_state.public_key))
			g_state.runtime_public_key = true;
		else
			TraceError("Invalid M2PACK_SIGN_PUBKEY_HEX value");
	}

	int argc = 0;
	PCHAR* argv = CommandLineToArgv(const_cast<PCHAR>(commandLine ? commandLine : ""), &argc);
	if (argv)
	{
		for (int i = 0; i < argc; ++i)
		{
			const std::string key = argv[i];
			if ((key == "--m2pack-key-hex" || key == "--m2pack-pubkey-hex" || key == "--m2pack-key-map" || key == "--m2pack-key-id") && i + 1 < argc)
			{
				ApplyCommandLineOption(key, argv[i + 1]);
				++i;
			}
		}

		SAFE_FREE_GLOBAL(argv);
	}

	if (g_state.runtime_master_key || g_state.runtime_public_key)
	{
		Tracef("M2Pack runtime key provider: override active (master=%d public=%d)\n",
			g_state.runtime_master_key ? 1 : 0,
			g_state.runtime_public_key ? 1 : 0);
	}
	else
	{
		Tracef("M2Pack runtime key provider: no runtime master key available; .m2p loading will be denied\n");
	}

	g_state.initialized = true;
	return true;
}

const std::array<uint8_t, M2PACK_KEY_SIZE>& GetM2PackActiveMasterKey()
{
	return g_state.master_key;
}

const std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE>* GetM2PackPublicKeyForKeyId(uint32_t keyId)
{
	if (g_state.runtime_public_key && g_state.master_key_id == keyId)
	{
		return &g_state.public_key;
	}

	return FindCompiledPublicKey(keyId);
}

uint32_t GetM2PackActiveMasterKeyId()
{
	return g_state.master_key_id;
}

bool HasM2PackRuntimeMasterKey()
{
	return g_state.runtime_master_key;
}

bool HasM2PackRuntimeKeysForArchiveLoad(uint32_t keyId)
{
	return g_state.runtime_master_key && g_state.master_key_id == keyId;
}

bool IsM2PackUsingRuntimeMasterKey()
{
	return g_state.runtime_master_key;
}

bool IsM2PackUsingRuntimePublicKey()
{
	return g_state.runtime_public_key;
}