#include "StdAfx.h"

#include "fs/zip/ZipPack.h"

#include "fs/NativeFilesystem.h"
#include "fs/StreamSourceToMapSourceAdapter.h"

#include "tl/crash.h"

#include "Logger.h"

#include <tl/vector.h>
#include <tl/finally.h>

#include "tl/narrow_cast.h"
extern "C"
{
#include "zip.h"
}
#include "ZipFileSource.h"
#include "ZipFileStreamSource.h"


//////////////////////////////////////////////////////////////////////////

namespace fs
{
namespace
{
void my_zip_fclose(zip_file* file)
{
	fs_zip_fclose(file);
}
void my_zip_close(zip* file)
{
	fs_zip_close(file);
}

uint64_t computePersistentHash(const AbsPath& path)
{
	uint64_t hash = 0;
	for (const tl::string& e : path)
		tl::hash_and_combine(hash, e);
	return hash;
}
uint64_t computePersistentHash(AbsPathView path)
{
	uint64_t hash = 0;
	for (const tl::string& e : path)
		tl::hash_and_combine(hash, e);
	return hash;
}

zip_int64_t mapSourceCallback(void* userdata, void* data, zip_uint64_t len, zip_source_cmd_t cmd)
{
	ZipPack::ZipData& zipData = *(ZipPack::ZipData*)userdata;

	switch (cmd)
	{
	case ZIP_SOURCE_SUPPORTS:
		return ZIP_SOURCE_SUPPORTS_READABLE | ZIP_SOURCE_SUPPORTS_SEEKABLE;
	case ZIP_SOURCE_OPEN:
	{
		zipData.fsMapSource->seekBeg(0);
		break;
	}
	case ZIP_SOURCE_READ:
	{
		return zipData.fsMapSource->read({ (uint8_t*)data, len });
	}
	case ZIP_SOURCE_SEEK:
	{
		zip_error error;
		fs_zip_error_init(&error);
		tl::finally finally([&] { fs_zip_error_fini(&error); });
		const zip_source_args_seek* args = ZIP_SOURCE_GET_ARGS(zip_source_args_seek, data, len, &error);
		if (!args)
			return -1;
		if (args->whence == SEEK_SET)
			zipData.fsMapSource->seekBeg(args->offset);
		else if (args->whence == SEEK_CUR)
			zipData.fsMapSource->seekRel(args->offset);
		else
			return -1;
		return zipData.fsMapSource->getLastError().has_value() ? -1 : 0;
	}
	case ZIP_SOURCE_TELL:
		return zipData.fsMapSource->tell();
	case ZIP_SOURCE_STAT:
	{
		zip_stat_t* stat = (zip_stat_t*)data;
		fs_zip_stat_init(stat);
		stat->valid = ZIP_STAT_SIZE;
		stat->size = zipData.fsMapSource->getSize();
		return sizeof(zip_stat_t);
	}
	}

	return 0;
}

zip_int64_t fileSourceCallback(void* userdata, void* data, zip_uint64_t len, zip_source_cmd_t cmd)
{
	ZipPack::ZipData& zipData = *(ZipPack::ZipData*)userdata;

	switch (cmd)
	{
	case ZIP_SOURCE_SUPPORTS:
		return ZIP_SOURCE_SUPPORTS_READABLE | ZIP_SOURCE_SUPPORTS_SEEKABLE;
	case ZIP_SOURCE_OPEN:
	{
		auto result = zipData.filesystem->openSource(zipData.path);
		if (result.has_value())
			zipData.fsSource = std::move(result.value());
		else
			LOGE("Error opening zip file source: {}", result.error());
		break;
	}
	case ZIP_SOURCE_CLOSE:
	{
		if (zipData.fsSource)
			zipData.fsSource = {};
		break;
	}
	case ZIP_SOURCE_READ:
	{
		if (zipData.fsSource)
			return zipData.fsSource->read({ (uint8_t*)data, len });
		return 0;
	}
	case ZIP_SOURCE_SEEK:
	{
		zip_error error;
		fs_zip_error_init(&error);
		tl::finally finally([&] { fs_zip_error_fini(&error); });
		const zip_source_args_seek* args = ZIP_SOURCE_GET_ARGS(zip_source_args_seek, data, len, &error);
		if (!args || !zipData.fsSource)
			return -1;
		if (args->whence == SEEK_SET)
			zipData.fsSource->seekBeg(args->offset);
		else if (args->whence == SEEK_CUR)
			zipData.fsSource->seekRel(args->offset);
		else if (args->whence == SEEK_END)
			zipData.fsSource->seekBeg(zipData.fsSource->getSize() + args->offset);
		else
			return -1;
		return zipData.fsSource->getLastError().has_value() ? -1 : 0;
	}
	case ZIP_SOURCE_TELL:
		if (!zipData.fsSource)
			return -1;
		return zipData.fsSource->tell();
	case ZIP_SOURCE_STAT:
	{
		zip_stat_t* stat = (zip_stat_t*)data;
		fs_zip_stat_init(stat);
		stat->valid = ZIP_STAT_SIZE;
		if (zipData.cachedSize == uint64_t(-1))
		{
			if (auto result = zipData.filesystem->getStat(zipData.path); result.has_value())
				zipData.cachedSize = result.value().size;
		}
		stat->size = zipData.cachedSize;
		return sizeof(zip_stat_t);
	}
	}

	return 0;
}

}

//////////////////////////////////////////////////////////////////////////

ZipPack::ZipPack(tl::shared_ptr<ZipData> zipData)
	: m_zipData(std::move(zipData))
{
	createEntries(*m_zipData->zipArchive);
}

//////////////////////////////////////////////////////////////////////////

ZipPack::CreateResult ZipPack::create(tl::unique_ref<IMapSource> source)
{
	tl::shared_ptr<ZipData> zipData = tl::make_shared<ZipData>();
	zipData->fsMapSource = std::move(source);

	zip_error error;
	fs_zip_error_init(&error);
	tl::finally finally([&] { fs_zip_error_fini(&error); });

	zip_source* zipSource = fs_zip_source_function_create(mapSourceCallback, zipData.get(), &error);

	zip* _zipArchive = fs_zip_open_from_source(zipSource, 0, &error);
	if (!_zipArchive)
	{
		zip_source_free(zipSource);
		const int ze = fs_zip_error_code_zip(&error);
		const int se = fs_zip_error_code_system(&error);
		return tl::make_error<Error>(ErrorCode::SystemError, "Failed to create zip archive: {}/{} ('{}')", ze, se, zip_error_strerror(&error));
	}
	tl::unique_ref<zip> zipArchive = tl::unique_ref<zip>(_zipArchive, &my_zip_close);

	zipData->zipArchive = std::move(zipArchive);

	return tl::unique_ref<ZipPack>(new ZipPack(std::move(zipData)));
}

//////////////////////////////////////////////////////////////////////////

ZipPack::CreateResult ZipPack::create(tl::lent_ref<const IFilesystem> filesystem, AbsPathView path)
{
	tl::shared_ptr<ZipData> zipData = tl::make_shared<ZipData>();
	zipData->filesystem = std::move(filesystem);
	zipData->path = path;

	zip_error error;
	fs_zip_error_init(&error);
	tl::finally finally([&] { fs_zip_error_fini(&error); });

	zip_source* zipSource = fs_zip_source_function_create(fileSourceCallback, zipData.get(), &error);
	zip* _zipArchive = fs_zip_open_from_source(zipSource, 0, &error);
	if (!_zipArchive)
	{
		zip_source_free(zipSource);
		const int ze = fs_zip_error_code_zip(&error);
		const int se = fs_zip_error_code_system(&error);
		return tl::make_error<Error>(ErrorCode::SystemError, "Failed to create zip archive: {}/{} ('{}')", ze, se, zip_error_strerror(&error));
	}

	tl::unique_ref<zip> zipArchive = tl::unique_ref<zip>(_zipArchive, &my_zip_close);

	zipData->zipArchive = std::move(zipArchive);

	return tl::unique_ref<ZipPack>(new ZipPack(std::move(zipData)));
}

//////////////////////////////////////////////////////////////////////////

void ZipPack::createEntries(zip& archive)
{
	m_files.clear();
	m_folders.clear();
	m_pathToIndex.clear();
	m_childrenIndices.clear();

	const size_t readerEntryCount = fs_zip_get_num_entries(&archive, 0);

	m_files.reserve(readerEntryCount + 1);
	m_folders.reserve(readerEntryCount / 2);
	m_pathToIndex.reserve(readerEntryCount + 1);

	// add the root
	m_folders.emplace_back();
	m_pathToIndex.emplace(computePersistentHash(AbsPath("/")), EntryIndex(true, 0));

	tl::vector<tl::vector<EntryIndex>> childrenIndices;
	childrenIndices.reserve(readerEntryCount + 1);

	AbsPath path;

	zip_stat_t stat;
	for (size_t readerEntryIndex = 0; readerEntryIndex < readerEntryCount; readerEntryIndex++)
	{
		if (fs_zip_stat_index(&archive, readerEntryIndex, 0, &stat) != 0)
		{
			LOGE("ZipPack entry '{}' stat failed", readerEntryIndex);
			continue;
		}

		const char* name = stat.name;
		if (!name || name[0] == 0)
		{
			LOGE("ZipPack entry '{}' name empty", readerEntryIndex);
			continue;
		}

		if (stat.comp_method != ZIP_CM_STORE && stat.comp_method != ZIP_CM_DEFLATE)
		{
			LOGE("ZipPack entry '{}' because it uses an unsupported compression", name);
			continue;
		}
		if (stat.encryption_method != ZIP_EM_NONE)
		{
			LOGE("ZipPack entry '{}' because it uses encryption", name);
			continue;
		}

		bool isFolder = false;
		if (name[0] == '/' || name[0] == '\\')
			isFolder = true;

		path = PosixRootTag::value();
		path += name;

		AbsPath parentPath = path;
		parentPath.pop_back();

		const uint16_t parentIndex = getOrAddFolder(std::move(parentPath), childrenIndices);
		uint16_t index = 0;

		if (isFolder)
		{
			index = tl::narrow<uint16_t>(m_folders.size());
			if (index >= 0xFFFF)
				TL_CRASH("Too many folders in zip pack");

			m_folders.push_back(Folder());

			Folder& folder = m_folders.back();
			folder.name = path.back();
			folder.zipIndex = static_cast<uint16_t>(readerEntryIndex);
			folder.parentIndex = parentIndex;
		}
		else
		{
			index = tl::narrow<uint16_t>(m_files.size());
			if (index >= 0xFFFF)
				TL_CRASH("Too many files/folders in zip pack");

			m_files.push_back(File());

			File& file = m_files.back();
			file.name = path.back();
			file.zipIndex = static_cast<uint16_t>(readerEntryIndex);
		}

		if (parentIndex >= childrenIndices.size())
			childrenIndices.resize(parentIndex + 1);

		tl::vector<EntryIndex>& children = childrenIndices[parentIndex];
		children.push_back(EntryIndex(isFolder, index));
		m_pathToIndex.emplace(computePersistentHash(path), EntryIndex(isFolder, index));
	}

	m_childrenIndices.reserve(8192);

	// compact the children indices
	for (size_t i = 0; i < childrenIndices.size(); i++)
	{
		Folder& folder = m_folders[i];
		tl::vector<EntryIndex>& children = childrenIndices[i];
		folder.childrenStartIndex = (uint16_t)m_childrenIndices.size();
		folder.childrenCount = (uint16_t)children.size();
		if (folder.childrenCount > 0)
		{
			m_childrenIndices.resize(m_childrenIndices.size() + children.size());
			memcpy(&m_childrenIndices[folder.childrenStartIndex], children.data(), folder.childrenCount * sizeof(EntryIndex));
		}
	}

	m_files.shrink_to_fit();
	m_folders.shrink_to_fit();
	m_childrenIndices.shrink_to_fit();
	m_pathToIndex.shrink_to_fit();
}

//////////////////////////////////////////////////////////////////////////

uint16_t ZipPack::getOrAddFolder(AbsPathView folderPath, tl::vector<tl::vector<EntryIndex>>& io_childrenIndices)
{
	if (folderPath.empty())
		return 0;

	uint64_t pathHash = computePersistentHash(folderPath);
	const auto it = m_pathToIndex.find(pathHash);
	if (it != m_pathToIndex.end())
	{
		const EntryIndex& entryIndex = it->second;
		if (!entryIndex.isFolder)
			TL_CRASH("Inconsistency");

		return entryIndex.index;
	}

	const tl::string name = std::move(folderPath.back());
	const AbsPathView parentPath = folderPath.parent();

	const uint16_t parentIndex = getOrAddFolder(parentPath, io_childrenIndices);

	const uint16_t index = tl::narrow<uint16_t>(m_folders.size());
	if (index >= 0xFFFF)
		TL_CRASH("Too many files/folders in zip pack");

	m_folders.push_back(Folder());
	Folder& folder = m_folders.back();
	folder.name = name;
	folder.parentIndex = parentIndex;

	if (parentIndex >= io_childrenIndices.size())
		io_childrenIndices.resize(parentIndex + 1);

	tl::vector<EntryIndex>& children = io_childrenIndices[parentIndex];
	children.push_back(EntryIndex(true, index));
	m_pathToIndex.emplace(pathHash, EntryIndex(true, index));

	return index;
}

//////////////////////////////////////////////////////////////////////////

void ZipPack::setEncryptionData(const tl::string& key, uint32_t rounds)
{
	m_encryptionKey = key;
	m_encryptionRounds = rounds;
}

//////////////////////////////////////////////////////////////////////////

OpenSourceResult ZipPack::openSource(AbsPathView path, IPack::SourceFlags flags) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	if (entryIndex.isFolder)
		return tl::make_error<Error>(ErrorCode::NotFound, "Trying to open a source on a folder");

	const File& file = m_files[entryIndex.index];

	return _openSource(file);
}

//////////////////////////////////////////////////////////////////////////

OpenStreamSourceResult ZipPack::openStreamSource(AbsPathView path, IPack::SourceFlags flags) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	if (entryIndex.isFolder)
		return tl::make_error<Error>(ErrorCode::NotFound, "Trying to open a source on a folder");

	return _openStreamSource(m_files[entryIndex.index]);
}

//////////////////////////////////////////////////////////////////////////

OpenMapSourceResult ZipPack::openMapSource(AbsPathView path, MapView mapView, SourceFlags flags) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	if (entryIndex.isFolder)
		return tl::make_error<Error>(ErrorCode::NotFound, "Trying to open a source on a folder");

	const File& file = m_files[entryIndex.index];

	return _openMapSource(file);
}

//////////////////////////////////////////////////////////////////////////

ConvertToNativePathResult ZipPack::convertToNativePath(AbsPathView path) const
{
	return tl::make_error<Error>(ErrorCode::NotAllowed, "Cannot convert zip paths to native ones '{}'", path);
}

//////////////////////////////////////////////////////////////////////////

IsFileResult ZipPack::isFile(AbsPathView path) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	return !entryIndex.isFolder;
}

//////////////////////////////////////////////////////////////////////////

IsFolderResult ZipPack::isFolder(AbsPathView path) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	return entryIndex.isFolder;
}

//////////////////////////////////////////////////////////////////////////

ExistsResult ZipPack::exists(AbsPathView path) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	return m_pathToIndex.find(computePersistentHash(path)) != m_pathToIndex.end();
}

//////////////////////////////////////////////////////////////////////////

GetStatResult ZipPack::getStat(AbsPathView path) const
{
	if (!path.is_valid())
		return tl::make_error<Error>(ErrorCode::BadPath, "Invalid Path");

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
		return tl::make_error<Error>(ErrorCode::NotFound, "Path '{}' doesn't exist", path);

	const EntryIndex& entryIndex = it->second;
	if (entryIndex.isFolder)
		return tl::make_error<Error>(ErrorCode::InvalidArgument, "Cannot stat a folder: '{}'", path);

	const File& file = m_files[entryIndex.index];
	zip_stat_t stat;
	if (fs_zip_stat_index(m_zipData->zipArchive.get(), file.zipIndex, 0, &stat) != 0)
		return tl::make_error<Error>(ErrorCode::SystemError, "Failed to stat: '{}'", path);

	return Stat{ stat.size, stat.mtime };
}

//////////////////////////////////////////////////////////////////////////

cppcoro::generator<EnumerateEntry> ZipPack::enumerate(AbsPath path) const
{
	TL_ASSERT(path.is_valid());

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
	{
		TL_FAIL("Path '{}' doesn't exist", path);
		co_return;
	}

	const EntryIndex& entryIndex = it->second;
	if (!entryIndex.isFolder)
	{
		TL_FAIL("Path '{}' is not a folder", path);
		co_return;
	}

	const Folder& folder = m_folders[entryIndex.index];
	for (size_t i = 0; i < folder.childrenCount; i++)
	{
		const EntryIndex& childEntryIndex = m_childrenIndices[folder.childrenStartIndex + i];
		if (childEntryIndex.isFolder)
		{
			const Folder& f = m_folders[childEntryIndex.index];
			co_yield{ RelPath(f.name), true };
		}
		else
		{
			const File& f = m_files[childEntryIndex.index];
			co_yield{ RelPath(f.name), false };
		}
	}
}

//////////////////////////////////////////////////////////////////////////

void ZipPack::buildUpPath(RelPath& path, const Folder& folder, uint32_t toIndex) const
{
	path = folder.name;
	{
		const Folder* f = &folder;
		while (f->parentIndex != toIndex && f->parentIndex != 0)
		{
			f = &m_folders[f->parentIndex];
			path.push_front(f->name);
		}
	}
}

void ZipPack::buildUpPath(RelPath& path, const File& file, uint32_t toIndex) const
{
	path = file.name;
	if (file.parentIndex != toIndex)
	{
		const Folder* f = &m_folders[file.parentIndex];
		while (true)
		{
			path.push_front(f->name);
			if (f->parentIndex == toIndex || f->parentIndex == 0)
				break;
			f = &m_folders[f->parentIndex];
		}
	}
}

cppcoro::generator<EnumerateEntry> ZipPack::enumerateRecursively(AbsPath path) const
{
	TL_ASSERT(path.is_valid());

	const auto it = m_pathToIndex.find(computePersistentHash(path));
	if (it == m_pathToIndex.end())
	{
		TL_FAIL("Path '{}' doesn't exist", path);
		co_return;
	}

	const EntryIndex& entryIndex = it->second;
	if (!entryIndex.isFolder)
	{
		TL_FAIL("Path '{}' is not a folder", path);
		co_return;
	}

	RelPath relPath;
	tl::fixed_vector<const Folder*, 64> stack;
	stack.reserve(64);
	stack.push_back(&m_folders[entryIndex.index]);
	for (size_t stackIndex = 0; stackIndex < stack.size(); stackIndex++)
	{
		const Folder* folder = stack[stackIndex];
		for (size_t i = 0; i < folder->childrenCount; i++)
		{
			const EntryIndex& childEntryIndex = m_childrenIndices[folder->childrenStartIndex + i];
			if (childEntryIndex.isFolder)
			{
				const Folder& f = m_folders[childEntryIndex.index];
				buildUpPath(relPath, f, entryIndex.index);
				co_yield{ relPath, true };
				stack.push_back(&f);
			}
			else
			{
				const File& f = m_files[childEntryIndex.index];
				buildUpPath(relPath, f, entryIndex.index);
				co_yield{ relPath, false };
			}
		}
	}
}

//////////////////////////////////////////////////////////////////////////

OpenMapSourceResult ZipPack::_openMapSource(const File& file) const
{
	zip_stat_t stat;
	if (fs_zip_stat_index(m_zipData->zipArchive.get(), file.zipIndex, 0, &stat) != 0)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': failed to stat", file.name);

	zip_file* zipFile = fs_zip_fopen_index(m_zipData->zipArchive.get(), file.zipIndex, 0);
	if (!zipFile)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': fopen failed", file.name);

	//we can use an adapter to convert a stream source into a regular source
	return tl::make_unique<StreamSourceToMapSourceAdapter>(tl::make_unique<ZipFileStreamSource>(m_zipData, tl::unique_ref<zip_file>(zipFile, my_zip_fclose), stat.size));
}

//////////////////////////////////////////////////////////////////////////

OpenSourceResult ZipPack::_openSource(const File& file) const
{
	zip_stat_t stat;
	if (fs_zip_stat_index(m_zipData->zipArchive.get(), file.zipIndex, 0, &stat) != 0)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': failed to stat", file.name);

	zip_file* zipFile = fs_zip_fopen_index(m_zipData->zipArchive.get(), file.zipIndex, 0);
	if (!zipFile)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': fopen failed", file.name);

	if (stat.comp_method == ZIP_CM_STORE) //we can use a ZipFileSource, since seeking is supported in this case
		return tl::make_unique<ZipFileSource>(m_zipData, tl::unique_ref<zip_file>(zipFile, my_zip_fclose), stat.size);

	//we can use an adapter to convert a stream source into a regular source
	return tl::make_unique<StreamSourceToMapSourceAdapter>(tl::make_unique<ZipFileStreamSource>(m_zipData, tl::unique_ref<zip_file>(zipFile, my_zip_fclose), stat.size));
}

//////////////////////////////////////////////////////////////////////////

OpenStreamSourceResult ZipPack::_openStreamSource(const File& file) const
{
	zip_stat_t stat;
	if (fs_zip_stat_index(m_zipData->zipArchive.get(), file.zipIndex, 0, &stat) != 0)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': failed to stat", file.name);

	zip_file* zipFile = fs_zip_fopen_index(m_zipData->zipArchive.get(), file.zipIndex, 0);
	if (!zipFile)
		return tl::make_error<Error>(ErrorCode::SystemError, "Cannot open file '{}': fopen failed", file.name);

	return tl::make_unique<ZipFileStreamSource>(m_zipData, tl::unique_ref<zip_file>(zipFile, my_zip_fclose), stat.size);
}
}
