// TJP COPYRIGHT HEADER

#include "zlib.hpp"

#include <tjp/core/threads/Lock.hpp>
#include <tjp/core/containers/Map.hpp>
#include <tjp/core/algorithm/mem_copy.hpp>
#include <tjp/core/algorithm/ExecuteOnDestruct.hpp>

#include <zlib/zlib.h>

namespace tjp::core::compression::zlib {

enum class ContextType {
	Compress,
	Decompress
} ;

namespace {

constexpr int ZlibWindowBits = 15;
constexpr int GzipWindowBits = ZlibWindowBits + 16;
constexpr int InflateAutoWindowBits = ZlibWindowBits + 32;

int window_bits_for(Format format)
{
	switch (format)
	{
		case Format::Zlib:
			return ZlibWindowBits;
		case Format::Gzip:
			return GzipWindowBits;
		default:
			return ZlibWindowBits;
	}
}

size_t compress_key(int level, Format format)
{
	return (size_t(level) << 2) | size_t(format);
}

} // namespace

struct CompressContexts {
	Mutex mutex;
	Map<size_t, Vector<void *>> contextsByKey;
	
	~CompressContexts()
	{
		for (auto &[k, contexts] : contextsByKey)
			for (auto *c : contexts)
				free((z_streamp)c);
	}
	
	void *getContext(int level, Format format)
	{
		auto key = compress_key(level, format);
		{
			auto lock = lock_of(mutex);

			auto &contexts = contextsByKey[key];
			if (!contexts.empty())
			{
				auto *v = contexts.back();
				contexts.pop_back();
				return v;
			}
		}
		
		auto z = (z_streamp)calloc(1, sizeof(z_stream));
		auto result = deflateInit2(
			z,
			level,
			Z_DEFLATED,
			window_bits_for(format),
			8,
			Z_DEFAULT_STRATEGY
		);
		if (result != Z_OK)
		{
			free((void *)z);
			return nullptr;
		}
		return z;
	}
	
	void putContext(int level, Format format, void *v)
	{
		auto key = compress_key(level, format);
		auto lock = lock_of(mutex);
		contextsByKey[key].push_back(v);
	}
} ;

struct DecompressContexts
{
	Mutex mutex;
	Vector<void *> contexts;
	
	~DecompressContexts()
	{
		for (auto *c : contexts)
			free((z_streamp)c);
	}
	
	void *getContext()
	{
		{
			auto lock = lock_of(mutex);

			if (!contexts.empty())
			{
				auto *v = contexts.back();
				contexts.pop_back();
				return v;
			}
		}
		
		auto z = (z_streamp)calloc(1, sizeof(z_stream));
		auto result = inflateInit2(z, InflateAutoWindowBits);
		if (result != Z_OK)
		{
			free((void *)z);
			return nullptr;
		}
		return z;
	}
	
	void putContext(void *v)
	{
		auto lock = lock_of(mutex);
		contexts.push_back(v);
	}
} ;

CompressContexts compressContexts;
DecompressContexts decompressContexts;

#ifdef TIMPREPSCIUS_ZLIB_USE_DICTIONARY

Dictionary::Dictionary()
{
	memset(block, 0, size);
}

bool setDict(void *p, ContextType type, const Dictionary &d)
{
	auto z = (z_streamp)p;
	
	int error = 0;
	if (type == ContextType::Decompress)
		error = inflateSetDictionary(z, (const Bytef *)d.block, (uInt)d.size);
	else
		error = deflateSetDictionary(z, (const Bytef *)d.block, (uInt)d.size);

	return error == Z_OK;
}

bool getDict(void *p, ContextType type, Dictionary &d)
{
	auto z = (z_streamp)p;
	
	constexpr auto DictStoreSize = 32768;
	char store[DictStoreSize];
	memset(store, 0, DictStoreSize);
	auto size_ = (uInt)DictStoreSize;

	int error = 0;
	if (type == ContextType::Decompress)
		error = inflateGetDictionary(z, (Bytef *)&store[0], &size_);
	else
		error = deflateGetDictionary(z, (Bytef *)&store[0], &size_);

	// zlib seems to store from end to front
	auto offset =
		size_ > d.size ?
			size_ - d.size :
			0;
		
	int firstIndex = -1, lastIndex = -1;
	for (auto i=0; i<DictStoreSize; ++i)
	{
		if (store[i] != 0)
		{
			if (firstIndex < 0)
				firstIndex = i;
			lastIndex = i;
		}
	}

	mem_copy(d.block, &store[offset], d.size);
	
	return error == Z_OK;
}

#else

constexpr
bool setDict(void *p, ContextType type, const Dictionary &d) { return true; }

constexpr
bool getDict(void *p, ContextType type, Dictionary &d) { return true; }

#endif

// --------------

bool clearStream(void *p, ContextType type)
{
	auto z = (z_streamp)p;

	if (type == ContextType::Decompress)
		return (inflateReset(z) == Z_OK);
	else
		return (deflateReset(z) == Z_OK);
}


// --------------

static size_t round_up_to_multiple(size_t number, size_t multiple) {
    if (multiple == 0) {
        return number;  // Avoid division by zero
    }
    auto remainder = number % multiple;
    if (remainder == 0) {
        return number;  // Already a multiple
    }
    return number + multiple - remainder;
}

Compressor::Compressor(int level_, Format format_) :
	level(level_),
	format(format_)
{
}

size_t Compressor::with(V &v, F &&f)
{
	error_ = 0;

	auto context = (z_streamp)compressContexts.getContext(level, format);
	if (!context)
		return 0;

	auto _ = ExecuteOnDestruct([&]() {
		compressContexts.putContext(level, format, context);
		i = nullptr;
		v_ = nullptr;
	});

	if (!clearStream(context, ContextType::Compress))
		return 0;
		
	if (!setDict(context, ContextType::Compress, dictionary))
		return 0;

	i = (I *)context;
	v_ = &v;
	p = v.size();
	
	auto p_ = p;
	
	f(*this);
	
	// finish the stream
	while(1)
	{
		auto newSize = round_up_to_multiple(p + 1, sizeMultiples);
		v.resize(newSize);

		context->next_out = (Bytef *)v.data() + p;
		context->avail_out = uInt(newSize - p);

		auto result = deflate(context, Z_FINISH);
		p = size_t(context->next_out - (Bytef *)v.data());

		if (result == Z_STREAM_END)
			break;

		if (result <= Z_ERRNO)
		{
			error_ = result;
			break;
		}
	}
	
	getDict(context, ContextType::Compress, dictionary);
	
	p = size_t(context->next_out - (Bytef *)v.data());
	v.resize(p);
	
	return p - p_;
}

size_t Compressor::write(const containers::MemorySegment<char> &s)
{
	auto &v = *v_;
	
	auto context = (z_streamp)i;
	
	context->next_in = (Bytef *)s.data();
	context->avail_in = (uInt)s.size();

	// beginning of stream
	int result = 0;
	size_t written = 0;

	while(1)
	{
		auto newSize = round_up_to_multiple(p + 1, sizeMultiples);
		v.resize(newSize);
		
		context->next_out = (Bytef *)v.data() + p;
		context->avail_out = uInt(newSize - p);

		result = deflate(context, Z_NO_FLUSH);
		p = size_t(context->next_out - (Bytef *)v.data());
		written = s.size() - (size_t)context->avail_in;

		if (result == Z_STREAM_END)
			break;

		if (result <= Z_ERRNO)
		{
			error_ = result;
			break;
		}
			
		if (context->avail_in == 0)
			break;
	}

	
	return written;
}

void Compressor::flush()
{
	auto &v = *v_;
	
	auto context = (z_streamp)i;
	
	context->next_in = nullptr;
	context->avail_in = 0;

	// beginning of stream
	int result = 0;

	while(1)
	{
		auto newSize = round_up_to_multiple(p + 1, sizeMultiples);
		v.resize(newSize);
		
		context->next_out = (Bytef *)v.data() + p;
		context->avail_out = uInt(newSize - p);

		result = deflate(context, Z_PARTIAL_FLUSH);
		p = size_t(context->next_out - (Bytef *)v.data());

		if (result == Z_STREAM_END)
			break;

		if (result <= Z_ERRNO)
		{
			error_ = result;
			break;
		}
			
		if (context->avail_out > 0)
			break;
	}
}

size_t Compressor::size() const
{
	return p;
}

bool Compressor::error () const
{
	return error_;
}

CompressStream::CompressStream(int level, Format format) :
	compressor(level, format)
{
}

size_t CompressStream::with(V &v, F &&f)
{
	return compressor.with(v, std::move(f));
}

CompressPacket::CompressPacket(int level_, Format format_) :
	level(level_),
	format(format_)
{

}

size_t CompressPacket::with(V &v, F &&f)
{
	Compressor compressor(level, format);
	return compressor.with(v, std::move(f));
}

Decompressor::Decompressor()
{
}

size_t Decompressor::with(V &v, F &&f)
{
	error_ = 0;
	
	auto context = (z_streamp)decompressContexts.getContext();
	if (!context)
		return 0;
	
	i = (I *)context;
	v_ = &v;
	
	auto _ = ExecuteOnDestruct([&]() {
		decompressContexts.putContext(context);
		i = nullptr;
		v_ = nullptr;
	});
	
	if (!clearStream(context, ContextType::Decompress))
		return 0;

	context->next_in = (Bytef *)v.data();
	context->avail_in = uInt(v.size());
	p = 0;
	
	f(*this);
	
	getDict(context, ContextType::Decompress, dictionary);
	
	auto written = context->total_out;
	auto consumed = size_t(context->next_in - (Bytef *)v.data());
	
	v.data_ = (char *)context->next_in;
	if (consumed > v.size_)
		v.size_ = 0;
	else
		v.size_ -= consumed;
	return written;
}

size_t Decompressor::read(const containers::MemorySegment<char> &s)
{
	auto &v = *v_;
	auto context = (z_streamp)i;
	
	context->next_out = (Bytef *)s.data();
	context->avail_out = (uInt)s.size();

	int result = 0;
	size_t written = 0;

	while(1)
	{
		result = inflate(context, Z_SYNC_FLUSH);
		written = size_t(context->next_out - (Bytef *)s.data());
		p = size_t(context->next_in - (Bytef *)v.data());
		
		if (result == Z_NEED_DICT)
		{
			if (!setDict(context, ContextType::Decompress, dictionary))
				return 0;
		}

		if (result == Z_STREAM_END)
			break;

		if (result <= Z_ERRNO)
		{
			error_ = result;
			break;
		}
			
		if (written == s.size())
			break;
	}

	return written;
}

bool Decompressor::eos() const
{
	return error() || p >= v_->size();
}

bool Decompressor::error () const
{
	return error_;
}

DecompressStream::DecompressStream()
{

}

size_t DecompressStream::with(V &v, F &&f)
{
	return decompressor.with(v, std::move(f));
}

DecompressPacket::DecompressPacket()
{

}

size_t DecompressPacket::with(V &v, F &&f)
{
	Decompressor decompressor;
	
	return decompressor.with(v, std::move(f));
}



} // namespace