#include "eterpack.h"

#include <sodium.h>
#include <zstd.h>

#include <algorithm>
#include <cstring>
#include <fstream>

namespace mtpack {

namespace {

constexpr size_t HEADER_SIZE = 8 + 8 + PACK_NONCE_SIZE; // 40
// entry tail after the name field: offset,file_size,compressed_size,encryption,nonce
constexpr size_t ENTRY_TAIL = 8 + 8 + 8 + 1 + PACK_NONCE_SIZE; // 49

void xchacha20(uint8_t *data, size_t len, const uint8_t *nonce) {
	crypto_stream_xchacha20_xor(data, data, len, nonce, PACK_KEY);
}

// Parse one decrypted entry blob of `field + ENTRY_TAIL` bytes.
Entry parse_entry(const uint8_t *p, int field) {
	Entry e;
	size_t nlen = strnlen(reinterpret_cast<const char *>(p), field);
	e.name.assign(reinterpret_cast<const char *>(p), nlen);
	const uint8_t *q = p + field;
	std::memcpy(&e.offset, q, 8);
	q += 8;
	std::memcpy(&e.file_size, q, 8);
	q += 8;
	std::memcpy(&e.compressed_size, q, 8);
	q += 8;
	e.encryption = *q++;
	std::memcpy(e.nonce, q, PACK_NONCE_SIZE);
	return e;
}

bool plausible(const Entry &e, uint64_t data_begin, uint64_t file_total) {
	if (e.encryption > 1) {
		return false;
	}
	// empty files are legal; a zstd frame is never 0 bytes though
	if (e.compressed_size == 0 || e.compressed_size > file_total) {
		return false;
	}
	if (data_begin + e.offset + e.compressed_size > file_total) {
		return false;
	}
	if (e.name.empty()) {
		return false;
	}
	// filenames may be CP949 (Korean) — reject only ASCII control bytes.
	for (unsigned char c : e.name) {
		if (c < 0x20) {
			return false;
		}
	}
	return true;
}

} // namespace

std::string EterPack::norm(std::string s) {
	std::replace(s.begin(), s.end(), '\\', '/');
	std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
	return s;
}

void EterPack::close() {
	m_file.clear();
	m_index.clear();
	m_by_name.clear();
	m_data_begin = 0;
}

bool EterPack::open(const std::string &path, std::string *err) {
	close();
	if (sodium_init() < 0) {
		if (err) *err = "sodium_init failed";
		return false;
	}
	std::ifstream f(path, std::ios::binary);
	if (!f) {
		if (err) *err = "cannot open " + path;
		return false;
	}
	f.seekg(0, std::ios::end);
	std::streamoff sz = f.tellg();
	f.seekg(0);
	if (sz < static_cast<std::streamoff>(HEADER_SIZE)) {
		if (err) *err = "file too small";
		return false;
	}
	m_file.resize(static_cast<size_t>(sz));
	f.read(reinterpret_cast<char *>(m_file.data()), sz);

	uint64_t entry_num = 0;
	std::memcpy(&entry_num, m_file.data(), 8);
	std::memcpy(&m_data_begin, m_file.data() + 8, 8);
	const uint8_t *hnonce = m_file.data() + 16;

	if (entry_num == 0 || m_data_begin < HEADER_SIZE || m_data_begin > m_file.size()) {
		if (err) *err = "bad header";
		return false;
	}
	const uint64_t index_bytes = m_data_begin - HEADER_SIZE;

	// Candidate name-field sizes: derived first, then the known platform values.
	std::vector<int> candidates;
	if (index_bytes % entry_num == 0) {
		int64_t es = static_cast<int64_t>(index_bytes / entry_num);
		if (es > static_cast<int64_t>(ENTRY_TAIL) + 1) {
			candidates.push_back(static_cast<int>(es - ENTRY_TAIL));
		}
	}
	for (int v : {261, 4097, 1025, 256}) {
		candidates.push_back(v);
	}

	for (int field : candidates) {
		const size_t entry_size = static_cast<size_t>(field) + ENTRY_TAIL;
		if (HEADER_SIZE + entry_num * entry_size > m_file.size()) {
			continue;
		}
		std::vector<Entry> idx;
		idx.reserve(entry_num);
		bool ok = true;
		std::vector<uint8_t> blob(entry_size);
		for (uint64_t i = 0; i < entry_num && ok; ++i) {
			std::memcpy(blob.data(), m_file.data() + HEADER_SIZE + i * entry_size, entry_size);
			xchacha20(blob.data(), entry_size, hnonce);
			Entry e = parse_entry(blob.data(), field);
			if (!plausible(e, m_data_begin, m_file.size())) {
				ok = false;
				break;
			}
			idx.push_back(std::move(e));
		}
		if (ok) {
			m_index = std::move(idx);
			m_name_field = field;
			for (size_t i = 0; i < m_index.size(); ++i) {
				m_by_name[norm(m_index[i].name)] = i;
			}
			return true;
		}
	}
	if (err) *err = "could not resolve index layout (name-field guess failed)";
	return false;
}

std::vector<std::string> EterPack::names() const {
	std::vector<std::string> v;
	v.reserve(m_index.size());
	for (const auto &e : m_index) {
		v.push_back(e.name);
	}
	return v;
}

bool EterPack::read(const std::string &name, std::vector<uint8_t> &out, std::string *err) const {
	auto it = m_by_name.find(norm(name));
	if (it == m_by_name.end()) {
		if (err) *err = "not in pack: " + name;
		return false;
	}
	const Entry &e = m_index[it->second];
	const uint8_t *src = m_file.data() + m_data_begin + e.offset;

	std::vector<uint8_t> comp(e.compressed_size);
	std::memcpy(comp.data(), src, e.compressed_size);
	if (e.encryption == 1) {
		crypto_stream_xchacha20_xor(comp.data(), comp.data(), comp.size(), e.nonce, PACK_KEY);
	}

	out.resize(e.file_size);
	size_t n = ZSTD_decompress(out.data(), out.size(), comp.data(), comp.size());
	if (ZSTD_isError(n) || n != e.file_size) {
		if (err) *err = std::string("zstd: ") + (ZSTD_isError(n) ? ZSTD_getErrorName(n) : "size mismatch");
		return false;
	}
	return true;
}

// --- writer ----------------------------------------------------------------

bool write_pack(const std::string &out_path, const std::vector<InputFile> &files, bool encrypt,
		std::string *err) {
	if (sodium_init() < 0) {
		if (err) *err = "sodium_init failed";
		return false;
	}
	const int field = PACK_NAME_FIELD_DEFAULT;
	const size_t entry_size = static_cast<size_t>(field) + ENTRY_TAIL;
	const uint64_t entry_num = files.size();
	const uint64_t data_begin = HEADER_SIZE + entry_num * entry_size;

	uint8_t header_nonce[PACK_NONCE_SIZE];
	randombytes_buf(header_nonce, sizeof(header_nonce));

	std::vector<uint8_t> index(entry_num * entry_size, 0);
	std::vector<uint8_t> data;
	uint64_t cursor = 0;

	for (uint64_t i = 0; i < entry_num; ++i) {
		const InputFile &in = files[i];
		size_t bound = ZSTD_compressBound(in.data.size());
		std::vector<uint8_t> comp(bound);
		size_t clen = ZSTD_compress(comp.data(), comp.size(), in.data.data(), in.data.size(), 3);
		if (ZSTD_isError(clen)) {
			if (err) *err = std::string("zstd compress: ") + ZSTD_getErrorName(clen);
			return false;
		}
		comp.resize(clen);

		Entry e;
		e.name = in.name;
		e.offset = cursor;
		e.file_size = in.data.size();
		e.compressed_size = clen;
		e.encryption = encrypt ? 1 : 0;
		if (encrypt) {
			randombytes_buf(e.nonce, sizeof(e.nonce));
			crypto_stream_xchacha20_xor(comp.data(), comp.data(), comp.size(), e.nonce, PACK_KEY);
		}

		// serialize entry (plaintext), then encrypt the whole index blob at the end
		uint8_t *p = index.data() + i * entry_size;
		std::string nm = e.name;
		std::replace(nm.begin(), nm.end(), '\\', '/');
		std::memcpy(p, nm.data(), std::min<size_t>(nm.size(), field - 1));
		uint8_t *q = p + field;
		std::memcpy(q, &e.offset, 8);
		q += 8;
		std::memcpy(q, &e.file_size, 8);
		q += 8;
		std::memcpy(q, &e.compressed_size, 8);
		q += 8;
		*q++ = e.encryption;
		std::memcpy(q, e.nonce, PACK_NONCE_SIZE);

		data.insert(data.end(), comp.begin(), comp.end());
		cursor += clen;
	}

	// encrypt the index with the header nonce
	for (uint64_t i = 0; i < entry_num; ++i) {
		crypto_stream_xchacha20_xor(index.data() + i * entry_size, index.data() + i * entry_size,
				entry_size, header_nonce, PACK_KEY);
	}

	std::ofstream f(out_path, std::ios::binary | std::ios::trunc);
	if (!f) {
		if (err) *err = "cannot write " + out_path;
		return false;
	}
	f.write(reinterpret_cast<const char *>(&entry_num), 8);
	f.write(reinterpret_cast<const char *>(&data_begin), 8);
	f.write(reinterpret_cast<const char *>(header_nonce), PACK_NONCE_SIZE);
	f.write(reinterpret_cast<const char *>(index.data()), static_cast<std::streamsize>(index.size()));
	f.write(reinterpret_cast<const char *>(data.data()), static_cast<std::streamsize>(data.size()));
	return static_cast<bool>(f);
}

} // namespace mtpack