// Copyright (c) 2012-2016 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "dbwrapper.h"
#include "random.h"
#include "test/test_bitcoin.h"
#include "uint256.h"

#include <boost/assert.hpp>
#include <boost/assign/std/vector.hpp> // for 'operator+=()'
#include <boost/test/unit_test.hpp>

#include <random>
#include <chrono>

// Test if a string consists entirely of null characters
bool is_null_key(const std::vector<uint8_t> &key) {
    bool isnull = true;

    for (unsigned int i = 0; i < key.size(); i++)
        isnull &= (key[i] == '\x00');

    return isnull;
}

BOOST_FIXTURE_TEST_SUITE(dbwrapper_tests, BasicTestingSetup)

namespace {

struct ScopedPathDeleter
{
    ScopedPathDeleter(fs::path& p)
    : p(p)
    {}

    ~ScopedPathDeleter()
    {
        fs::remove_all(p);
    }

    fs::path& p;
};

}

BOOST_AUTO_TEST_CASE(dbwrapper) {
    // Perform tests both obfuscated and non-obfuscated.
    for (int i = 0; i < 2; i++) {
        bool obfuscate = (bool)i;
        fs::path ph = fs::temp_directory_path() / fs::unique_path();
        CDBWrapper dbw(ph, (1 << 20), true, false, obfuscate);
        char key = 'k';
        uint256 in = InsecureRand256();
        uint256 res;

        // Ensure that we're doing real obfuscation when obfuscate=true
        BOOST_CHECK(obfuscate !=
                    is_null_key(dbwrapper_private::GetObfuscateKey(dbw)));

        BOOST_CHECK(dbw.Write(key, in));
        BOOST_CHECK(dbw.Read(key, res));
        BOOST_CHECK_EQUAL(res.ToString(), in.ToString());
    }
}

// Test batch operations
BOOST_AUTO_TEST_CASE(dbwrapper_batch) {
    // Perform tests both obfuscated and non-obfuscated.
    for (int i = 0; i < 2; i++) {
        bool obfuscate = (bool)i;
        fs::path ph = fs::temp_directory_path() / fs::unique_path();
        CDBWrapper dbw(ph, (1 << 20), true, false, obfuscate);

        char key = 'i';
        uint256 in = InsecureRand256();
        char key2 = 'j';
        uint256 in2 = InsecureRand256();
        char key3 = 'k';
        uint256 in3 = InsecureRand256();

        uint256 res;
        CDBBatch batch(dbw);

        batch.Write(key, in);
        batch.Write(key2, in2);
        batch.Write(key3, in3);

        // Remove key3 before it's even been written
        batch.Erase(key3);

        dbw.WriteBatch(batch);

        BOOST_CHECK(dbw.Read(key, res));
        BOOST_CHECK_EQUAL(res.ToString(), in.ToString());
        BOOST_CHECK(dbw.Read(key2, res));
        BOOST_CHECK_EQUAL(res.ToString(), in2.ToString());

        // key3 should've never been written
        BOOST_CHECK(dbw.Read(key3, res) == false);
    }
}

BOOST_AUTO_TEST_CASE(dbwrapper_iterator) {
    // Perform tests both obfuscated and non-obfuscated.
    for (int i = 0; i < 2; i++) {
        bool obfuscate = (bool)i;
        fs::path ph = fs::temp_directory_path() / fs::unique_path();
        CDBWrapper dbw(ph, (1 << 20), true, false, obfuscate);

        // The two keys are intentionally chosen for ordering
        char key = 'j';
        uint256 in = InsecureRand256();
        BOOST_CHECK(dbw.Write(key, in));
        char key2 = 'k';
        uint256 in2 = InsecureRand256();
        BOOST_CHECK(dbw.Write(key2, in2));

        std::unique_ptr<CDBIterator> it(
            const_cast<CDBWrapper &>(dbw).NewIterator());

        // Be sure to seek past the obfuscation key (if it exists)
        it->Seek(key);

        char key_res;
        uint256 val_res;

        it->GetKey(key_res);
        it->GetValue(val_res);
        BOOST_CHECK_EQUAL(key_res, key);
        BOOST_CHECK_EQUAL(val_res.ToString(), in.ToString());

        it->Next();

        it->GetKey(key_res);
        it->GetValue(val_res);
        BOOST_CHECK_EQUAL(key_res, key2);
        BOOST_CHECK_EQUAL(val_res.ToString(), in2.ToString());

        it->Next();
        BOOST_CHECK_EQUAL(it->Valid(), false);
    }
}

// Test that we do not obfuscation if there is existing data.
BOOST_AUTO_TEST_CASE(existing_data_no_obfuscate) {
    // We're going to share this fs::path between two wrappers
    fs::path ph = fs::temp_directory_path() / fs::unique_path();
    create_directories(ph);
    ScopedPathDeleter ph_deleter(ph);

    // Set up a non-obfuscated wrapper to write some initial data.
    CDBWrapper *dbw = new CDBWrapper(ph, (1 << 10), false, false, false);
    char key = 'k';
    uint256 in = InsecureRand256();
    uint256 res;

    BOOST_CHECK(dbw->Write(key, in));
    BOOST_CHECK(dbw->Read(key, res));
    BOOST_CHECK_EQUAL(res.ToString(), in.ToString());

    // Call the destructor to free leveldb LOCK
    delete dbw;

    // Now, set up another wrapper that wants to obfuscate the same directory
    CDBWrapper odbw(ph, (1 << 10), false, false, true);

    // Check that the key/val we wrote with unobfuscated wrapper exists and
    // is readable.
    uint256 res2;
    BOOST_CHECK(odbw.Read(key, res2));
    BOOST_CHECK_EQUAL(res2.ToString(), in.ToString());

    // There should be existing data
    BOOST_CHECK(!odbw.IsEmpty());
    // The key should be an empty string
    BOOST_CHECK(is_null_key(dbwrapper_private::GetObfuscateKey(odbw)));

    uint256 in2 = InsecureRand256();
    uint256 res3;

    // Check that we can write successfully
    BOOST_CHECK(odbw.Write(key, in2));
    BOOST_CHECK(odbw.Read(key, res3));
    BOOST_CHECK_EQUAL(res3.ToString(), in2.ToString());
}

// Ensure that we start obfuscating during a reindex.
BOOST_AUTO_TEST_CASE(existing_data_reindex) {
    // We're going to share this fs::path between two wrappers
    fs::path ph = fs::temp_directory_path() / fs::unique_path();
    create_directories(ph);
    ScopedPathDeleter ph_deleter(ph);

    // Set up a non-obfuscated wrapper to write some initial data.
    CDBWrapper *dbw = new CDBWrapper(ph, (1 << 10), false, false, false);
    char key = 'k';
    uint256 in = InsecureRand256();
    uint256 res;

    BOOST_CHECK(dbw->Write(key, in));
    BOOST_CHECK(dbw->Read(key, res));
    BOOST_CHECK_EQUAL(res.ToString(), in.ToString());

    // Call the destructor to free leveldb LOCK
    delete dbw;

    // Simulate a -reindex by wiping the existing data store
    CDBWrapper odbw(ph, (1 << 10), false, true, true);

    // Check that the key/val we wrote with unobfuscated wrapper doesn't exist
    uint256 res2;
    BOOST_CHECK(!odbw.Read(key, res2));
    BOOST_CHECK(!is_null_key(dbwrapper_private::GetObfuscateKey(odbw)));

    uint256 in2 = InsecureRand256();
    uint256 res3;

    // Check that we can write successfully
    BOOST_CHECK(odbw.Write(key, in2));
    BOOST_CHECK(odbw.Read(key, res3));
    BOOST_CHECK_EQUAL(res3.ToString(), in2.ToString());
}

BOOST_AUTO_TEST_CASE(iterator_ordering) {
    fs::path ph = fs::temp_directory_path() / fs::unique_path();
    CDBWrapper dbw(ph, (1 << 20), true, false, false);
    for (int x = 0x00; x < 256; ++x) {
        uint8_t key = x;
        uint32_t value = x * x;
        BOOST_CHECK(dbw.Write(key, value));
    }

    std::unique_ptr<CDBIterator> it(
        const_cast<CDBWrapper &>(dbw).NewIterator());
    for (int c = 0; c < 2; ++c) {
        int seek_start;
        if (c == 0)
            seek_start = 0x00;
        else
            seek_start = 0x80;
        it->Seek((uint8_t)seek_start);
        for (int x = seek_start; x < 256; ++x) {
            uint8_t key;
            uint32_t value;
            BOOST_CHECK(it->Valid());
            // Avoid spurious errors about invalid iterator's  key and value in
            // case of failure
            if (!it->Valid()) break;
            BOOST_CHECK(it->GetKey(key));
            BOOST_CHECK(it->GetValue(value));
            BOOST_CHECK_EQUAL(key, x);
            BOOST_CHECK_EQUAL(value, x * x);
            it->Next();
        }
        BOOST_CHECK(!it->Valid());
    }
}

struct StringContentsSerializer {
    // Used to make two serialized objects the same while letting them have a
    // different lengths. This is a terrible idea.
    std::string str;
    StringContentsSerializer() {}
    StringContentsSerializer(const std::string &inp) : str(inp) {}

    StringContentsSerializer &operator+=(const std::string &s) {
        str += s;
        return *this;
    }
    StringContentsSerializer &operator+=(const StringContentsSerializer &s) {
        return *this += s.str;
    }

    ADD_SERIALIZE_METHODS;

    template <typename Stream, typename Operation>
    inline void SerializationOp(Stream &s, Operation ser_action) {
        if (ser_action.ForRead()) {
            str.clear();
            char c = 0;
            while (true) {
                try {
                    READWRITE(c);
                    str.push_back(c);
                } catch (const std::ios_base::failure &e) {
                    break;
                }
            }
        } else {
            for (size_t i = 0; i < str.size(); i++)
                READWRITE(str[i]);
        }
    }
};

BOOST_AUTO_TEST_CASE(iterator_string_ordering) {
    char buf[10];

    fs::path ph = fs::temp_directory_path() / fs::unique_path();
    CDBWrapper dbw(ph, (1 << 20), true, false, false);
    for (int x = 0x00; x < 10; ++x) {
        for (int y = 0; y < 10; y++) {
            sprintf(buf, "%d", x);
            StringContentsSerializer key(buf);
            for (int z = 0; z < y; z++)
                key += key;
            uint32_t value = x * x;
            BOOST_CHECK(dbw.Write(key, value));
        }
    }

    std::unique_ptr<CDBIterator> it(
        const_cast<CDBWrapper &>(dbw).NewIterator());
    for (int c = 0; c < 2; ++c) {
        int seek_start;
        if (c == 0)
            seek_start = 0;
        else
            seek_start = 5;
        sprintf(buf, "%d", seek_start);
        StringContentsSerializer seek_key(buf);
        it->Seek(seek_key);
        for (int x = seek_start; x < 10; ++x) {
            for (int y = 0; y < 10; y++) {
                sprintf(buf, "%d", x);
                std::string exp_key(buf);
                for (int z = 0; z < y; z++)
                    exp_key += exp_key;
                StringContentsSerializer key;
                uint32_t value;
                BOOST_CHECK(it->Valid());
                // Avoid spurious errors about invalid iterator's key and value
                // in case of failure
                if (!it->Valid()) break;
                BOOST_CHECK(it->GetKey(key));
                BOOST_CHECK(it->GetValue(value));
                BOOST_CHECK_EQUAL(key.str, exp_key);
                BOOST_CHECK_EQUAL(value, x * x);
                it->Next();
            }
        }
        BOOST_CHECK(!it->Valid());
    }
}

//
// Write large number of small records (simulating coins) directly to database
// on disk, read them back and show elapsed time.
//
BOOST_AUTO_TEST_CASE(bulk_rw)
{
    // Create temporary directory and delete it when we're done.
    fs::path ph = fs::temp_directory_path() / fs::unique_path();
    ScopedPathDeleter ph_deleter(ph);

    // Create a new empty database on disk in a temporary directory
    CDBWrapper dbw(ph, (1 << 20), false, true, true); // do not use memory environment, wipe existing data and use obfuscation

    // Helper to generate random 256-bit ids so that the same seed produces same ids.
    class
    {
        std::mt19937 engine;
        std::uniform_int_distribution<unsigned int> uniform_dist{0, 255};

    public:
        void ResetSeed(decltype(engine)::result_type seed)
        {
            this->engine.seed(seed);
        }

        uint256 GenerateRandomId()
        {
            uint256 txid;
            for(auto& b: txid)
            {
                b = static_cast<std::uint8_t>(this->uniform_dist(this->engine));
            }
            return txid;
        };
    } rtg;

    using clock = std::chrono::steady_clock;

    // Number of records
    const unsigned int N = 100000;

    // Value of each record in database
    const std::vector<uint8_t> value(100, 0xde);

    // Write records
    rtg.ResetSeed(1);
    auto tp0=clock::now();
    for(unsigned int i=0; i<N; ++i)
    {
        bool res = dbw.Write(rtg.GenerateRandomId(), value);
        if(!res)
        {
            BOOST_TEST_FAIL("Error writing record "+std::to_string(i)+" to LevelDB!");
            return;
        }
    }
    BOOST_TEST_MESSAGE("Writing " + std::to_string(N) + " records to LevelDB database took: "+std::to_string(std::chrono::duration_cast<std::chrono::milliseconds>(clock::now()-tp0 ).count())+"ms" );

    // Read records
    std::vector<uint8_t> value1;
    rtg.ResetSeed(1);
    tp0=clock::now();
    for(unsigned int i=0; i<N; ++i)
    {
        bool res = dbw.Read(rtg.GenerateRandomId(), value1);
        if(!res)
        {
            BOOST_TEST_FAIL("Error reading record "+std::to_string(i)+" from LevelDB!");
            return;
        }
    }
    BOOST_TEST_MESSAGE("Reading " + std::to_string(N) + " records from LevelDB database took: "+std::to_string(std::chrono::duration_cast<std::chrono::milliseconds>(clock::now()-tp0 ).count())+"ms" );
    BOOST_TEST(value1==value); // only check value of the last record to minimize effect on timings
}

BOOST_AUTO_TEST_SUITE_END()
