#include "M2Pack.h"
#include "PackProfile.h"
#include "M2PackRuntimeKeyProvider.h"

#include <chrono>
#include <cstring>

#include <zstd.h>

#include "EterBase/Debug.h"
#include "EterLib/BufferPool.h"

namespace
{

thread_local ZSTD_DCtx* g_m2pack_zstd_dctx = nullptr;

ZSTD_DCtx* GetThreadLocalZstdContext()
{
	if (!g_m2pack_zstd_dctx)
	{
		g_m2pack_zstd_dctx = ZSTD_createDCtx();
	}

	return g_m2pack_zstd_dctx;
}

template <typename T>
bool ReadPod(const uint8_t* bytes, std::size_t size, std::size_t& offset, T& out)
{
	if (offset + sizeof(T) > size)
	{
		return false;
	}

	memcpy(&out, bytes + offset, sizeof(T));
	offset += sizeof(T);
	return true;
}

}

bool CM2Pack::Load(const std::string& path)
{
	m_source_path = path;
	std::error_code ec;
	m_file.map(path, ec);

	if (ec)
	{
		TraceError("CM2Pack::Load: map failed for '%s': %s", path.c_str(), ec.message().c_str());
		return false;
	}

	if (m_file.size() < sizeof(TM2PackHeader))
	{
		TraceError("CM2Pack::Load: file too small '%s'", path.c_str());
		return false;
	}

	memcpy(&m_header, m_file.data(), sizeof(TM2PackHeader));
	static constexpr char kMagic[M2PACK_MAGIC_SIZE] = {'M', '2', 'P', 'A', 'C', 'K', '2', '\0'};
	if (memcmp(m_header.magic, kMagic, sizeof(kMagic)) != 0)
	{
		TraceError("CM2Pack::Load: invalid magic in '%s'", path.c_str());
		return false;
	}

	if (m_header.version != 1)
	{
		TraceError("CM2Pack::Load: unsupported version %u in '%s'", m_header.version, path.c_str());
		return false;
	}

	if (!HasM2PackRuntimeKeysForArchiveLoad(m_header.key_id))
	{
		TraceError("CM2Pack::Load: runtime master key with key_id=%u required for '%s' (active key_id=%u)",
			m_header.key_id,
			path.c_str(),
			GetM2PackActiveMasterKeyId());
		return false;
	}

	if (m_header.manifest_offset + m_header.manifest_size > m_file.size())
	{
		TraceError("CM2Pack::Load: manifest out of bounds in '%s'", path.c_str());
		return false;
	}

	m_manifest_bytes.assign(
		m_file.data() + m_header.manifest_offset,
		m_file.data() + m_header.manifest_offset + m_header.manifest_size);

	if (!ValidateManifest())
	{
		TraceError("CM2Pack::Load: manifest validation failed for '%s'", path.c_str());
		return false;
	}

	return true;
}

bool CM2Pack::ValidateManifest()
{
	std::array<uint8_t, M2PACK_HASH_SIZE> manifest_hash {};
	const auto hashStart = std::chrono::steady_clock::now();
	crypto_generichash(
		manifest_hash.data(),
		manifest_hash.size(),
		m_manifest_bytes.data(),
		m_manifest_bytes.size(),
		nullptr,
		0);
	RecordPackProfileStage(
		"m2p",
		"manifest_hash",
		m_manifest_bytes.size(),
		manifest_hash.size(),
		static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
			std::chrono::steady_clock::now() - hashStart).count()));

	if (memcmp(manifest_hash.data(), m_header.manifest_hash, manifest_hash.size()) != 0)
	{
		TraceError("CM2Pack::ValidateManifest: manifest hash mismatch");
		return false;
	}

	const auto* publicKey = GetM2PackPublicKeyForKeyId(m_header.key_id);
	if (!publicKey)
	{
		TraceError("CM2Pack::ValidateManifest: no public key configured for key_id=%u", m_header.key_id);
		return false;
	}

	const auto verifyStart = std::chrono::steady_clock::now();
	if (crypto_sign_verify_detached(
		m_header.manifest_signature,
		m_manifest_bytes.data(),
		m_manifest_bytes.size(),
		publicKey->data()) != 0)
	{
		RecordPackProfileStage(
			"m2p",
			"manifest_signature",
			m_manifest_bytes.size(),
			M2PACK_SIGNATURE_SIZE,
			static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
				std::chrono::steady_clock::now() - verifyStart).count()));
		TraceError("CM2Pack::ValidateManifest: manifest signature mismatch");
		return false;
	}
	RecordPackProfileStage(
		"m2p",
		"manifest_signature",
		m_manifest_bytes.size(),
		M2PACK_SIGNATURE_SIZE,
		static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
			std::chrono::steady_clock::now() - verifyStart).count()));

	std::size_t offset = 0;
	const auto parseStart = std::chrono::steady_clock::now();
	TM2PackManifestHeader manifest_header {};
	if (!ReadPod(m_manifest_bytes.data(), m_manifest_bytes.size(), offset, manifest_header))
	{
		return false;
	}

	m_index.clear();
	m_index.reserve(manifest_header.entry_count);

	for (uint32_t i = 0; i < manifest_header.entry_count; ++i)
	{
		TM2PackManifestEntryFixed fixed {};
		if (!ReadPod(m_manifest_bytes.data(), m_manifest_bytes.size(), offset, fixed))
		{
			return false;
		}

		if (offset + fixed.path_size > m_manifest_bytes.size())
		{
			return false;
		}

		TM2PackEntry entry;
		entry.path.assign(reinterpret_cast<const char*>(m_manifest_bytes.data() + offset), fixed.path_size);
		offset += fixed.path_size;
		entry.compression = fixed.compression;
		entry.data_offset = fixed.data_offset;
		entry.original_size = fixed.original_size;
		entry.stored_size = fixed.stored_size;
		memcpy(entry.nonce.data(), fixed.nonce, entry.nonce.size());
		memcpy(entry.plaintext_hash.data(), fixed.plaintext_hash, entry.plaintext_hash.size());

		const uint64_t payload_begin = sizeof(TM2PackHeader);
		const uint64_t payload_end = m_header.manifest_offset;
		const uint64_t begin = payload_begin + entry.data_offset;
		const uint64_t end = begin + entry.stored_size;

		if (entry.path.empty() || entry.path.find("..") != std::string::npos || entry.path[0] == '/')
		{
			TraceError("CM2Pack::ValidateManifest: invalid path '%s'", entry.path.c_str());
			return false;
		}

		if (begin > payload_end || end > payload_end || end < begin)
		{
			TraceError("CM2Pack::ValidateManifest: invalid entry bounds '%s'", entry.path.c_str());
			return false;
		}

		m_index.push_back(std::move(entry));
	}

	RecordPackProfileStage(
		"m2p",
		"manifest_parse",
		m_manifest_bytes.size(),
		m_index.size(),
		static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
			std::chrono::steady_clock::now() - parseStart).count()));
	return true;
}

bool CM2Pack::DecryptEntryPayload(const TM2PackEntry& entry, std::vector<uint8_t>& decrypted, CBufferPool* pPool)
{
	const uint64_t begin = sizeof(TM2PackHeader) + entry.data_offset;
	const auto* ciphertext = reinterpret_cast<const unsigned char*>(m_file.data() + begin);
	if (pPool)
	{
		decrypted = pPool->Acquire(entry.stored_size);
	}
	decrypted.resize(entry.stored_size);
	const auto decryptStart = std::chrono::steady_clock::now();
	unsigned long long written = 0;
	if (crypto_aead_xchacha20poly1305_ietf_decrypt(
		decrypted.data(),
		&written,
		nullptr,
		ciphertext,
		entry.stored_size,
		reinterpret_cast<const unsigned char*>(entry.path.data()),
		entry.path.size(),
		entry.nonce.data(),
		GetM2PackActiveMasterKey().data()) != 0)
	{
		if (pPool)
		{
			pPool->Release(std::move(decrypted));
		}
		return false;
	}
	RecordPackProfileStage(
		"m2p",
		"aead_decrypt",
		entry.stored_size,
		written,
		static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
			std::chrono::steady_clock::now() - decryptStart).count()));

	decrypted.resize(static_cast<std::size_t>(written));
	return true;
}

bool CM2Pack::GetFile(const TM2PackEntry& entry, std::vector<uint8_t>& result)
{
	return GetFileWithPool(entry, result, nullptr);
}

bool CM2Pack::GetFileWithPool(const TM2PackEntry& entry, std::vector<uint8_t>& result, CBufferPool* pPool)
{
	std::vector<uint8_t> compressed;
	if (!DecryptEntryPayload(entry, compressed, pPool))
	{
		TraceError("CM2Pack::GetFileWithPool: decrypt failed for '%s'", entry.path.c_str());
		return false;
	}

	switch (entry.compression)
	{
	case 0:
		result = std::move(compressed);
		break;

	case 1:
	{
		result.resize(entry.original_size);
		ZSTD_DCtx* dctx = GetThreadLocalZstdContext();
		const auto decompressStart = std::chrono::steady_clock::now();
		size_t written = ZSTD_decompressDCtx(dctx, result.data(), result.size(), compressed.data(), compressed.size());
		RecordPackProfileStage(
			"m2p",
			"zstd_decompress",
			compressed.size(),
			ZSTD_isError(written) ? 0 : written,
			static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
				std::chrono::steady_clock::now() - decompressStart).count()));
		if (pPool)
		{
			pPool->Release(std::move(compressed));
		}
		if (ZSTD_isError(written) || written != entry.original_size)
		{
			TraceError("CM2Pack::GetFileWithPool: zstd failed for '%s'", entry.path.c_str());
			return false;
		}
		break;
	}

	default:
		if (pPool)
		{
			pPool->Release(std::move(compressed));
		}
		TraceError("CM2Pack::GetFileWithPool: unsupported compression %u for '%s'", entry.compression, entry.path.c_str());
		return false;
	}

	if (!ShouldVerifyM2PackPlaintextHash())
	{
		return true;
	}

	std::array<uint8_t, M2PACK_HASH_SIZE> plain_hash {};
	const auto hashStart = std::chrono::steady_clock::now();
	crypto_generichash(
		plain_hash.data(),
		plain_hash.size(),
		result.data(),
		result.size(),
		nullptr,
		0);
	RecordPackProfileStage(
		"m2p",
		"plaintext_hash",
		result.size(),
		plain_hash.size(),
		static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::microseconds>(
			std::chrono::steady_clock::now() - hashStart).count()));

	if (memcmp(plain_hash.data(), entry.plaintext_hash.data(), plain_hash.size()) != 0)
	{
		TraceError("CM2Pack::GetFileWithPool: plaintext hash mismatch for '%s'", entry.path.c_str());
		return false;
	}

	return true;
}