#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 packet_payload = packet->data();
	auto header_length = packet->length() - packet_payload.length();

	size_t max_compressed_payload_size = max(min(packet_payload.length() * 2, (size_t) (packet_payload.length() + 400ULL)), (size_t) 24ULL); /* at least 12 bytes (QLZ header) */
	auto target_buffer = buffer::allocate_buffer(max_compressed_payload_size + header_length);

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

	memcpy(target_buffer.data_ptr(), packet->buffer().data_ptr(), header_length);
	packet->buffer(target_buffer.range(0, actual_length + 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->has_flag(protocol::PacketFlag::Compressed) && !packet->isCompressed()) {
		if(!this->compress(packet, error)) return false;
		packet->setCompressed(true);
	}
	return true;
}

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