#include "StdAfx.h"
#include "fs/DeflateStreamSink.h"
#include "tl/narrow_cast.h"

extern "C"
{
#include "zlib-ng.h"
}

namespace fs
{

static constexpr size_t k_preferredBufferSize = 100 * 1024;
static constexpr size_t k_outBufferSize = 100 * 1024;

struct DeflateStreamSink::ZLib
{
	zng_stream streamer;
};

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

DeflateStreamSink::DeflateStreamSink(tl::unique_ref<IStreamSink> sink)
	: m_ownedSink(std::move(sink))
	, m_sink(*m_ownedSink)
{

}

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

DeflateStreamSink::DeflateStreamSink(IStreamSink& sink)
	: m_sink(sink)
{

}

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

DeflateStreamSink::~DeflateStreamSink()
{
	if (!m_isFinished)
		flush(true);

	if (m_zlib)
		(void)zng_deflateEnd(&m_zlib->streamer);
}

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

void DeflateStreamSink::setCompressionLevel(uint8_t level)
{
	m_level = tl::min<uint8_t>(level, 9u);
}

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

tl::optional<Error> DeflateStreamSink::getLastError() const
{
	tl::optional<Error> error = m_optLastError;
	m_optLastError = tl::nullopt;
	return error;
}

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

void DeflateStreamSink::initializeDeflate()
{
	m_buffer.resize_uninitialized(getPreferredBufferSize());
	m_bufferOffset = 0;
	m_outBuffer.resize_uninitialized(m_sink.getPreferredBufferSize() > 0 ? m_sink.getPreferredBufferSize() : k_outBufferSize);

	m_zlib = tl::make_unique<ZLib>();
	m_zlib->streamer.zalloc = nullptr;
	m_zlib->streamer.zfree = nullptr;
	m_zlib->streamer.opaque = nullptr;
	const int ret = zng_deflateInit2(&m_zlib->streamer, m_level, Z_DEFLATED, -MAX_WBITS, 9, Z_DEFAULT_STRATEGY);
	if (ret != Z_OK)
	{
		(void)zng_deflateEnd(&m_zlib->streamer);
		m_zlib = nullptr;
	}
}

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

size_t DeflateStreamSink::write(tl::span<const uint8_t> data)
{
	if (m_optLastError.has_value())
	{
		TL_FAIL(); //the previous file access generated an error
		return 0;
	}

	if (data.empty())
		return 0;

	if (!data.data())
	{
		m_optLastError = tl::make_error<Error>(ErrorCode::InvalidArgument, "Null buffer.");
		TL_FAIL("{}", m_optLastError);
		return 0;
	}

	size_t left = data.size();
	const uint8_t* src = data.data();

	if (!m_zlib)
		initializeDeflate();

	if (!m_zlib)
	{
		m_optLastError = tl::make_error<Error>(ErrorCode::SystemError, "Failed to initialize the compressor");
		TL_FAIL("{}", m_optLastError);
		return 0;
	}

	while (left > 0)
	{
		if (const size_t s = tl::min(left, m_buffer.size() - m_bufferOffset))
		{
			memcpy(m_buffer.data() + m_bufferOffset, src, s);
			src += s;
			left -= s;
			m_bufferOffset += s;
			m_size += s;
		}

		if (!flush(false))
			return 0;
	}

	return data.size();
}

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

bool DeflateStreamSink::flush(bool finalize)
{
	if (!m_zlib)
		return true;

	if (finalize || (m_bufferOffset > 0 && m_bufferOffset >= m_buffer.size()))
	{
		m_zlib->streamer.avail_in = tl::narrow<uInt>(m_bufferOffset);
		m_zlib->streamer.next_in = m_buffer.data();

		do
		{
			m_zlib->streamer.avail_out = tl::narrow<uInt>(m_outBuffer.size());
			m_zlib->streamer.next_out = m_outBuffer.data();
			const int ret = zng_deflate(&m_zlib->streamer, finalize ? Z_FINISH : Z_NO_FLUSH);
			TL_ASSERT(ret != Z_STREAM_ERROR);
			(void)ret;
			const size_t available = m_outBuffer.size() - m_zlib->streamer.avail_out;
			const uint64_t written = m_sink.write({ m_outBuffer.data(), available });
			tl::optional<Error> optLastError = m_sink.getLastError();
			if (written != available || optLastError.has_value())
			{
				m_buffer.clear();
				m_bufferOffset = 0;
				m_optLastError = tl::make_error<Error>(ErrorCode::SystemError, "Subsink failed while writting data");
				TL_FAIL("{}", m_optLastError);
				return false;
			}
		} while (m_zlib->streamer.avail_out == 0);

		m_bufferOffset = 0;
	}
	return true;
}

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

void DeflateStreamSink::finish()
{
	TL_ASSERT(!m_isFinished);
	m_isFinished = true;
	flush(true);
}

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

tl::unique_ref<IStreamSink> DeflateStreamSink::finishAndExtract()
{
	finish();
	return tl::promote(std::move(m_ownedSink));
}

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

size_t DeflateStreamSink::getPreferredBufferSize() const
{
	return m_sink.getPreferredBufferSize() > 0 ? m_sink.getPreferredBufferSize() : k_preferredBufferSize;
}

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

uint64_t DeflateStreamSink::getSize() const
{
	return m_size;
}

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