#include "assert.hpp"
#include "../include/odhtdb/Log.hpp"
#include "../include/odhtdb/Database.hpp"
#include "../include/odhtdb/Group.hpp"
#include "../include/odhtdb/Encryption.hpp"
#include "../include/odhtdb/Hash.hpp"
#include "../include/odhtdb/hex2bin.hpp"
#include "../include/odhtdb/bin2hex.hpp"
#include "../include/odhtdb/DataView.hpp"
#include <vector>
#include <chrono>
#include <thread>
#include <boost/filesystem.hpp>
#include <map>
#include <sibs/BootstrapNode.hpp>

using namespace std;
using namespace chrono_literals;
using namespace odhtdb;

static const u16 PORT = 22111;

static void testBinHexConvert()
{
    DataView input { (void*)"hello", 5 };
    
    string inputHex = bin2hex((const char*)input.data, input.size);
    assertEquals<string>("68656c6c6f", inputHex);
    
    string inputBin = hex2bin(inputHex.c_str(), inputHex.size());
    if(inputBin.size() != input.size)
    {
        string errMsg = "Expected input converted to hex and then back to binary size to be 5, was: ";
        errMsg += to_string(inputBin.size());
        fail(errMsg);
    }
    
    if(memcmp(input.data, inputBin.data(), input.size) != 0)
    {
        string errMsg = "Expected input converted to hex and then back to binary to be same as original input ('hello'), was: ";
        errMsg += string(inputBin.c_str(), inputBin.size());
        fail(errMsg);
    }
}

static void testHash()
{
    Hash hash("odhtdb", 6);
    string hashHex = hash.toString();
    assertEquals<string>("a7b30ec8ab92de60e551b26bb8f78d315697f84dd7f5549a143477e095ec934f", hashHex);
    Log::debug("hash of 'odhtdb' is: a7b30ec8ab92de60e551b26bb8f78d315697f84dd7f5549a143477e095ec934f");
}

static void testSignData(const Signature::KeyPair &localUserKeyPair)
{
    std::string publicKeyStr = localUserKeyPair.getPublicKey().toString();
    Log::debug("Local user public key: %s", publicKeyStr.c_str());
    
    std::string privateKeyStr = localUserKeyPair.getPrivateKey().toString();
    Log::debug("Local user private key: %s", privateKeyStr.c_str());
    
    string expectedUnsignedData = "hello, world!";
    string signedData = localUserKeyPair.getPrivateKey().sign(DataView((void*)expectedUnsignedData.data(), expectedUnsignedData.size()));
    assertEquals(SIGNED_HASH_SIZE + expectedUnsignedData.size(), signedData.size());
    
    string unsignedData = localUserKeyPair.getPublicKey().unsign(DataView((void*)signedData.data(), signedData.size()));
    assertEquals(expectedUnsignedData, unsignedData);
    
    try
    {
        Signature::KeyPair anotherKeyPair;
        anotherKeyPair.getPublicKey().unsign(DataView((void*)signedData.data(), signedData.size()));
        fail("Expected unsign to fail since the data was not signed with the private key matching given public key");
    }
    catch (UnsignWrongKeyException &e)
    {
        
    }
    catch (exception &e)
    {
        string errMsg = "Expected unsign to fail with UnsignWrongKeyException, fail reason: ";
        errMsg += e.what();
        fail(errMsg);
    }
    
    try
    {
        Signature::KeyPair anotherKeyPair;
        anotherKeyPair.getPublicKey().unsign(DataView((void*)signedData.data(), 3));
        fail("Expected unsign to fail since the (signed) data is too small to have been signed (signed hash is 64 bytes)");
    }
    catch (UnsignInvalidSizeException &e)
    {
        
    }
    catch (exception &e)
    {
        string errMsg = "Expected unsign to fail with UnsignInvalidSizeException, fail reason: ";
        errMsg += e.what();
        fail(errMsg);
    }
}

static void testEncryption()
{
    const char *message = "hello, world!";
    const unsigned long long messageLength = 13;
    Encryption encryption(DataView((void*)message, messageLength));
    
    Decryption decryption(encryption.getCipherText(), encryption.getNonce(), encryption.getKey());
    assertEquals<unsigned long long>(messageLength, decryption.getDecryptedText().size);
    assertEquals(0, strncmp(message, (const char*)decryption.getDecryptedText().data, messageLength));
}

static void testTimestamp(const std::unique_ptr<Database> &database)
{
    auto timestamp1 = database->getSyncedTimestampUtc();
    this_thread::sleep_for(chrono::milliseconds(100));
    auto timestamp2 = database->getSyncedTimestampUtc();
    
    if(timestamp2.getCombined() > timestamp1.getCombined())
        Log::debug("Second timestamp is more than first one, as expected");
    else
        fail("Second timestamp is not more than first one for some reason");
}

static void testStandard()
{
    testBinHexConvert();
    testHash();
    testEncryption();
    
    boost::filesystem::path storagePath("/tmp/odhtdbTest");
    boost::filesystem::remove_all(storagePath);
    boost::filesystem::create_directory(storagePath);
    
    int createNodeCounter = 0;
    int addDataCounter = 0;
    int addUserCounter = 0;
    
    auto createNodeCallback = [&createNodeCounter](const DatabaseCreateNodeRequest &request)
    {
        Log::debug("Create node callback");
        ++createNodeCounter;
    };
    
    auto addNodeCallback = [&addDataCounter](const DatabaseAddNodeRequest &request)
    {
        Log::debug("Add node callback");
        ++addDataCounter;
    };
    
    auto addUserCallback = [&addUserCounter](const DatabaseAddUserRequest &request)
    {
        Log::debug("Add user callback");
        ++addUserCounter;
    };

    DatabaseCallbackFuncs callbackFuncs { createNodeCallback, addNodeCallback, addUserCallback };
    DatabaseNode databaseNode;
    
    {
        Signature::KeyPair localUserKeyPair;
        testSignData(localUserKeyPair);
        
        std::unique_ptr<Database> database = Database::connect("127.0.0.1", PORT, storagePath, callbackFuncs).get();
        testTimestamp(database);

        auto databaseCreateResponse = database->create();
        databaseNode = { databaseCreateResponse->getNodeEncryptionKey(), databaseCreateResponse->getRequestHash() };
        auto adminUserKey = databaseCreateResponse->getNodeAdminKeyPair();
        database->addData(databaseNode, *adminUserKey, DataView{ (void*)"hello, world!", 13 });
        database->addUser(databaseNode, *adminUserKey, localUserKeyPair.getPublicKey(), databaseCreateResponse->getNodeAdminGroupId()->getView());
        database->addData(databaseNode, localUserKeyPair, DataView{ (void*)"hello, aaald!", 13 });
        this_thread::sleep_for(chrono::seconds(3));
        
        assertEquals(1, createNodeCounter);
        assertEquals(2, addDataCounter);
        assertEquals(1, addUserCounter);
        
        string username = "dec05eba";
        string password = "secretPassword";
        if(database->doesStoredUserExist(username))
            fail("Expected stored to not exist until it has been added");
        
        database->storeNodeInfoForUserEncrypted(databaseNode, username, password, *adminUserKey);
        try
        {
            database->storeNodeInfoForUserEncrypted(databaseNode, username, password, localUserKeyPair);
            fail("Expected store user password to fail since we have already stored an user in the node");
        }
        catch(SqlExecException &e)
        {
            Log::debug("Failed with sql exception as expected, since we already have an user in the node: %s", e.what());
        }
        
        if(!database->doesStoredUserExist(username))
            fail("Expected stored to exist after it has been added");
        
        auto nodeUserData = database->getStoredNodeUserInfoDecrypted(username, password);
        assertEquals((size_t)1, nodeUserData.size());
        auto userDataIt = nodeUserData.find(*databaseNode.getRequestHash());
        if(userDataIt == nodeUserData.end())
            fail("Expected stored node hash to match node hash");
        if(userDataIt->second.userKeyPair->getPublicKey() != adminUserKey->getPublicKey())
            fail("Expected stored public key to match admin user public key");
        if(userDataIt->second.userKeyPair->getPrivateKey() != adminUserKey->getPrivateKey())
            fail("Expected stored private key to match admin user private key");
        if(userDataIt->second.nodeEncryptionKey->getView() != databaseCreateResponse->getNodeEncryptionKey()->getView())
            fail("Expected stored node encryption key to match node encryption key");
        
        try
        {
            database->storeUserWithoutNodes(username, password);
            fail("Expected store user to fail since the user already exists in database");
        }
        catch(SqlExecException &e)
        {
            
        }
        
        database->storeUserWithoutNodes("anotherUser", password);
        if(!database->doesStoredUserExist("anotherUser"))
            fail("Added user 'anotherUser' to database without any nodes, but it doesn't seem to be stored");
        
        auto adminUserGroups = database->getUserGroups(*databaseNode.getRequestHash(), adminUserKey->getPublicKey());
        if(adminUserGroups.size() != 1 || adminUserGroups[0].getView() != databaseCreateResponse->getNodeAdminGroupId()->getView())
            fail("Admin group doesn't match group stored in database");
        
        auto userGroups = database->getUserGroups(*databaseNode.getRequestHash(), localUserKeyPair.getPublicKey());
        if(userGroups.size() != 1 || userGroups[0].getView() != databaseCreateResponse->getNodeAdminGroupId()->getView())
            fail("User group doesn't match group stored in database");

        // Give time for us to disconnect
        std::this_thread::sleep_for(std::chrono::seconds(5));
    }
    Log::debug("Callback works when adding data while connected, now testing to reconnect and check if data remains...");
    {
        createNodeCounter = 0;
        addDataCounter = 0;
        addUserCounter = 0;
        
        std::unique_ptr<Database> database = Database::connect("127.0.0.1", PORT, storagePath, callbackFuncs).get();
        database->loadNode(*databaseNode.getRequestHash());

        database->seed(databaseNode, DatabaseFetchOrder::OLDEST_FIRST);
        this_thread::sleep_for(chrono::seconds(3));
        
        assertEquals(1, createNodeCounter);
        assertEquals(2, addDataCounter);
        assertEquals(1, addUserCounter);
        
        InfoHash customMessageKey = Database::getInfoHash("asdf", 4);
        sibs::SafeSerializer messageToSendSerializer;
        messageToSendSerializer.add((u32)10);
        
        u32 receivedNumber = 0;
        database->receiveCustomMessage(customMessageKey, [&receivedNumber](const void *data, usize size)
        {
            sibs::SafeDeserializer deserializer((const u8*)data, size);
            receivedNumber = deserializer.extract<u32>();
            sibs::SafeSerializer serializer;
            serializer.add((u32)20);
            return serializer;
        });
        
        u32 sendCustomMessageResponseNumber = 0;
        database->sendCustomMessage(customMessageKey, messageToSendSerializer.getBuffer().data(), messageToSendSerializer.getBuffer().size(), [&sendCustomMessageResponseNumber](bool gotResponse, const void *data, usize size)
        {
            if(!gotResponse)
            {
                Log::error("Didn't get reponse!");
                return true;
            }
            
            sibs::SafeDeserializer deserializer((const u8*)data, size);
            sendCustomMessageResponseNumber = deserializer.extract<u32>();
            return false;
        });
        
        
        this_thread::sleep_for(chrono::seconds(3));
        assertEquals((u32)10, receivedNumber);
        assertEquals((u32)20, sendCustomMessageResponseNumber);
    }
}

static void testTwoLocalNodes()
{
    boost::filesystem::path storagePath1("/tmp/odhtdbTest1");
    boost::filesystem::remove_all(storagePath1);
    boost::filesystem::create_directory(storagePath1);
    
    boost::filesystem::path storagePath2("/tmp/odhtdbTest2");
    boost::filesystem::remove_all(storagePath2);
    boost::filesystem::create_directory(storagePath2);
    
    auto createNodeCallback = [](const DatabaseCreateNodeRequest &request)
    {
        Log::debug("Create node callback");
    };
    
    auto addNodeCallback = [](const DatabaseAddNodeRequest &request)
    {
        Log::debug("Add node callback");
    };
    
    auto addUserCallback = [](const DatabaseAddUserRequest &request)
    {
        Log::debug("Add user callback");
    };

    DatabaseCallbackFuncs callbackFuncs { createNodeCallback, addNodeCallback, addUserCallback };
    
    std::unique_ptr<Database> database1 = Database::connect("127.0.0.1", PORT, storagePath1, callbackFuncs).get();
    auto databaseCreateResponse = database1->create();
    DatabaseNode databaseNode = { databaseCreateResponse->getNodeEncryptionKey(), databaseCreateResponse->getRequestHash() };
    auto adminUserKey = databaseCreateResponse->getNodeAdminKeyPair();
    database1->addData(databaseNode, *adminUserKey, DataView{ (void*)"hello, world!", 13 });
    
    std::unique_ptr<Database> database2 = Database::connect("127.0.0.1", PORT, storagePath2, callbackFuncs).get();
    database2->seed(databaseNode);

    this_thread::sleep_for(chrono::seconds(5));
}

static void testMemoryUsage()
{
    boost::filesystem::path storagePath("/tmp/odhtdbTestMemoryUsage");
    boost::filesystem::remove_all(storagePath);
    boost::filesystem::create_directory(storagePath);
    
    auto createNodeCallback = [](const DatabaseCreateNodeRequest &request)
    {
        Log::debug("Create node callback");
    };
    
    auto addNodeCallback = [](const DatabaseAddNodeRequest &request)
    {
        Log::debug("Add node callback");
    };
    
    auto addUserCallback = [](const DatabaseAddUserRequest &request)
    {
        Log::debug("Add user callback");
    };

    DatabaseCallbackFuncs callbackFuncs { createNodeCallback, addNodeCallback, addUserCallback };
    
    std::unique_ptr<Database> database = Database::connect("127.0.0.1", PORT, storagePath, callbackFuncs).get();
    auto databaseCreateResponse = database->create();
    DatabaseNode databaseNode = { databaseCreateResponse->getNodeEncryptionKey(), databaseCreateResponse->getRequestHash() };
    auto adminUserKey = databaseCreateResponse->getNodeAdminKeyPair();
    
    const char *msg = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.";
    const usize msgLength = strlen(msg);
    
    const int iterations = 500;
    for(int i = 0; i < iterations; ++i)
    {
        Log::debug("Memory usage test %d/%d", 1 + i, iterations);
        database->addData(databaseNode, *adminUserKey, DataView{ (void*)msg, msgLength });
        this_thread::sleep_for(chrono::milliseconds(250));
    }
}

static pair<bool, shared_ptr<OwnedByteArray>> __attribute__((optimize("O0"))) getDataNoCopy()
{
    u8 *decryptionKeyRawCopy = new u8[1024 * 1024 * 64];
    memcpy(decryptionKeyRawCopy, "hello, world!", 14);
    shared_ptr<OwnedByteArray> decryptionKey = make_shared<OwnedByteArray>(decryptionKeyRawCopy, ENCRYPTION_KEY_BYTE_SIZE);
    return make_pair(true, decryptionKey);
}

static void __attribute__((optimize("O0"))) testMemoryLeak()
{
    {
        auto data = getDataNoCopy();
        printf("data: %s\n", data.second->data);
    }
}

static void testStoreAccount()
{
    boost::filesystem::path storagePath("/tmp/odhtdb");
    boost::filesystem::remove_all(storagePath);
    boost::filesystem::create_directory(storagePath);

    auto createNodeCallback = [](const DatabaseCreateNodeRequest &request)
    {
        Log::debug("Create node callback");
    };
    
    auto addNodeCallback = [](const DatabaseAddNodeRequest &request)
    {
        Log::debug("Add node callback");
    };
    
    auto addUserCallback = [](const DatabaseAddUserRequest &request)
    {
        Log::debug("Add user callback");
    };

    {
        DatabaseCallbackFuncs callbackFuncs { createNodeCallback, addNodeCallback, addUserCallback };
        std::unique_ptr<Database> database = Database::connect("127.0.0.1", PORT, storagePath, callbackFuncs).get();

        std::string username = "username";
        std::string password = "password";

        database->storeUserWithoutNodes(username, password);
        database->getStoredNodeUserInfoDecrypted(username, password);
    }
    {
        DatabaseCallbackFuncs callbackFuncs { createNodeCallback, addNodeCallback, addUserCallback };
        std::unique_ptr<Database> database = Database::connect("127.0.0.1", PORT, storagePath, callbackFuncs).get();

        std::string username = "username";
        std::string password = "password";

        database->getStoredNodeUserInfoDecrypted(username, password);
    }
}

struct Test
{
    function<void()> testFunc;
    bool optional;
};

int main(int argc, char **argv)
{
    sibs::BootstrapNode bootstrapNode(sibs::Ipv4(nullptr, PORT));
    // Wait until bootstrap node is ready to accept connections
    this_thread::sleep_for(chrono::seconds(2));

    map<string, Test> testByName;
    testByName["standard"] = { testStandard, false };
    testByName["two_local_nodes"] = { testTwoLocalNodes, false };
    testByName["memory_usage"] = { testMemoryUsage, true };
    testByName["test_memory_leak"] = { testMemoryLeak, true };
    testByName["store_account"] = { testStoreAccount, false };
    
    const char *testName = "all";
    if(argc > 1)
        testName = argv[1];
    
    if(strcmp(testName, "all") == 0)
    {
        Log::debug("Running all non-optional tests");
        for(auto &testIt : testByName)
        {
            if(!testIt.second.optional)
            {
                Log::debug("Running test: %s", testIt.first.c_str());
                testIt.second.testFunc();
                Log::debug("Finished test: %s", testIt.first.c_str());
            }
        }
    }
    else
    {
        auto testIt = testByName.find(testName);
        if(testIt == testByName.end())
        {
            Log::error("There is no test called %s", testName);
            exit(1);
        }
        
        Log::debug("Running test: %s", testIt->first.c_str());
        testIt->second.testFunc();
        Log::debug("Finished test: %s", testIt->first.c_str());
    }
    
    return 0;
}