#include "StdAfx.h"

#include "fs/zip/ZipWriter.h"

#include "fs/CRCStream.h"
#include "fs/IStreamSource.h"
#include "fs/ISink.h"
#include "tl/narrow_cast.h"

// This writer is implemented pretty close to the https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT standard

namespace fs
{
//////////////////////////////////////////////////////////////////////////

ZipWriter::ZipWriter(ISink& sink, size_t fileDataAlignment)
	: m_sink(sink)
	, m_fileDataAlignment(fileDataAlignment == 0 ? 1 : fileDataAlignment)
{}

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

tl::result<void> ZipWriter::addFile(tl::string name, const DataWriter& dataWriter)
{
	return addFile(std::move(name),
	               //
	               {},
	               {},
	               GeneralBitFlags(),
	               tl::nullopt,
	               dataWriter);
}

tl::result<void> ZipWriter::addFile(tl::string name, tl::optional<time_t> lastModificationTime, const DataWriter& dataWriter)
{
	return addFile(std::move(name),
	               //
	               {},
	               {},
	               GeneralBitFlags(),
	               std::move(lastModificationTime),
	               dataWriter);
}

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

tl::result<void> ZipWriter::addFile(tl::string name,
                                    tl::span<const uint8_t> comment,
                                    tl::span<const uint8_t> extraField,
                                    GeneralBitFlags generalBitFlags,
                                    tl::optional<time_t> lastModificationTime,
                                    const DataWriter& dataWriter)
{
	if (name.size() > tl::numeric_limits<uint16_t>::max())
		return tl::make_generic_error("Name too long");

	if (comment.size() > tl::numeric_limits<uint16_t>::max())
		return tl::make_generic_error("Comment too long");

	if (extraField.size() + sizeof(ExtendedInformationExtraField64) > tl::numeric_limits<uint16_t>::max())
		return tl::make_generic_error("Extra field too long");

	//////////////////////////////////////////////////////////////////////////
	// Cannot compute the whole header as we don't know how big the data will be.
	// So we compute as much as possible from the header, enough to figure out the padding.
	// The we skip to the data section and write it - computing its size at the same time.
	// Then we go back to the header section, fill in all the info (since we now know the size, and implictly if it's zip32/64 format) and write the header, name
	// etc) NOTE - because of this we always have to write the ExtendedInformationExtraField64 section, even if the final format will not be zip64

	LocalFileHeader header;
	header.extraFieldLength = tl::narrow<uint16_t>(extraField.size()) + sizeof(ExtendedInformationExtraField64);
	header.generalBitFlag = generalBitFlags.value() | tl::narrow<uint16_t>(GeneralBitFlag::LanguageEncodingFlag);

	header.lastModFileDate = 0;
	header.lastModFileTime = 0;
	if (lastModificationTime != tl::nullopt)
		computeMsDosTime(*lastModificationTime, header.lastModFileDate, header.lastModFileTime);

	header.filenameLength = tl::narrow<uint16_t>(name.size());

	const uint64_t padding = computePadding(m_sink.tell(), header);
	m_sink.seekRel(padding);

	const uint64_t localHeaderOffset = m_sink.tell();

	// now skip to the data
	m_sink.seekRel(sizeof(LocalFileHeader) + header.filenameLength + header.extraFieldLength);

	const uint64_t startOffset = m_sink.tell();

	// write the data - this computes the size as well
	DataWriterResult dataProviderRes = dataWriter(m_sink);
	if (!dataProviderRes)
		return dataProviderRes.error();

	const DataWriterPayload payload = dataProviderRes.value();

	header.compressionMethod = tl::narrow<uint16_t>(payload.compressionMethod);
	header.crc32 = payload.uncompressedCRC32;

	const uint64_t endOffset = m_sink.tell();
	const uint64_t compressedSize = endOffset - startOffset;
	const uint64_t uncompressedSize = payload.uncompressedSize ? payload.uncompressedSize : compressedSize;

	// write the size fields (in zip64, these are -1)
	header.compressedSize = static_cast<uint32_t>(-1);
	header.uncompressedSize = static_cast<uint32_t>(-1);

	// now go back and write the header
	m_sink.seekBeg(localHeaderOffset);
	m_sink.write({ reinterpret_cast<const uint8_t*>(&header), sizeof(LocalFileHeader) });

	// the filename
	if (!name.empty())
		m_sink.write({ reinterpret_cast<const uint8_t*>(name.data()), name.size() });

	// extended field - always present!
	{
		ExtendedInformationExtraField64 extended;
		extended.size = static_cast<uint16_t>(sizeof(ExtendedInformationExtraField64) - 4); // Substract the tag + size fields.
		extended.compressedSize = compressedSize;
		extended.originalSize = uncompressedSize;
		extended.localHeaderOffset = localHeaderOffset;
		m_sink.write({ reinterpret_cast<const uint8_t*>(&extended), sizeof(ExtendedInformationExtraField64) });
	}

	// custom fields
	if (!extraField.empty())
		m_sink.write(extraField);

	m_sink.seekBeg(endOffset);

	// store the central directory record
	CentralDirectoryRecord record;
	record.header = computeCentralDirectoryFileHeader(header, true);
	record.compressedSize = compressedSize;
	record.uncompressedSize = uncompressedSize;
	record.localHeaderOffset = localHeaderOffset;
	record.header.localHeaderOffset = static_cast<uint32_t>(-1); //zip64 mandates -1
	record.zip64 = true;
	record.header.fileCommentLength = tl::narrow<uint16_t>(comment.size());
	record.name = std::move(name);
	record.comment.assign(comment.begin(), comment.end());
	record.extraField.assign(extraField.begin(), extraField.end());
	m_records.push_back(std::move(record));

	return tl::success();
}

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

void ZipWriter::finish()
{
	const uint64_t centralDirectoryRecordsOffset = m_sink.tell();

	for (const CentralDirectoryRecord& record : m_records)
	{
		TL_ASSERT(record.header.filenameLength == record.name.size());
		TL_ASSERT(record.header.extraFieldLength == record.extraField.size() + (record.zip64 ? sizeof(ExtendedInformationExtraField64) : 0));
		TL_ASSERT(record.header.fileCommentLength == record.comment.size());

		m_sink.write({ reinterpret_cast<const uint8_t*>(&record.header), sizeof(CentralDirectoryFileHeader) });
		if (!record.name.empty())
			m_sink.write({ reinterpret_cast<const uint8_t*>(record.name.data()), record.name.size() });

		if (!record.extraField.empty())
			m_sink.write(record.extraField);

		if (record.zip64)
		{
			ExtendedInformationExtraField64 extended;
			extended.size = tl::narrow<uint16_t>(sizeof(ExtendedInformationExtraField64) - 4); // Subtract the tag + size fields.
			extended.originalSize = record.uncompressedSize;
			extended.compressedSize = record.compressedSize;
			extended.localHeaderOffset = record.localHeaderOffset;
			m_sink.write({ reinterpret_cast<const uint8_t*>(&extended), sizeof(ExtendedInformationExtraField64) });
		}
		if (!record.comment.empty())
			m_sink.write(record.comment);
	}
	const uint64_t sizeOfCentralDirectoryRecords = m_sink.tell() - centralDirectoryRecordsOffset;

	const bool zip64 = centralDirectoryRecordsOffset >= tl::numeric_limits<uint32_t>::max() || //
		sizeOfCentralDirectoryRecords >= tl::numeric_limits<uint32_t>::max() || //
		m_records.size() >= tl::numeric_limits<uint16_t>::max();

	if (zip64)
	{
		const uint64_t endOfCentralDirectoryRecord64Offset = m_sink.tell();

		{
			EndOfCentralDirectoryRecord64 header;
			header.endOfCentralDirectorySize = sizeof(EndOfCentralDirectoryRecord64) - 12;
			header.centralDirectoryRecordCountOnThisDisk = m_records.size();
			header.centralDirectoryRecordCount = m_records.size();
			header.centralDirectorySize = sizeOfCentralDirectoryRecords;
			header.centralDirectoryOffset = centralDirectoryRecordsOffset;
			m_sink.write({ reinterpret_cast<const uint8_t*>(&header), sizeof(EndOfCentralDirectoryRecord64) });
		}
		{
			EndOfCentralDirectoryLocator64 header;
			header.endOfCentralDirectoryRecordOffset = endOfCentralDirectoryRecord64Offset;
			header.diskCount = 1;
			m_sink.write({ reinterpret_cast<const uint8_t*>(&header), sizeof(EndOfCentralDirectoryLocator64) });
		}
	}

	EndOfCentralDirectoryRecord header;
	header.centralDirectoryRecordCountOnThisDisk = zip64 ? static_cast<uint16_t>(-1) : tl::narrow<uint16_t>(m_records.size());
	header.centralDirectoryRecordCount = zip64 ? static_cast<uint16_t>(-1) : tl::narrow<uint16_t>(m_records.size());
	header.centralDirectorySize = zip64 ? static_cast<uint32_t>(-1) : tl::narrow<uint32_t>(sizeOfCentralDirectoryRecords);
	header.centralDirectoryOffset = zip64 ? static_cast<uint32_t>(-1) : tl::narrow<uint32_t>(centralDirectoryRecordsOffset);
	m_sink.write({ reinterpret_cast<const uint8_t*>(&header), sizeof(EndOfCentralDirectoryRecord) });
}

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

inline uint64_t ZipWriter::computePadding(uint64_t offset, const LocalFileHeader& header) const
{
	const uint64_t dataOffset = offset + sizeof(LocalFileHeader) + header.filenameLength + header.extraFieldLength;
	const uint64_t m = dataOffset % m_fileDataAlignment;
	if (m == 0)
		return 0;

	return m_fileDataAlignment - m;
}

//////////////////////////////////////////////////////////////////////////
}