#include <csignal>
#include <misc/memtracker.h>
#include "CompressionHandler.h"
#define QLZ_COMPRESSION_LEVEL 1
#define QLZ_MEMORY_SAFE
#include "qlz/QuickLZ.h"
#include "buffers.h"

using namespace ts;
using namespace ts::connection;
using namespace std;

bool CompressionHandler::compress(protocol::BasicPacket* packet, std::string &error) {
	//// "Always allocate size + 400 bytes for the destination buffer when compressing." <= http://www.quicklz.com/manual.html
	auto length = packet->data().length();
	auto header_length = packet->header().length() + packet->mac().length();
	size_t expected_bytes = min(length * 2, length + 400);
	auto buffer = buffer::allocate_buffer(expected_bytes + header_length);

	qlz_state_compress state_compress{};
	size_t actualLength = qlz_compress(packet->data().data_ptr(), (char*) &buffer[header_length], packet->data().length(), &state_compress);
	if(actualLength > buffer.length()) {
		logCritical(0, "Buffer overflow! Compressed data is longer than expected. (Expected: {}, Written: {}, Source length: {}, Allocated block size: {})",
				expected_bytes,
				actualLength,
				packet->data().length(),
				buffer.capacity()
		);
		error = "overflow";
		return false;
	}
	if(actualLength <= 0){
		error = "Cloud not compress packet";
		return false;
	}

	memcpy(buffer.data_ptr(), packet->buffer().data_ptr(), header_length);
	packet->buffer(buffer.range(0, actualLength + header_length));
	return true;
}

bool CompressionHandler::decompress(protocol::BasicPacket* packet, std::string &error) {
	qlz_state_decompress state_decompress{};

	size_t expected_length = qlz_size_decompressed((char*) packet->data().data_ptr());
	if(expected_length > this->max_packet_size){ //Max 16MB. (97% Compression!)
		error = "Invalid packet size. (Calculated target length of " + to_string(expected_length) + ". Max length: " + to_string(this->max_packet_size) + ")";
		return false;
	}

	auto header_length = packet->header().length() + packet->mac().length();
	auto buffer = buffer::allocate_buffer(expected_length + header_length);
	size_t data_length = qlz_decompress((char*) packet->data().data_ptr(), &buffer[header_length], &state_decompress);
	if(data_length <= 0){
		error = "Could not decompress packet.";
		return false;
	}

	memcpy(buffer.data_ptr(), packet->buffer().data_ptr(), header_length);
	packet->buffer(buffer.range(0, data_length + header_length));
	return true;
}

bool CompressionHandler::progressPacketIn(protocol::BasicPacket* packet, std::string &error) {
	if(packet->isCompressed()) {
		if(!this->decompress(packet, error)) return false;
		packet->setCompressed(false);
	}
	return true;
}

bool CompressionHandler::progressPacketOut(protocol::BasicPacket* packet, std::string& error) {
	if(packet->hasFlag(protocol::PacketFlag::Compressed) && !packet->isCompressed()) {
		if(!this->compress(packet, error)) return false;
		packet->setCompressed(true);
	}
	return true;
}

CompressionHandler::CompressionHandler() { }
CompressionHandler::~CompressionHandler() { }