//
// Created by wolverindev on 08.10.17.
//

#include <tommath.h>
#include <bitset>
#include <openssl/sha.h>
#include "Connection.h"
#include "misc/base64.h"
#include "misc/endianness.h"

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

const int InitVersionLength = 4;
const uint8_t InitVersion[InitVersionLength] = {0x09, 0x83, 0x8C, 0xCF};

/**
 * Maybe memset to 0 for security?
 */
#define RESET_DATA \
bufferIndex = 0;   \
delete pkt;        \
pkt = nullptr;


inline ClientPacket *solvePuzzle(shared_ptr<ServerPacket> response, Identity *, std::string &);

inline std::string toString(mp_int* num){
    char buffer[2048];
    memset(buffer, 0, 2048);
    auto len = mp_todecimal(num, buffer);
    return string(buffer);
}


extern void hexdump(std::ostream& outs, const std::string& s, size_t line_len = 16);
bool ServerConnection::handshake(std::string &errorMessage) {
    //setup the init mac
    /**
     * Low level
     */
    ts::protocol::ClientPacket *pkt;
    shared_ptr<ServerPacket> response;
    int maxBufferSize = 512;
    size_t bufferIndex = 0;
    uint8_t buffer[maxBufferSize];
	memset(buffer, 0, maxBufferSize);
    int err = 0;
    string error = "success";

	beginCoocie:
    memcpy(buffer, InitVersion, InitVersionLength);
    bufferIndex += InitVersionLength;
    buffer[bufferIndex++] = 0x00; //Login state

    auto millis = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    memcpy(&buffer[bufferIndex], &millis, 4);
    bufferIndex += 4;
    //generate the alpha key
    for (int i = 0; i < 4; i++) buffer[bufferIndex++] = (uint8_t) std::rand();
    bufferIndex += 8; //Reserved bytes

	pkt = new ts::protocol::ClientPacket(ts::protocol::PacketTypeInfo::Init1, pipes::buffer_view((void *) buffer, bufferIndex));
	pkt->clientId(0);
	pkt->toggle(ts::protocol::PacketFlag::Unencrypted, true);
	pkt->applyPacketId(101, 0);
	this->sendPacket(*pkt);
	RESET_DATA;

	response = readNextPacket();
	if (!response) {
		errorMessage = "could not get a valid response!";
		return false;
	}
	if (response->type() != protocol::PacketTypeInfo::Init1) {
		errorMessage = "invalid response type. Got: " + response->type().name();
		return false;
	}
	if (response->data()[0] != 1) {
		errorMessage = "iInvalid requested login type (" + to_string((int) response->data()[0]) + " == 1)";
		return false;
	}

    //the second request of the manager
    memcpy(buffer, InitVersion, InitVersionLength);
    bufferIndex += InitVersionLength;
    buffer[bufferIndex++] = 0x02; //Login state
    if(response)
        memcpy(&buffer[bufferIndex], response->data().string().substr(1, 16).data(), 16);
    bufferIndex += 16; //Servers 16 bytes
	if(response)
        memcpy(&buffer[bufferIndex], response->data().string().substr(17, 4).data(), 4);
    bufferIndex += 4; //My own 16 bytes, reversed

    pkt = new ts::protocol::ClientPacket(ts::protocol::PacketTypeInfo::Init1, pipes::buffer_view((void *) buffer, bufferIndex));
    pkt->clientId(0);
    pkt->toggle(ts::protocol::PacketFlag::Unencrypted, true);
    pkt->applyPacketId(101, 0);
	this->sendPacket(*pkt);
    RESET_DATA;

    //We got the RSA challenge
    response = readNextPacket();
    if (!response) {
        errorMessage = "could not get a valid response!";
        return false;
    }
    if (response->type() != protocol::PacketTypeInfo::Init1) {
        errorMessage = "invalid response type";
        return false;
    }
    if (response->data()[0] != 3) {
	    if(response->data()[0] == 127) {
		    cout << "COOCIE RESET!" << endl;
		    goto beginCoocie;
	    }
        hexdump(cout, response->data().string());
        errorMessage = "Invalid requested login type (" + to_string((int) response->data()[0]) + " == 3 | unencripted -> " + (response->hasFlag(PacketFlag::Unencrypted) ? "true" : "false") + ")";
        return false;
    }

    //Generate puzzel response
	std::string alpha;
    pkt = solvePuzzle(response, this->clientIdentity, alpha);
	pkt->applyPacketId(101, 0);
	this->sendPacket(*pkt);
	RESET_DATA;

    cout << "manager init done" << endl;
    this->encriptAck = true;

    response = readNextPacket();
    if (!response) {
        errorMessage = "could not get a valid response!";
        return false;
    }
    if (response->type() != protocol::PacketTypeInfo::Command) {
        errorMessage = "invalid response type: " + response->type().name();
        return false;
    }


    auto command = response->asCommand();
    if (command.getCommand().compare("initivexpand") != 0) {
       // errorMessage = "invalid response command. Got: " + command.getCommand() + " Expected: initivexpand";
        return this->handshakeNew(command, alpha, errorMessage);
    }

    //std::string alpha = base64::decode(command[0]["alpha"]);
    std::string beta = base64::decode(command[0]["beta"]);
    std::string omega = base64::decode(command[0]["omega"]); //Remotes public key
    cout << "RESPONSE! -> " << command.build() << endl;
    //Read public key
    ecc_key remotePublicKey{};
    if ((err = ecc_import((const unsigned char *) omega.data(), omega.length(), &remotePublicKey)) != CRYPT_OK) {
        errorMessage = "ecc_import(...) returned " + to_string(err) + "/" + error_to_string(err);
        return false;
    }

    if(strcmp(remotePublicKey.dp->name, "ECC-256") != 0){
        errorMessage = "invalid imported public key! Curve found " + string(remotePublicKey.dp->name);
        return false;
    }

    size_t sharedSecretLength = 32;
    char sharedSecret[sharedSecretLength];

    if ((err = ecc_shared_secret(clientIdentity->getKeyPair(), &remotePublicKey, (unsigned char *) sharedSecret, &sharedSecretLength)) != CRYPT_OK) {
        errorMessage = "ecc_shared_secret(...) returned " + to_string(err) + "/" + error_to_string(err);
        return false;
    }

    if (!setupSharedSecret(alpha, beta, string(sharedSecret, sharedSecretLength), error)) {
        errorMessage = "setupSharedSecret(...) failed: " + error;
        return false;
    }

    //this->readQueue[PacketType::Command.type()]->reset();

    //TS 3.1
    /*
    response = readNextPacket();
    if (!response) {
        errorMessage = "could not get a valid response!";
        return false;
    }
    if (response->type() != protocol::PacketType::Command) {
        errorMessage = "invalid response type: " + response->type().name();
        return false;
    }
    command = response->asCommand();
    cout << "Having initiv2 -> " << response->data() << endl;
     */

    this->idManager.nextPacketId(PacketTypeInfo::Command);
    Command clientinit("clientinit");
    //94ec66de-5940-4e38-b002-970df0cf6c94,62444179-0d99-42ba-a45c-c6b1557d079a,d95f9901-c42d-4bac-8849-7164fd9e2310
    //clientinit["client_badges"] = "badges=450f81c1-ab41-4211-a338-222fa94ed157,c9e97536-5a2d-4c8e-a135-af404587a472,94ec66de-5940-4e38-b002-970df0cf6c94"; //,62444179-0d99-42ba-a45c-c6b1557d079a
    clientinit["client_nickname"] = "Wolf C++ XXXX";
    clientinit["client_version"] = "3.1 [Build: 1471417187]";
    clientinit["client_platform"] = "Windows";
    clientinit["client_version_sign"] = "Vr9F7kbVorcrkV5b/Iw+feH9qmDGvfsW8tpa737zhc1fDpK5uaEo6M5l2DzgaGqqOr3GKl5A7PF9Sj6eTM26Aw==";
    clientinit["client_input_hardware"] = true;
    clientinit["client_output_hardware"] = true;
    clientinit["client_default_channel"] = "";
    clientinit["client_default_channel_password"] = "";

    string password;
    if(!password.empty()){
        char passwordBuffer[SHA_DIGEST_LENGTH];
        SHA1((const unsigned char *) password.data(), password.length(), (unsigned char *) passwordBuffer);
        password = base64_encode(string(passwordBuffer, SHA_DIGEST_LENGTH));
    }
    clientinit["client_server_password"] = password;
    clientinit["client_meta_data"] = "";
    clientinit["client_key_offset"] = this->clientIdentity->lastValidKeyOffset();
    clientinit["client_nickname_phonetic"] = "";
    clientinit["client_default_token"] = "";
    clientinit["hwid"] = "123,456123123123";
    sendCommand(clientinit);

    while(true){
        response = readNextPacket();
        if (!response) {
            errorMessage = "could not get a valid response!";
            return false;
        }
        if(response->type() == PacketTypeInfo::Ack) continue;
        break;
    }

    //TODO check ack id

    if (!response) {
        errorMessage = "could not get a valid response!";
        return false;
    }
    if (response->type() != protocol::PacketTypeInfo::Command) {
        errorMessage = "invalid response type: " + response->type().name();
        return false;
    }
    if(response->asCommand().getCommand() == "initserver"){ //Got success
        this->handleQueueLock.lock();
        this->handleQueue.push_front(response);
        this->handleQueueLock.unlock();

        this->setClientId(response->asCommand()["aclid"]);

        this->autoHandle = true;
        cout << "Successfull connected!" << endl;

        /*
        std::thread([&](){
            usleep(1000 * 1000);
            cout << " -> send extra command" << endl;
            //this->sendCommand(Command("channelsubscribeall return_code=1:i"));

            while(true){
                //this->sendCommand(Command("getconnectioninfo clid=320 return_code=1:112"));
                //this->sendCommand(Command("ftgetfilelist cid=0 cpw path=\\/icons return_code=1:z0"));
                this->sendCommand(Command("servergrouppermlist sgid=6 return_code=1:112"));
                usleep(10 * 1000 * 1000);
            }
            //Command cmd("channelgetdescription cid=1 return_code=1:3o");
            //Command cmd("clientupdate client_nickname=WolverinDEV22 return_code=__1_");
            //Command cmd("clientdisconnect reasonid=8 reasonmsg=leaving");
            //this->sendCommand(Command("permissionlist return_code=__1_"));
            //this->sendCommand(Command("clientgetvariables clid=" + to_string(this->clientId)));
        }).detach();
         */


        return true;
    }
    cout << "Invalid connect: " << response->data() << endl;
    //TODO error handling
    return true;
}

inline ClientPacket* solvePuzzle(shared_ptr<ServerPacket> response, Identity *identity, std::string &alpha) {
    uint32_t puzzelLength = be2le32(&((char*) response->data().data_ptr())[1 + 128]); //1 for the first byte (the state byte)

    auto buffer = (char*) response->data().data_ptr();

    char alphaBuffer[10];
    for (int index = 0; index < 10; index++)
        alphaBuffer[index] = 0; //rand();
	alpha = string(alphaBuffer, 10);

    //Generating command
    auto pkey = identity->publicKey();
    ts::Command command("clientinitiv", {
            {"alpha", base64_encode(alphaBuffer, 10)},
            {"omega", pkey},
            {"ip", ""},
            {"ot", 1}                       //Required by 3.1
    });
    std::string cmd = command.build();

    //Sloving puzzel
    mp_int x{};
    mp_int n{};
    mp_int result{};

    //mp_init_multi(&x, &n, &result);
    mp_init(&x);
    mp_init(&n);
    mp_init(&result);

    char numBuffer[2048];
    mp_read_unsigned_bin(&x, (const unsigned char *) &response->data()[1], 64); //One offset
    mp_read_unsigned_bin(&n, (const unsigned char *) &response->data()[1 + 64], 64); //1 + 64 offset

    cout << "X: " << toString(&x) << endl;
    cout << "N: " << toString(&n) << endl;
    cout << "Length: " << puzzelLength << endl;

    mp_int exp{};
    mp_init(&exp);
    mp_2expt(&exp, puzzelLength);


    //x ** (2 ** puzzelLength) mod n
    int err = 0;
    if ((err = mp_exptmod(&x, &exp, &n, &result)) != CRYPT_OK) {
        cerr << "Invalid crypt: " << err << "/" << error_to_string(err) << endl;
    }

    int resultBufferLength = mp_unsigned_bin_size(&result);
    char resultBuffer[resultBufferLength];
    mp_to_unsigned_bin(&result, (unsigned char *) resultBuffer);


    //mp_clear_multi(&x, &n, &exp, &result);
    mp_clear(&x);
    mp_clear(&n);
    mp_clear(&exp);
    mp_clear(&result);

    size_t packetBufferLength = InitVersionLength + 1 + 232 + 64 + cmd.length();
    char packetBuffer[packetBufferLength];
    memset(packetBuffer, 0, packetBufferLength);

    memcpy(packetBuffer, InitVersion, InitVersionLength);
    packetBuffer[InitVersionLength] = 0x04;

    //Copy old data
    memcpy(&packetBuffer[InitVersionLength + 1], &response->data()[1], 232);
    memcpy(&packetBuffer[InitVersionLength + 1 + 232 + (64 - resultBufferLength)], resultBuffer, resultBufferLength);
    memcpy(&packetBuffer[InitVersionLength + 1 + 232 + 64], cmd.data(), cmd.length());

    cout << "sending puzzel sulution" << endl;
    auto pkt = new ts::protocol::ClientPacket(ts::protocol::PacketTypeInfo::Init1, pipes::buffer_view((void *) packetBuffer, packetBufferLength));
    pkt->clientId(0);
    pkt->toggle(ts::protocol::PacketFlag::Unencrypted, true);
    return pkt;
}