#include "StdAfx.h"
#include "fs/DeflateStreamSource.h"

#include "tl/narrow_cast.h"

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

namespace fs
{

static constexpr size_t k_preferredBufferSize = 100 * 1024;

struct DeflateStreamSource::ZLib
{
	zng_stream streamer;
	bool needsData = true;
};

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

DeflateStreamSource::DeflateStreamSource(tl::unique_ref<IStreamSource> source)
	: m_ownedSource(std::move(source))
	, m_source(*m_ownedSource)
{

}

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

DeflateStreamSource::DeflateStreamSource(IStreamSource& source)
	: m_source(source)
{

}

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

DeflateStreamSource::~DeflateStreamSource()
{
	if (m_zlib)
		(void)zng_inflateEnd(&m_zlib->streamer);
}

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

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

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

void DeflateStreamSource::initializeDeflate()
{
	m_buffer.resize_uninitialized(getPreferredBufferSize());

	m_zlib = tl::make_unique<ZLib>();
	m_zlib->streamer.zalloc = nullptr;
	m_zlib->streamer.zfree = nullptr;
	m_zlib->streamer.opaque = nullptr;
	m_zlib->streamer.next_in = nullptr;
	m_zlib->streamer.avail_in = 0;
	const int ret = zng_inflateInit2(&m_zlib->streamer, -MAX_WBITS);
	if (ret != Z_OK)
	{
		(void)zng_inflateEnd(&m_zlib->streamer);
		m_zlib = nullptr;
	}
}

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

size_t DeflateStreamSource::read(tl::span<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 (!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;
	}

	size_t left = data.size();
	uint8_t* dst = data.data();

	do
	{
		if (m_zlib->needsData)
		{
			//this end of source is valid, and it means end of compressed source.
			if (m_source.isEOS()) 
				return data.size() - left;

			m_buffer.resize_uninitialized(getPreferredBufferSize());

			const uint64_t read = m_source.read(m_buffer);
			if (tl::optional<fs::Error> optLastError = m_source.getLastError(); optLastError.has_value())
			{
				m_buffer.clear();
				m_optLastError = tl::make_error<Error>(ErrorCode::SystemError, "Unexpected end of encrypted source");
				TL_FAIL("{}", m_optLastError);
				return 0;
			}
			m_buffer.resize_uninitialized(static_cast<size_t>(read));

			m_zlib->streamer.avail_in = tl::narrow<uInt>(m_buffer.size());
			m_zlib->streamer.next_in = m_buffer.data();
		}

		m_zlib->streamer.avail_out = tl::narrow<uInt>(left);
		m_zlib->streamer.next_out = reinterpret_cast<Bytef*>(dst);
		const int ret = zng_inflate(&m_zlib->streamer, Z_NO_FLUSH);
		TL_ASSERT(ret != Z_STREAM_ERROR);
		switch (ret)
		{
		case Z_OK: break; //all good
		case Z_STREAM_END: break; //all good
		case Z_NEED_DICT:
		{
			m_optLastError = tl::make_error<Error>(ErrorCode::SystemError, "deflate dictionary error ({})", ret);
			TL_FAIL("{}", m_optLastError);
			return 0;
		}
		case Z_DATA_ERROR:
		case Z_MEM_ERROR:
		default:
		{
			m_optLastError = tl::make_error<Error>(ErrorCode::SystemError, "deflate data error ({})", ret);
			TL_FAIL("{}", m_optLastError);
			return 0;
		}
		}

		const size_t read = left - m_zlib->streamer.avail_out;
		left -= read;
		dst += read;

		m_zlib->needsData = m_zlib->streamer.avail_out != 0;
	} while (left > 0);

	return data.size();
}

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

tl::unique_ref<IStreamSource> DeflateStreamSource::extract()
{
	return tl::promote(std::move(m_ownedSource));
}

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

bool DeflateStreamSource::isEOS() const
{
	return m_source.isEOS() && (!m_zlib || m_zlib->needsData);
}

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

size_t DeflateStreamSource::getPreferredBufferSize() const
{
	return m_source.getPreferredBufferSize() > 0 ? m_source.getPreferredBufferSize() : k_preferredBufferSize;
}

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