#include "Archive.h"

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <memory>
#include <vector>

#include <zlib.h>


namespace Archive
{


using UniqueFileHandle = std::unique_ptr<std::FILE, decltype(&std::fclose)>;


static inline std::size_t BytesToBlocks(std::size_t bytes);
static inline std::size_t ConvertPathSeparator(const char* src, char* dest);


static constexpr std::uint32_t VERSION {0x0010};
static constexpr std::size_t BLOCK_SIZE {1024};
static constexpr std::size_t DIRECTORY_MARGIN {64};
static constexpr int UNUSED_BLOCK {-1};
static constexpr int DIRECTORY_BLOCK {-2};


Archive::Archive(const char* p) :
	path {p},
	header {},
	entries {}
{
	UniqueFileHandle file {std::fopen(path, "rb"), &std::fclose};
	if (!file)
		throw "could not open archive";
	std::size_t length = std::fread(&header, sizeof(Header), 1, file.get());
	if (1 != length)
		throw "could not read";
	if (VERSION != header.version)
		throw "bad version";
	int err = std::fseek(file.get(), sizeof(Header) + header.directory.offset, SEEK_SET);
	if (-1 == err)
		throw "could not find directory in archive";
	std::vector<Bytef> compressed(header.directory.compressed_size);
	length = std::fread(compressed.data(), 1, header.directory.compressed_size, file.get());
	if (header.directory.compressed_size != length)
		throw "could not read compressed directory";
	const std::size_t total_entries = header.total_entries + DIRECTORY_MARGIN;
	entries.resize(header.total_entries);
	entries.reserve(total_entries);  // In original impl entries were switched on/off with their uncompressed size value.
	uLongf uncompressed_size = sizeof(Entry) * total_entries;
	err = uncompress(
		reinterpret_cast<Bytef*>(entries.data()),
		&uncompressed_size,
		compressed.data(),
		header.directory.compressed_size);
	if (Z_OK != err)
		throw "could not uncompress directory";
	UpdateBlockMap();
}


int
Archive::Expand(const int index, std::uint8_t*& buffer, const bool null_terminated) const
{
	if (0 > index && static_cast<int>(header.total_entries) <= index)
		return -1;  // Can be replaced by error-checked access to std::vector holding entries.
	UniqueFileHandle file {std::fopen(path, "rb"), &std::fclose};
	if (!file)
		return -1;
	const auto& entry = entries[index];
	std::vector<Bytef> compressed(entry.compressed_size);
	int err = std::fseek(file.get(), sizeof(Header) + entry.offset, SEEK_SET);
	if (-1 == err)
		return -1;
	const std::size_t length = std::fread(compressed.data(), 1, entry.compressed_size, file.get());
	if (entry.compressed_size != length)
		return -1;
	uLongf output_length = entry.original_size;
	auto uncompressed = std::make_unique<std::uint8_t[]>(output_length + null_terminated);
	err = uncompress(uncompressed.get(), &output_length, compressed.data(), entry.compressed_size);
	if (Z_OK != err)
		return -1;
	buffer = uncompressed.release();
	if (null_terminated)
		buffer[output_length] = 0;
	return output_length;
}


int
Archive::Expand(const char* filepath, std::uint8_t*& buffer, const bool null_terminated) const
{
	const int index = Find(filepath);
	if (-1 == index)
		return -1;
	return Expand(index, buffer, null_terminated);
}


int
Archive::Find(const char* filepath) const
{
	std::vector<char> path(std::strlen(filepath) + 1);
	ConvertPathSeparator(filepath, path.data());
	for (std::size_t i = 0; i < entries.size(); ++i) {
		if (0 == std::strcmp(entries[i].name, path.data()))
			return i;
	}
	return -1;
}


void
Archive::UpdateBlockMap()
{
	const std::size_t last_directory_block = header.directory.offset / BLOCK_SIZE + DirectoryBlocks() - 1;
	std::size_t last_block = last_directory_block;
	for (const auto& entry : entries) {
		const std::size_t block = entry.offset / BLOCK_SIZE + BytesToBlocks(entry.compressed_size) - 1;
		if (last_block < block)
			last_block = block;
	}
	blocks.clear();
	blocks.resize(last_block + 1, UNUSED_BLOCK);
	for (std::size_t n = header.directory.offset / BLOCK_SIZE; n <= last_directory_block; ++n)
		blocks[n] = DIRECTORY_BLOCK;
	for (std::size_t i = 0; i < entries.size(); ++i) {
		const std::size_t from = entries[i].offset / BLOCK_SIZE;
		const std::size_t exclusive_to = from + BytesToBlocks(entries[i].compressed_size);
		for (std::size_t n = from; n < exclusive_to; ++n)
			blocks[n] = i;
	}
}


void
Archive::PrintNamesOfEntries() const
{
	for (const auto& entry : entries)
		std::printf("%s\n", entry.name);
}


void
Archive::PrintBlocks() const
{
	for (std::size_t i = 0; i < blocks.size(); ++i) {
		std::printf("%ld\t", i);
		switch (blocks[i]) {
		case UNUSED_BLOCK:
			std::printf("(unused)\n");
			break;
		case DIRECTORY_BLOCK:
			std::printf("(dir)\n");
			break;
		default:
			std::printf("%s\n", entries[blocks[i]].name);
		}
	}
}


std::size_t
Archive::DirectoryBlocks() const
{
	const std::size_t blocks = BytesToBlocks(header.total_entries * sizeof(Entry));
	return blocks == 0 ? 1 : blocks;
}


std::size_t
BytesToBlocks(const std::size_t bytes)
{
	const std::size_t full = bytes / BLOCK_SIZE;
	const std::size_t partial = (bytes % BLOCK_SIZE) > 0;
	return full + partial;
}


std::size_t
ConvertPathSeparator(const char* src, char* dest)
{
	std::size_t i;
	for (i = 0; 0 != src[i]; ++i) {
		if ('\\' == src[i])
			dest[i] = '/';
		else
			dest[i] = src[i];
	}
	return i;
}


}  // namespace Archive