node_connection.h

#pragma once

#include <utility>

#include "common/task_system/progress_bar.h"
#include "main/lbug.h"
#include "node_database.h"
#include "node_prepared_statement.h"
#include "node_progress_bar_display.h"
#include "node_query_result.h"
#include <napi.h>

using namespace lbug::main;
using namespace lbug::common;

class NodeConnection : public Napi::ObjectWrap<NodeConnection> {
    friend class ConnectionInitAsyncWorker;
    friend class ConnectionExecuteAsyncWorker;
    friend class ConnectionQueryAsyncWorker;
    friend class NodePreparedStatement;

public:
    static Napi::Object Init(Napi::Env env, Napi::Object exports);
    explicit NodeConnection(const Napi::CallbackInfo& info);
    ~NodeConnection() override = default;

private:
    Napi::Value InitAsync(const Napi::CallbackInfo& info);
    Napi::Value InitSync(const Napi::CallbackInfo& info);
    void InitCppConnection();
    void SetMaxNumThreadForExec(const Napi::CallbackInfo& info);
    void SetQueryTimeout(const Napi::CallbackInfo& info);
    Napi::Value ExecuteAsync(const Napi::CallbackInfo& info);
    Napi::Value QueryAsync(const Napi::CallbackInfo& info);
    Napi::Value QueryArrowAsync(const Napi::CallbackInfo& info);
    Napi::Value ExecuteSync(const Napi::CallbackInfo& info);
    Napi::Value QuerySync(const Napi::CallbackInfo& info);
    Napi::Value QueryArrowSync(const Napi::CallbackInfo& info);
    Napi::Value CreateArrowTableSync(const Napi::CallbackInfo& info);
    Napi::Value CreateArrowRelTableSync(const Napi::CallbackInfo& info);
    Napi::Value DropArrowTableSync(const Napi::CallbackInfo& info);
    void Close(const Napi::CallbackInfo& info);

private:
    std::shared_ptr<Database> database;
    std::shared_ptr<Connection> connection;
};

class ConnectionInitAsyncWorker : public Napi::AsyncWorker {
public:
    ConnectionInitAsyncWorker(Napi::Function& callback, NodeConnection* nodeConnection)
        : Napi::AsyncWorker(callback), nodeConnection(nodeConnection) {}

    ~ConnectionInitAsyncWorker() override = default;

    void Execute() override {
        try {
            nodeConnection->InitCppConnection();
        } catch (const std::exception& exc) {
            SetError(std::string(exc.what()));
        }
    }

    void OnOK() override { Callback().Call({Env().Null()}); }

    void OnError(Napi::Error const& error) override { Callback().Call({error.Value()}); }

private:
    NodeConnection* nodeConnection;
};

namespace lbug {
namespace main {

class ConnectionExecuteAsyncWorker : public Napi::AsyncWorker {
public:
    ConnectionExecuteAsyncWorker(Napi::Function& callback, std::shared_ptr<Connection>& connection,
        std::shared_ptr<Database>& database,
        std::shared_ptr<PreparedStatement> preparedStatement, NodeQueryResult* nodeQueryResult,
        std::unordered_map<std::string, std::unique_ptr<Value>> params,
        Napi::Value progressCallback)
        : Napi::AsyncWorker(callback), connection(connection), database(database),
          preparedStatement(std::move(preparedStatement)), nodeQueryResult(nodeQueryResult),
          params(std::move(params)) {
        if (progressCallback.IsFunction()) {
            this->progressCallback = Napi::ThreadSafeFunction::New(Env(),
                progressCallback.As<Napi::Function>(), "ProgressCallback", 0, 1);
        }
    }

    ~ConnectionExecuteAsyncWorker() override = default;

    void Execute() override {
        auto clientContext = connection->getClientContext();
        uint64_t queryID = clientContext->getDatabase()->getNextQueryID();
        auto progressBar = common::ProgressBar::Get(*clientContext);
        auto trackProgress = progressBar->getProgressBarPrinting();
        auto display = progressBar->getDisplay().get();
        NodeProgressBarDisplay* nodeDisplay =
            dynamic_cast_checked<NodeProgressBarDisplay*>(display);
        if (progressCallback) {
            nodeDisplay->setCallbackFunction(queryID, *progressCallback);
            progressBar->toggleProgressBarPrinting(true);
        }
        try {
            auto result =
                connection
                    ->executeWithParamsWithID(preparedStatement.get(), std::move(params), queryID);
            if (!result->isSuccess()) {
                SetError(result->getErrorMessage());
                return;
            }
            nodeQueryResult->AdoptQueryResult(std::move(result), connection, database);
        } catch (const std::exception& exc) {
            SetError(std::string(exc.what()));
        }
        if (progressCallback) {
            if (nodeDisplay->getNumCallbacks() == 0) {
                progressBar->toggleProgressBarPrinting(trackProgress);
            }
        }
    }

    void OnOK() override { Callback().Call({Env().Null()}); }

    void OnError(Napi::Error const& error) override { Callback().Call({error.Value()}); }

private:
    std::shared_ptr<Connection> connection;
    std::shared_ptr<Database> database;
    std::shared_ptr<PreparedStatement> preparedStatement;
    NodeQueryResult* nodeQueryResult;
    std::unordered_map<std::string, std::unique_ptr<Value>> params;
    std::optional<Napi::ThreadSafeFunction> progressCallback;
};

class ConnectionQueryAsyncWorker : public Napi::AsyncWorker {
public:
    ConnectionQueryAsyncWorker(Napi::Function& callback, std::shared_ptr<Connection>& connection,
        std::shared_ptr<Database>& database, std::string statement,
        NodeQueryResult* nodeQueryResult, Napi::Value progressCallback)
        : Napi::AsyncWorker(callback), connection(connection), database(database),
          statement(std::move(statement)), nodeQueryResult(nodeQueryResult) {
        if (progressCallback.IsFunction()) {
            this->progressCallback = Napi::ThreadSafeFunction::New(Env(),
                progressCallback.As<Napi::Function>(), "ProgressCallback", 0, 1);
        }
    }

    ~ConnectionQueryAsyncWorker() override = default;

    void Execute() override {
        auto clientContext = connection->getClientContext();
        uint64_t queryID = clientContext->getDatabase()->getNextQueryID();
        auto progressBar = common::ProgressBar::Get(*clientContext);
        auto trackProgress = progressBar->getProgressBarPrinting();
        auto display = progressBar->getDisplay().get();
        NodeProgressBarDisplay* nodeDisplay =
            dynamic_cast_checked<NodeProgressBarDisplay*>(display);
        if (progressCallback) {
            nodeDisplay->setCallbackFunction(queryID, *progressCallback);
            progressBar->toggleProgressBarPrinting(true);
        }
        try {
            auto result = connection->queryWithID(statement, queryID);
            if (!result->isSuccess()) {
                SetError(result->getErrorMessage());
                return;
            }
            nodeQueryResult->AdoptQueryResult(std::move(result), connection, database);
        } catch (const std::exception& exc) {
            SetError(std::string(exc.what()));
        }
        if (progressCallback) {
            if (nodeDisplay->getNumCallbacks() == 0) {
                progressBar->toggleProgressBarPrinting(trackProgress);
            }
        }
    }

    void OnOK() override { Callback().Call({Env().Null()}); }

    void OnError(Napi::Error const& error) override { Callback().Call({error.Value()}); }

private:
    std::shared_ptr<Connection> connection;
    std::shared_ptr<Database> database;
    std::string statement;
    NodeQueryResult* nodeQueryResult;
    std::optional<Napi::ThreadSafeFunction> progressCallback;
};

class ConnectionQueryArrowAsyncWorker : public Napi::AsyncWorker {
public:
    ConnectionQueryArrowAsyncWorker(Napi::Function& callback,
        std::shared_ptr<Connection>& connection, std::shared_ptr<Database>& database,
        std::string statement, int64_t chunkSize, NodeQueryResult* nodeQueryResult)
        : Napi::AsyncWorker(callback), connection(connection), database(database),
          statement(std::move(statement)), chunkSize(chunkSize), nodeQueryResult(nodeQueryResult) {}

    ~ConnectionQueryArrowAsyncWorker() override = default;

    void Execute() override {
        try {
            auto result = connection->queryAsArrow(statement, chunkSize);
            if (!result->isSuccess()) {
                SetError(result->getErrorMessage());
                return;
            }
            nodeQueryResult->AdoptQueryResult(std::move(result), connection, database);
        } catch (const std::exception& exc) {
            SetError(std::string(exc.what()));
        }
    }

    void OnOK() override { Callback().Call({Env().Null()}); }

    void OnError(Napi::Error const& error) override { Callback().Call({error.Value()}); }

private:
    std::shared_ptr<Connection> connection;
    std::shared_ptr<Database> database;
    std::string statement;
    int64_t chunkSize;
    NodeQueryResult* nodeQueryResult;
};

} // namespace main
} // namespace lbug