#pragma once

#include <unordered_map>

#include "main/storage_driver.h"

#include "py_database.h"

#include "py_prepared_statement.h"

#include "py_query_result.h"

using lbug::common::LogicalType;

using lbug::common::LogicalTypeID;

using lbug::common::Value;

class PyConnection {

public:

static void initialize(py::handle& m);

explicit PyConnection(PyDatabase* pyDatabase, uint64_t numThreads);

void close();

~PyConnection() = default;

void setQueryTimeout(uint64_t timeoutInMS);

void interrupt();

std::unique_ptr<PyQueryResult> execute(PyPreparedStatement* preparedStatement,

const py::dict& params);

std::unique_ptr<PyQueryResult> query(const std::string& statement);

void setMaxNumThreadForExec(uint64_t numThreads);

PyPreparedStatement prepare(const std::string& query, const py::dict& parameters);

uint64_t getNumNodes(const std::string& nodeName);

uint64_t getNumRels(const std::string& relName);

void getAllEdgesForTorchGeometric(py::array_t<int64_t>& npArray,

const std::string& srcTableName, const std::string& relName,

const std::string& dstTableName, size_t queryBatchSize);

static bool isPandasDataframe(const py::handle& object);

static bool isPolarsDataframe(const py::handle& object);

static bool isPyArrowTable(const py::handle& object);

void createScalarFunction(const std::string& name, const py::function& udf,

const py::list& params, const std::string& retval, bool defaultNull, bool catchExceptions);

void removeScalarFunction(const std::string& name);

std::unique_ptr<PyQueryResult> createArrowTable(const std::string& tableName,

py::object arrowTable);

std::unique_ptr<PyQueryResult> createArrowRelTable(const std::string& tableName,

py::object arrowTable, const std::string& srcTableName, const std::string& dstTableName);

std::unique_ptr<PyQueryResult> dropArrowTable(const std::string& tableName);

static Value transformPythonValue(const py::handle& val);

static Value transformPythonValueAs(const py::handle& val, const LogicalType& type);

static Value transformPythonValueFromParameter(const py::handle& val);

static Value transformPythonValueFromParameterAs(const py::handle& val,

const LogicalType& type);

private:

std::unique_ptr<StorageDriver> storageDriver;

std::unique_ptr<Connection> conn;

std::unordered_map<std::string, py::object> arrowTableRefs;

static std::unique_ptr<PyQueryResult> checkAndWrapQueryResult(

std::unique_ptr<QueryResult>& queryResult);

};