#include "td/misc/Compression.h"

#include <zlib.h>
#include <cassert>

#define COMPRESSION_THRESHOLD 64

namespace td {
namespace utils {

std::uint64_t Inflate(const std::string& source, std::string& dest) {
	uLongf size = dest.size();
	uncompress(reinterpret_cast<Bytef*>(dest.data()), reinterpret_cast<uLongf*>(&size), reinterpret_cast<const Bytef*>(source.c_str()), source.length());
	return size;
}

std::uint64_t Deflate(const std::string& source, std::string& dest) {
	uLongf size = source.length();
	dest.resize(source.size());  // Resize for the compressed data to fit into
	compress(reinterpret_cast<Bytef*>(dest.data()), reinterpret_cast<uLongf*>(&size), reinterpret_cast<const Bytef*>(source.c_str()), source.length());
	dest.resize(size);  // Resize to cut useless data
	return size;
}

DataBuffer Compress(const DataBuffer& buffer) {
	std::string compressedData;
	DataBuffer packet;

	if (buffer.GetSize() < COMPRESSION_THRESHOLD) {
		// Don't compress since it's a small packet
		std::uint64_t dataLength = 0;
		std::uint64_t packetLength = buffer.GetSize() + sizeof(dataLength);

		packet << packetLength;
		packet << dataLength;
		packet << buffer;
		return packet;
	}

	Deflate(buffer.ToString(), compressedData);

	std::uint64_t dataLength = buffer.GetSize();
	std::uint64_t packetLength = compressedData.length() + sizeof(dataLength);

	packet << packetLength;
	packet << dataLength;
	packet << compressedData;
	return packet;
}

DataBuffer Decompress(DataBuffer& buffer, std::uint64_t packetLength) {
	std::uint64_t uncompressedLength;

	buffer >> uncompressedLength;

	std::uint64_t compressedLength = packetLength - sizeof(uncompressedLength);

	if (uncompressedLength == 0) {
		// Uncompressed
		DataBuffer ret;
		buffer.ReadSome(ret, compressedLength);
		return ret;
	}

	assert(buffer.GetReadOffset() + compressedLength <= buffer.GetSize());

	std::string deflatedData;
	buffer.ReadSome(deflatedData, compressedLength);

	std::string inflated;
	inflated.resize(uncompressedLength);

	Inflate(deflatedData, inflated);

	assert(inflated.length() == uncompressedLength);
	return DataBuffer(inflated);
}

DataBuffer Decompress(DataBuffer& buffer) {
	std::uint64_t packetLength;

	buffer >> packetLength;

	return Decompress(buffer, packetLength);
}

} // namespace utils
} // namespace td
