#include <iostream>
#include <string>
#include <vector>
#include <unordered_map>
#include <thread>
#include <mutex>
#include <queue>
#include <variant>
#include <sqlite3.h>
#include <dlang.hpp>

#if defined(_WIN32) || defined(_WIN64)
#define EXPORT __declspec(dllexport)
#else
#define EXPORT __attribute__((visibility("default")))
#endif

using SqlValue = std::variant<std::nullptr_t, double, std::string>;
using SqlRow = std::unordered_map<std::string, SqlValue>;
using SqlQueryResult = std::vector<SqlRow>;

struct DbResult {
    bool success;
    std::string errorMessage;
    SqlQueryResult rows;
    int rowsAffected;
};

struct AsyncTask {
    StackValue callbackVal;
    DbResult result;
};

static sqlite3* g_db = nullptr;
static std::queue<AsyncTask> g_completedTasks;
static std::mutex g_queueMutex;

static DbResult ExecuteQueryInternal(sqlite3* db, const std::string& query) {
    DbResult res{ true, "", {}, 0 };
    sqlite3_stmt* stmt = nullptr;

    if (sqlite3_prepare_v2(db, query.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        res.success = false;
        res.errorMessage = sqlite3_errmsg(db);
        return res;
    }

    int colCount = sqlite3_column_count(stmt);

    while (true) {
        int stepRes = sqlite3_step(stmt);
        if (stepRes == SQLITE_ROW) {
            SqlRow row;
            for (int i = 0; i < colCount; ++i) {
                const char* colName = sqlite3_column_name(stmt, i);
                int type = sqlite3_column_type(stmt, i);

                switch (type) {
                case SQLITE_INTEGER:
                    row[colName] = static_cast<double>(sqlite3_column_int64(stmt, i));
                    break;
                case SQLITE_FLOAT:
                    row[colName] = sqlite3_column_double(stmt, i);
                    break;
                case SQLITE_TEXT:
                    row[colName] = std::string(reinterpret_cast<const char*>(sqlite3_column_text(stmt, i)));
                    break;
                case SQLITE_NULL:
                default:
                    row[colName] = nullptr;
                    break;
                }
            }
            res.rows.push_back(row);
        }
        else if (stepRes == SQLITE_DONE) {
            break;
        }
        else {
            res.success = false;
            res.errorMessage = sqlite3_errmsg(db);
            break;
        }
    }

    res.rowsAffected = sqlite3_changes(db);
    sqlite3_finalize(stmt);
    return res;
}

static void PushDbResultToVM(DlangVM* vm, const DbResult& res) {
    dlang_newtable(vm);

    dlang_pushboolean(vm, res.success);
    dlang_setfield(vm, "success");

    if (!res.success) {
        dlang_pushstring(vm, res.errorMessage.c_str());
        dlang_setfield(vm, "error");
    }

    dlang_pushnumber(vm, res.rowsAffected);
    dlang_setfield(vm, "rowsAffected");

    dlang_newlist(vm);

    for (const auto& row : res.rows) {
        dlang_newtable(vm);
        
        for (const auto& [colName, val] : row) {
            if (std::holds_alternative<double>(val)) {
                dlang_pushnumber(vm, std::get<double>(val));
            }
            else if (std::holds_alternative<std::string>(val)) {
                dlang_pushstring(vm, std::get<std::string>(val).c_str());
            }
            else {
                dlang_pushnil(vm);
            }
            dlang_setfield(vm, colName.c_str());
        }

        dlang_push(vm, dlang_pop(vm));
        dlang_pushindex(vm);
    }

    dlang_setfield(vm, "rows");
}

void NativeSqliteOpen(DlangVM* vm) {
    if (!dlang_isstring(vm)) {
        dlang_pushboolean(vm, false);
        return;
    }

    std::string dbPath = dlang_tostring(vm);
    if (g_db) {
        sqlite3_close(g_db);
        g_db = nullptr;
    }

    int rc = sqlite3_open(dbPath.c_str(), &g_db);
    if (rc == SQLITE_OK) {
        sqlite3_exec(g_db, "PRAGMA journal_mode = WAL;", nullptr, nullptr, nullptr);
        sqlite3_exec(g_db, "PRAGMA synchronous = NORMAL;", nullptr, nullptr, nullptr);
        dlang_pushboolean(vm, true);
    }
    else {
        dlang_pushboolean(vm, false);
    }
}

void NativeSqliteQuery(DlangVM* vm) {
    if (!g_db || !dlang_isstring(vm)) {
        DbResult err{ false, "Geen DB of ongeldige query", {}, 0 };
        PushDbResultToVM(vm, err);
        return;
    }

    std::string query = dlang_tostring(vm);
    DbResult res = ExecuteQueryInternal(g_db, query);
    PushDbResultToVM(vm, res);
}

void NativeSqliteQueryAsync(DlangVM* vm) {
    if (!g_db) return;

    if (!dlang_isstring(vm)) return;
    std::string query = dlang_tostring(vm);

    if (!dlang_isfunction(vm)) return;
    StackValue callbackVal = dlang_totoken(vm);

    std::thread([query, callbackVal]() {
        DbResult res = ExecuteQueryInternal(g_db, query);

        std::lock_guard<std::mutex> lock(g_queueMutex);
        g_completedTasks.push({ callbackVal, res });
    }).detach();
}

void NativeSqlitePoll(DlangVM* vm) {
    std::queue<AsyncTask> localTasks;

    {
        std::lock_guard<std::mutex> lock(g_queueMutex);
        localTasks.swap(g_completedTasks);
    }

    while (!localTasks.empty()) {
        auto task = localTasks.front();
        localTasks.pop();

        PushDbResultToVM(vm, task.result);
        dlang_push(vm, task.callbackVal);
        dlang_call(vm, 1);
    }
}

void NativeSqliteClose(DlangVM* vm) {
    if (g_db) {
        sqlite3_close(g_db);
        g_db = nullptr;
    }
    dlang_pushboolean(vm, true);
}

extern "C" {
    EXPORT void initialize(DlangVM* vm) {
        dlang_newtable(vm);

        dlang_pushcfunction(vm, NativeSqliteOpen);
        dlang_setfield(vm, "open");

        dlang_pushcfunction(vm, NativeSqliteQuery);
        dlang_setfield(vm, "query");

        dlang_pushcfunction(vm, NativeSqliteQueryAsync);
        dlang_setfield(vm, "query_async");

        dlang_pushcfunction(vm, NativeSqlitePoll);
        dlang_setfield(vm, "poll");

        dlang_pushcfunction(vm, NativeSqliteClose);
        dlang_setfield(vm, "close");

        dlang_setglobal(vm, "sqlite");

        std::cout << ">>> SQLite Module Initialized (Queue-based Async) <<<" << std::endl;
    }
}