//
// Created by WolverinDEV on 10/03/2020.
//

#include "PacketDecoder.h"

#include <protocol/buffers.h>
#include <protocol/AcknowledgeManager.h>
#include <protocol/CompressionHandler.h>
#include <protocol/CryptHandler.cpp>

#include "../../ConnectionStatistics.h"

using namespace ts;
using namespace ts::server::server::udp;

PacketDecoder::PacketDecoder(ts::connection::CryptHandler *crypt_handler)
        : crypt_handler_{crypt_handler} {
    memtrack::allocated<PacketDecoder>(this);
}

PacketDecoder::~PacketDecoder() {
    memtrack::freed<PacketDecoder>(this);
    this->reset();
}

void PacketDecoder::reset() {
    std::lock_guard buffer_lock(this->packet_buffer_lock);
    for(auto& buffer : this->_command_fragment_buffers)
        buffer.reset();
}

bool PacketDecoder::decode_incoming_data(const pipes::buffer_view &buffer) {
    std::string error{};
    bool needs_command_reassemble{false};

    auto result = this->decode_incoming_data_(error, needs_command_reassemble, buffer);
    if(result != PacketDecodeResult::SUCCESS)
        if(auto callback{this->callback_decode_failed}; callback)
            callback(this->callback_argument, result, error);

    return needs_command_reassemble;
}

PacketDecodeResult PacketDecoder::decode_incoming_data_(std::string& error, bool& commands_pending, const pipes::buffer_view &buffer) {
#ifdef FUZZING_TESTING_INCOMMING
    #ifdef FIZZING_TESTING_DISABLE_HANDSHAKE
    if (this->client->state == ConnectionState::CONNECTED) {
    #endif
        if ((rand() % FUZZING_TESTING_DROP_MAX) < FUZZING_TESTING_DROP) {
            debugMessage(this->client->getServerId(), "{}[FUZZING] Dropping incoming packet of length {}", CLIENT_STR_LOG_PREFIX_(this->client), buffer.length());
            return;
        }
    #ifdef FIZZING_TESTING_DISABLE_HANDSHAKE
    }
    #endif
#endif

    ClientPacketParser packet_parser{buffer};
    if(!packet_parser.valid())
        return PacketDecodeResult::INVALID_PACKET;

    assert(packet_parser.type() >= 0 && packet_parser.type() < this->incoming_generation_estimators.size());
    packet_parser.set_estimated_generation(this->incoming_generation_estimators[packet_parser.type()].visit_packet(packet_parser.packet_id()));

    auto is_command = packet_parser.type() == protocol::COMMAND || packet_parser.type() == protocol::COMMAND_LOW;
    /* pretest if the packet is worth the effort of decoding it */
    if(is_command) {
        /* handle the order stuff */
        auto& fragment_buffer = this->_command_fragment_buffers[PacketDecoder::command_fragment_buffer_index(packet_parser.type())];

        unique_lock queue_lock(fragment_buffer.buffer_lock);
        auto result = fragment_buffer.accept_index(packet_parser.packet_id());
        if(result != 0) { /* packet index is ahead buffer index */
            error = "pid: " + std::to_string(packet_parser.packet_id()) + ",";
            error += "bidx: " + std::to_string(fragment_buffer.current_index()) + ",";
            error += "bcap: " + std::to_string(fragment_buffer.capacity());

            if(result == -1) { /* underflow */
                /* we've already got the packet, but the client dosn't know that so we've to send the acknowledge again */
                this->callback_send_acknowledge(packet_parser.packet_id(), packet_parser.type() == protocol::COMMAND_LOW);
                return PacketDecodeResult::DUPLICATED_PACKET;
            }

            return PacketDecodeResult::BUFFER_OVERFLOW;
        }
    }

    //NOTICE I found out that the Compressed flag is set if the packet contains an audio header
    /* decrypt the packet if needed */
    if(packet_parser.is_encrypted()) {
        CryptHandler::key_t crypt_key{};
        CryptHandler::nonce_t crypt_nonce{};

        auto data = (uint8_t*) packet_parser.mutable_data_ptr();
        bool use_default_key{!this->protocol_encrypted}, decrypt_result;

        decrypt_packet:
        if(use_default_key) {
            crypt_key = CryptHandler::default_key;
            crypt_nonce = CryptHandler::default_nonce;
        } else {
            if(!this->crypt_handler_->generate_key_nonce(true, packet_parser.type(), packet_parser.packet_id(), packet_parser.estimated_generation(), crypt_key, crypt_nonce))
                return PacketDecodeResult::DECRYPT_KEY_GEN_FAILED;
        }

        decrypt_result = this->crypt_handler_->decrypt(
                data + ClientPacketParser::kHeaderOffset, ClientPacketParser::kHeaderLength,
                data + ClientPacketParser::kPayloadOffset, packet_parser.payload_length(),
                data,
                crypt_key, crypt_nonce,
                error
        );

        if(!decrypt_result) {
            if(packet_parser.packet_id() < 10 && packet_parser.estimated_generation() == 0) {
                if(use_default_key) {
                    return PacketDecodeResult::DECRYPT_FAILED;
                } else {
                    use_default_key = true;
                    goto decrypt_packet;
                }
            } else {
                return PacketDecodeResult::DECRYPT_FAILED;
            }
        }
        packet_parser.set_decrypted();
    }

    if(auto statistics{this->statistics_}; statistics)
        statistics->logIncomingPacket(stats::ConnectionStatistics::category::from_type(packet_parser.type()), buffer.length());

#ifdef LOG_INCOMPING_PACKET_FRAGMENTS
    debugMessage(lstream << CLIENT_LOG_PREFIX << "Recived packet. PacketId: " << packet->packetId() << " PacketType: " << packet->type().name() << " Flags: " << packet->flags() << " - " << packet->data() << endl);
#endif
    if(is_command) {
        auto& fragment_buffer = this->_command_fragment_buffers[command_fragment_buffer_index(packet_parser.type())];
        CommandFragment fragment_entry{
                packet_parser.packet_id(),
                packet_parser.estimated_generation(),

                packet_parser.flags(),
                (uint32_t) packet_parser.payload_length(),
                packet_parser.payload().own_buffer()
        };

        {
            unique_lock queue_lock(fragment_buffer.buffer_lock);

            if(!fragment_buffer.insert_index(packet_parser.packet_id(), std::move(fragment_entry)))
                return PacketDecodeResult::COMMAND_INSTERT_FAILED;
        }

        this->callback_send_acknowledge(this->callback_argument, packet_parser.packet_id(), packet_parser.type() == protocol::COMMAND_LOW);
        commands_pending = true;
    } else {
        this->callback_decoded_packet(this->callback_argument, packet_parser);
    }
}

bool PacketDecoder::verify_encryption(const pipes::buffer_view &buffer) {
    ClientPacketParser packet_parser{buffer};
    if(!packet_parser.valid() || !packet_parser.is_encrypted()) return false;

    assert(packet_parser.type() >= 0 && packet_parser.type() < this->incoming_generation_estimators.size());
    return this->crypt_handler_->verify_encryption(buffer, packet_parser.packet_id(), this->incoming_generation_estimators[packet_parser.type()].generation());
}

CommandReassembleResult PacketDecoder::reassemble_command(pipes::buffer &result, bool &is_command_low) {
    bool more_commands_pending{false};
    command_fragment_buffer_t* buffer{nullptr};
    unique_lock<std::recursive_timed_mutex> buffer_lock; /* general buffer lock */

    {
        //FIXME: Currently command low packets cant be handled if there is a command packet stuck in reassemble queue

        /* handle commands before command low packets */
        for(auto& buf : this->_command_fragment_buffers) {
            unique_lock ring_lock(buf.buffer_lock, try_to_lock);
            if(!ring_lock.owns_lock()) continue;

            if(buf.front_set()) {
                if(!buffer) { /* lets still test for reexecute */
                    buffer_lock = move(ring_lock);
                    buffer = &buf;
                } else {
                    more_commands_pending = true;
                    break;
                }
            }
        }
    }

    if(!buffer)
        return CommandReassembleResult::NO_COMMANDS_PENDING;

    uint8_t packet_flags{0};
    pipes::buffer payload{};

    /* lets find out if we've to reassemble the packet */
    auto& first_buffer = buffer->slot_value(0);
    if(first_buffer.packet_flags & PacketFlag::Fragmented) {
        uint16_t sequence_length{1};
        size_t total_payload_length{first_buffer.payload_length};
        do {
            if(sequence_length >= buffer->capacity())
                return CommandReassembleResult::SEQUENCE_LENGTH_TOO_LONG;

            if(!buffer->slot_set(sequence_length))
                return CommandReassembleResult::NO_COMMANDS_PENDING; /* we need more packets */

            auto& packet = buffer->slot_value(sequence_length++);
            total_payload_length += packet.payload_length;
            if(packet.packet_flags & PacketFlag::Fragmented) {
                /* yep we find the end */
                break;
            }
        } while(true);
        /* ok we have all fragments lets reassemble */

        /*
         * Packet sequence could never be so long. If it is so then the data_length() returned an invalid value.
         * We're checking it here because we dont want to make a huge allocation
         */
        assert(total_payload_length < 512 * 1024 * 1024);

        pipes::buffer packet_buffer{total_payload_length};
        char* packet_buffer_ptr = &packet_buffer[0];
        size_t packet_count{0};

        packet_flags = buffer->slot_value(0).packet_flags;
        while(packet_count < sequence_length) {
            auto fragment = buffer->pop_front();
            memcpy(packet_buffer_ptr, fragment.payload.data_ptr(), fragment.payload_length);

            packet_buffer_ptr += fragment.payload_length;
            packet_count++;
        }

#ifndef _NDEBUG
        if((packet_buffer_ptr - 1) != &packet_buffer[packet_buffer.length() - 1]) {
            logCritical(0,
                        "Buffer over/underflow: packet_buffer_ptr != &packet_buffer[packet_buffer.length() - 1]; packet_buffer_ptr := {}; packet_buffer.end() := {}",
                        (void*) packet_buffer_ptr,
                        (void*) &packet_buffer[packet_buffer.length() - 1]
            );
        }
#endif
        payload = packet_buffer;
    } else {
        auto packet = buffer->pop_front();
        packet_flags = packet.packet_flags;
        payload = packet.payload;
    }

    more_commands_pending |= buffer->front_set(); /* set the more flag if we have more to process */
    buffer_lock.unlock();

    if(packet_flags & PacketFlag::Compressed) {
        std::string error{};

        auto decompressed_size = compression::qlz_decompressed_size(payload.data_ptr(), payload.length());
        auto uncompressed_buffer = buffer::allocate_buffer(decompressed_size);
        if(!compression::qlz_decompress_payload(payload.data_ptr(), uncompressed_buffer.data_ptr(), &decompressed_size))
            return CommandReassembleResult::COMMAND_DECOMPRESS_FAILED;

        payload = uncompressed_buffer.range(0, decompressed_size);
    }

    result = std::move(payload);
    return more_commands_pending ? CommandReassembleResult::MORE_COMMANDS_PENDING : CommandReassembleResult::SUCCESS;
}

void PacketDecoder::force_insert_command(const pipes::buffer_view &buffer) {
    CommandFragment fragment_entry{
            0,
            0,

            PacketFlag::Unencrypted,
            (uint32_t) buffer.length(),
            buffer.own_buffer()
    };


    {
        auto& fragment_buffer = this->_command_fragment_buffers[command_fragment_buffer_index(protocol::COMMAND)];
        unique_lock queue_lock(fragment_buffer.buffer_lock);
        fragment_buffer.push_front(std::move(fragment_entry));
    }
}

void PacketDecoder::register_initiv_packet() {
    auto& fragment_buffer = this->_command_fragment_buffers[command_fragment_buffer_index(protocol::COMMAND)];
    unique_lock buffer_lock(fragment_buffer.buffer_lock);
    fragment_buffer.set_full_index_to(1); /* the first packet (0) is already the clientinitiv packet */
}