/**
 * Copyright (c) 2009-2025, quasardb SAS. All rights reserved.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *    * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *    * Neither the name of quasardb nor the names of its contributors may
 *      be used to endorse or promote products derived from this software
 *      without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY QUASARDB AND CONTRIBUTORS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <qdb/client.hpp>
#include <qdb/query.hpp>
#include <qdb/ts.hpp>
#include <algorithm>
#include <array>
#include <chrono>
#include <iostream>
#include <locale>
#include <string>
#include <tuple>

namespace
{

static constexpr unsigned long long batch_size = 1'000'000;

[[nodiscard]] constexpr unsigned long long align_to_batch_size(unsigned long long v) noexcept
{
    static_assert(batch_size > 0u, "batch size cannot be zero");
    return ((v + batch_size - 1u) / batch_size) * batch_size;
}

// the columns of our table
static constexpr qdb_size_t col_count = 2;
static const std::array<qdb_exp_batch_push_column_schema_t, col_count> table_columns{
    qdb_exp_batch_push_column_schema_t{qdb_ts_column_double, 0, nullptr, "price"}, //
    qdb_exp_batch_push_column_schema_t{qdb_ts_column_int64, 1, nullptr, "volume"}  //
};

// schema of our table
static const std::array<qdb_exp_batch_push_table_schema_t, 1> table_schema{
    qdb_exp_batch_push_table_schema_t{qdb_d_hour, qdb_ttl_disabled, table_columns.data(), table_columns.size()} //
};

[[nodiscard]] inline qdb_timespec_t operator+(qdb_timespec_t ts, qdb_time_t ns_increment) noexcept
{
    qdb_timespec_t result = ts;
    result.tv_nsec += ns_increment;

    if (result.tv_nsec >= qdb_time_t{1'000'000'000})
    {
        // we assume we started from 0 nsec
        result.tv_nsec = 0;
        result.tv_sec += 1;
    }

    return result;
}

// in order to get thousand delimiters in the output we set the locale
void set_compatible_locale()
{
    static const std::vector<std::string> compatible_locales = {
#ifdef _MSC_VER
        "en-US",
#endif
        "en_US.UTF-8", //
        "en_GB.UTF-8", //
        "en",          //
        "C.UTF-8",     //
        "C"};

    if (std::none_of(compatible_locales.cbegin(), compatible_locales.cend(), [](const auto & loc) {
            const auto res = std::setlocale(LC_ALL, loc.data());
            if (res)
            {
                std::locale::global(std::locale(loc.data()));
                std::cout.imbue(std::locale(loc));
            }
            return res;
        }))
    {
        std::cerr << "cannot set any compatible locale" << std::endl;
    }
}

void drop_table(qdb::handle & h, std::string_view table_name)
{
    // ignore errors for this example
    // you can also do a DROP TABLE query
    qdb_remove(h, table_name.data());
}

void create_table(qdb::handle & h, std::string_view table_name)
{
    qdb_query_result_t * res = nullptr;

    const auto query_str = "CREATE TABLE " + std::string{table_name} + " (price DOUBLE, volume INT64) SHARD_SIZE=1h";

    // table has been dropped before
    const auto err = qdb_query(h, query_str.data(), &res);
    if (QDB_FAILURE(err))
    {
        throw err;
    }
    if (res)
    {
        qdb_release(h, res);
    }
}

// placeholder to keep alive the generated data while sending it
struct generated_data
{
    std::array<qdb_exp_batch_push_column_t, 2> columns;

    std::vector<qdb_timespec_t> timestamps;
    std::vector<double> prices;
    std::vector<qdb_int_t> volumes;
};

[[nodiscard]] qdb_exp_batch_options_t make_options() noexcept
{
    qdb_exp_batch_options_t options;
    // qdb_exp_batch_push_async is the fastest but will delay data visibility
    // if you need immediate data visibility, use qdb_exp_batch_push_fast
    // qdb_exp_batch_push_transactional is typically 2 to 3 times slower than fast and can create transaction conflicts
    options.mode = qdb_exp_batch_push_fast;
    // informs the server to not keep the newly written data in cache, which is usually what you want
    options.push_flags = qdb_exp_batch_push_flag_write_through;

    return options;
}

// makes the structure without the data
[[nodiscard]] qdb_exp_batch_push_table_t make_batch_push_table(std::string_view table_name) noexcept
{
    qdb_exp_batch_push_table_t table_data{};

    table_data.name = table_name.data();
    // no truncation
    table_data.truncate_ranges      = nullptr;
    table_data.truncate_range_count = 0;
    // no deduplication
    table_data.deduplication_mode    = qdb_exp_batch_deduplication_mode_disabled;
    table_data.where_duplicate       = nullptr;
    table_data.where_duplicate_count = 0;
    // we already created the table, it's faster to create the table before than within the batch
    table_data.creation = qdb_exp_batch_dont_create;

    return table_data;
}

// we keep data generation very simple, otherwise it's going to be the bottleneck
// random data generation is slower than you think...
// feel free to replace with real world data for your tests
void generate_data(generated_data & context, qdb_timespec_t start_point, unsigned long long count)
{
    // one point every 20 millisecond for this example, that's 180,000 points per hour
    static constexpr qdb_time_t time_increment = 20'000'000; // 20 milliseconds

    static_assert(time_increment > 0u, "time increment cannot be zero");

    context.timestamps.resize(count);
    context.prices.resize(count);
    context.volumes.resize(count);

    std::generate(context.timestamps.begin(), context.timestamps.end(), [&start_point]() {
        start_point.tv_nsec += time_increment;
        if (start_point.tv_nsec >= qdb_time_t{1'000'000'000})
        {
            // it's not super accurate, but we don't really need great accuracy in the timestamp generation, just a good overal density
            ++start_point.tv_sec;
            start_point.tv_nsec = 0;
        }

        return start_point;
    });

    std::transform(context.timestamps.cbegin(), context.timestamps.cend(), context.prices.begin(), [](const auto & ts) {
        // we just need values for our insertion tests
        return static_cast<double>(ts.tv_sec) + static_cast<double>(ts.tv_nsec);
    });

    std::transform(context.timestamps.cbegin(), context.timestamps.cend(), context.volumes.begin(), [](const auto & ts) {
        // we just need values for our insertion tests
        return static_cast<qdb_int_t>(ts.tv_sec) ^ static_cast<qdb_int_t>(ts.tv_nsec);
    });

    // update the structure
    context.columns[0].data.doubles = context.prices.data();
    context.columns[0].data_type    = qdb_ts_column_double;
    context.columns[0].name         = "price";

    context.columns[1].data.ints = context.volumes.data();
    context.columns[1].data_type = qdb_ts_column_int64;
    context.columns[1].name      = "volume";
}

void generate_data(generated_data & data, unsigned long long rows_count)
{
    // start at January 2nd 2025 - 18:00 UTC - arbitrary
    static constexpr qdb_timespec_t start_point{1735840800, 0};

    generate_data(data, start_point, rows_count);
}

void insert_rows(qdb::handle & h, std::string_view table_name, const generated_data & data)
{
    auto options = make_options();

    auto table_data = make_batch_push_table(table_name);

    // you can specify nullptr, in which case the schema will be retrieved from the server and the structure will be filled with an
    // allocated structure you will need to call qdb_release on
    // that allocated structure will contain the correct schema information
    // you can either
    // - Provide it statically as in this example
    // - Retrieve it from the server on the first call, and pass it to the next calls
    const qdb_exp_batch_push_table_schema_t * table_schemas = table_schema.data();

    table_data.data.row_count    = data.prices.size();
    table_data.data.timestamps   = data.timestamps.data();
    table_data.data.columns      = data.columns.data();
    table_data.data.column_count = data.columns.size();

    // the batch push will break down your large data sets into smaller pieces and send that to the server
    const auto err = qdb_exp_batch_push_with_options(h, &options, &table_data, &table_schemas, 1);
    if (QDB_FAILURE(err))
    {
        throw err;
    }
}

} // namespace

struct timer
{
    timer() noexcept
        : _start{std::chrono::high_resolution_clock::now()}
    {}

    [[nodiscard]] std::chrono::microseconds elapsed() const noexcept
    {
        return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::high_resolution_clock::now() - _start);
    }

private:
    std::chrono::high_resolution_clock::time_point _start;
};

int main(int argc, char ** argv)
{
    try
    {
        static constexpr int expected_args = 4;

        if (argc < expected_args)
        {
            std::cerr << "usage: " << argv[0] << " qdb_url table_name rows_count_in_millions [threads]" << std::endl;
            return EXIT_FAILURE;
        }

        set_compatible_locale();

        // this assumes a cluster with authentication disabled, you will need to set your credentials if authentication is on
        qdb::handle h;

        std::cout << "connecting to " << argv[1] << "..." << std::endl;

        // the default is half std::thread::hardware_concurrency();
        // here we set the default to one thread max
        qdb_size_t threads = 1;

        if (argc > 4)
        {
            threads = std::stoul(argv[4]);
        }

        // parallelism must be configured before connecting, data will be sent in parallel using at most the given number of threads
        // tasks are split on shard boundaries, meaning that if you only insert in a single shard, parallelism brings no improvement
        auto err = h.set_client_max_parallelism(threads);
        if (QDB_FAILURE(err))
        {
            throw err;
        }

        // recommended protocol compression setting starting with 3.14.2, big impact when inserting data
        // set it before connecting
        // for backward compatibility reason it's set by default to qdb_comp_none
        err = h.set_compression(qdb_comp_balanced);
        if (QDB_FAILURE(err))
        {
            throw err;
        }

        // here we connect to the cluster
        err = h.connect(argv[1]);
        if (QDB_FAILURE(err))
        {
            throw err;
        }

        // configures how tasks should be batched, a value of 1 means tasks will be sent in batches of 1
        // the default value of 23 is a good value for complex SELECT queries, for batch insert, a value of 5 can help speed up if your
        // batches are relatively small you can even go down to 1 if you think your inserts will typically span only a couple of shards
        err = qdb_option_set_client_max_batch_load(h, 5);
        if (QDB_FAILURE(err))
        {
            throw err;
        }

        const auto table_name = std::string_view{argv[2]};

        {
            timer t;

            std::cout << "dropping table table " << table_name << "..." << std::endl;
            drop_table(h, table_name);

            std::cout << "table dropped in " << t.elapsed().count() << " us" << std::endl;
        }

        {
            timer t;

            std::cout << "creating table " << table_name << "..." << std::endl;
            create_table(h, table_name);

            std::cout << "table created in " << t.elapsed().count() << " us" << std::endl;
        }

        const auto rows_count = align_to_batch_size(std::stoull(argv[3]) * 1'000'000);

        generated_data data;

        {
            timer t;
            std::cout << "generating " << rows_count << " rows..." << std::endl;
            generate_data(data, rows_count);
            std::cout << "rows generated in " << t.elapsed().count() << " us" << std::endl;
        }

        {
            timer t;

            std::cout << "inserting " << rows_count << " rows using " << threads << " thread(s)..." << std::endl;
            insert_rows(h, table_name, data);

            const auto elapsed      = std::max<unsigned long long>(1u, static_cast<unsigned long long>(t.elapsed().count()));
            const auto rows_per_sec = (rows_count * 1'000'000u) / elapsed;

            std::cout << "inserted " << rows_count << " rows in " << elapsed << " us - " << rows_per_sec << " rows/s" << std::endl;
        }

        return EXIT_SUCCESS;
    }
    catch (qdb_error_t err)
    {
        std::cerr << "quasardb error: " << qdb_error(err) << std::endl;
        return EXIT_FAILURE;
    }
    catch (const std::exception & e)
    {
        std::cerr << "exception caught: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }
}
