#include <tommath.h>
#include <tomcrypt.h>
#include <misc/endianness.h>
#include <misc/digest.h>
#include <misc/base64.h>
#include <openssl/sha.h>
#include <log/LogUtils.h>
#include <src/client/SpeakingClient.h>

#include "VoiceClient.h"
#include "../../InstanceHandler.h"
#include "../../server/VoiceServer.h"

using namespace std;
using namespace std::chrono;
using namespace ts::server;
using namespace ts::protocol;
using namespace ts::connection;

inline void generate_random(uint8_t *destination, size_t length) {
    while(length-- > 0)
        *(destination++) = (uint8_t) rand();
}

ts::command_result VoiceClient::handleCommandClientInitIv(Command& command) {
    this->last_packet_handshake = system_clock::now();

    if(this->state == ConnectionState::CONNECTED) { /* we've a reconnect */
        if(system_clock::now() - this->lastPingResponse < seconds(5)) {
            logMessage(this->getServerId(), "{} Client initialized session reconnect, but last ping response is not older then 5 seconds ({}). Ignoring attempt",
                    CLIENT_STR_LOG_PREFIX,
                    duration_cast<milliseconds>(system_clock::now() - this->lastPingResponse).count()
            );
            return ts::command_result{error::ok};
        } else if(!config::voice::allow_session_reinitialize) {
            logMessage(this->getServerId(), "{} Client initialized session reconnect and last ping response is older then 5 seconds ({}). Dropping attempt because its not allowed due to config settings",
                       CLIENT_STR_LOG_PREFIX,
                       duration_cast<milliseconds>(system_clock::now() - this->lastPingResponse).count()
            );
            return ts::command_result{error::ok};
        }
        logMessage(this->getServerId(), "{} Client initialized reconnect and last ping response is older then 5 seconds ({}). Allowing attempt",
                   CLIENT_STR_LOG_PREFIX,
                   duration_cast<milliseconds>(system_clock::now() - this->lastPingResponse).count()
        );
        {
            unique_lock server_channel_lock(this->server->channel_tree_lock); /* we cant get moved if this is locked! */
            if(this->currentChannel) {
                this->server->client_move(this->ref(), nullptr, nullptr, config::messages::timeout::connection_reinitialized, ViewReasonId::VREASON_TIMEOUT, false, server_channel_lock);
            }
        }

        this->finalDisconnect();
    } else if(this->state == ConnectionState::DISCONNECTING) {
        this->closeConnection(); /* executing final disconnect right now! */
    }
    this->connection->reset();
    {
        auto& read_queue = this->connection->packet_buffers()[protocol::COMMAND];
        unique_lock buffer_lock(read_queue.buffer_lock);
        read_queue.set_full_index_to(1); /* the first packet (0) is already the clientinitiv packet */
    }
    this->state = ConnectionState::INIT_HIGH;

    bool use_teaspeak = command.hasParm("teaspeak");
    if(use_teaspeak) {
        debugMessage(this->getServerId(), "{} Client using TeaSpeak with auth type {}", CLIENT_STR_LOG_PREFIX, command["verify_type"].string());
        this->properties()[property::CLIENT_TYPE_EXACT] = ClientType::CLIENT_TEASPEAK;
    }

    /* normal TeamSpeak handling */
    string clientAlpha = base64::decode(command["alpha"]); //random
    if(clientAlpha.length() != 10) return ts::command_result{error::parameter_invalid};

    string clientOmega = base64::decode(command["omega"]); //The identity public key
    string ip = command["ip"];
    bool ot = command[0].has("ot") ? command["ot"] : false;

    this->crypto.remote_key = std::shared_ptr<ecc_key>(new ecc_key{}, [](ecc_key* key){
        if(!key) return;
        ecc_free(key);
        delete key;
    });

    auto state = ecc_import((const unsigned char *) clientOmega.data(), clientOmega.length(), this->crypto.remote_key.get());
    if(state != CRYPT_OK) {
        this->crypto.remote_key.reset();
        return ts::command_result{error::client_could_not_validate_identity};
    }
    this->properties()[property::CLIENT_UNIQUE_IDENTIFIER] = base64::encode(digest::sha1(command["omega"].string()));

    this->crypto.alpha = clientAlpha;
    this->crypto.new_protocol = !use_teaspeak && ot && config::experimental_31 && (this->crypto.client_time >= 173265950 || this->crypto.client_time == 5680278000UL);
    {
        size_t beta_length = this->crypto.new_protocol ? 54 : 10;
        char beta[beta_length];
        generate_random((uint8_t *) beta, beta_length);
        this->crypto.beta = string(beta, beta_length);
    }

    if(this->crypto.new_protocol) {
        //Pre setup
        //Generate chain
        debugMessage(this->getServerId(), "{} Got client 3.1 protocol with build timestamp {}", CLIENT_STR_LOG_PREFIX, this->crypto.client_time);
        this->crypto.chain_data = serverInstance->getTeamSpeakLicense()->license();
        this->crypto.chain_data->chain->addEphemeralEntry();
        auto rawLicense = this->crypto.chain_data->chain->exportChain();

        //Sign license
        auto serverOmega = this->getServer()->publicServerKey();
        auto rawServerOmega = base64::decode(serverOmega);

        auto licenseHash = digest::sha256(rawLicense);
        size_t signBufferLength = 128;
        char signBuffer[signBufferLength];
        prng_state prngState{};
        memset(&prngState, 0, sizeof(prngState));
        if(ecc_sign_hash((u_char*) licenseHash.data(), licenseHash.length(), (u_char*) signBuffer, &signBufferLength, &prngState, find_prng("sprng"), this->getServer()->serverKey()) != CRYPT_OK) {
            logError(this->getServerId(), "Failed to sign crypto chain!");
            return ts::command_result{error::vs_critical};
        }
        auto proof = base64::encode(signBuffer, signBufferLength);

        Command initivexpand2("initivexpand2");
        initivexpand2["time"] = duration_cast<seconds>(system_clock::now().time_since_epoch()).count();
        initivexpand2["l"] = base64::encode(rawLicense);
        initivexpand2["beta"] = base64::encode(this->crypto.beta);
        initivexpand2["omega"] = serverOmega;
        initivexpand2["proof"] = proof;
        initivexpand2["tvd"] = "";
        initivexpand2["root"] = base64::encode((char*) this->crypto.chain_data->public_key, 32);
        initivexpand2["ot"] = 1;

        this->sendCommand(initivexpand2);
        this->handshake.state = HandshakeState::SUCCEEDED; /* we're doing the verify via TeamSpeak */
    } else {
        debugMessage(this->getServerId(), "{} Got non client 3.1 protocol with build timestamp {}", CLIENT_STR_LOG_PREFIX, this->crypto.client_time);

        auto serverOmega = this->getServer()->publicServerKey();
        if(serverOmega.empty()) {
            logError(this->getServerId(), "Failed to export server public key!");
            return ts::command_result{error::vs_critical};;
        }

        {
            Command initivexpand("initivexpand");
            initivexpand["alpha"] = base64::encode(clientAlpha);
            initivexpand["beta"] = base64::encode(this->crypto.beta);
            initivexpand["omega"] = serverOmega;
            if(use_teaspeak) {
                initivexpand["teaspeak"] = 1;
                this->handshake.state = HandshakeState::BEGIN; /* we need to start the handshake */
            } else {
                this->handshake.state = HandshakeState::SUCCEEDED; /* we're doing the verify via TeamSpeak */
            }
            this->sendCommand0(initivexpand.build(), false, true); //If we setup the encryption now
        }

        {
            this->connection->getCryptHandler()->block(); //Block until the key is setuped
            string error;
            if(!this->connection->getCryptHandler()->setupSharedSecret(this->crypto.alpha, this->crypto.beta, this->crypto.remote_key.get(), this->server->serverKey(), error)){
                this->connection->getCryptHandler()->unblock();
                logError(this->server->getServerId(), "Could not setup shared secret! (" + error + ")");
                return ts::command_result{error::vs_critical};
            }
            this->connection->getCryptHandler()->unblock();
            this->crypto.protocol_encrypted = true;
        }
    }
    return ts::command_result{error::ok};
}

ts::command_result VoiceClient::handleCommandClientEk(const std::unique_ptr<protocol::ClientPacket>& packet, Command& cmd) {
    this->last_packet_handshake = system_clock::now();
    debugMessage(this->getServerId(), "{} Got client ek!", CLIENT_STR_LOG_PREFIX);

    auto client_key = base64::decode(cmd["ek"]);
    auto x = this->crypto.chain_data->chain->generatePrivateKey(this->crypto.chain_data->root_key, this->crypto.chain_data->root_index);

    this->connection->getCryptHandler()->setupSharedSecretNew(this->crypto.alpha, this->crypto.beta, (char*) x.data(), client_key.data());
    this->connection->acknowledge_handler.reset();
    this->crypto.protocol_encrypted = true;
    this->sendAcknowledge(packet->packetId()); //Send the encrypted acknowledge
    return ts::command_result{error::ok};
}