bool createDatabase(TypeDB::Driver& driver, const std::string& dbName) {
    std::cout << "Creating a new database...";
    driver.databases.create(dbName);
    std::cout << "OK" << std::endl;
    TypeDB::Options options;
    {
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::SCHEMA, options);
        dbSchemaSetup(session);
    }
    {
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        dbDatasetSetup(session);
    }
    return true;
}

bool replaceDatabase(TypeDB::Driver& driver, const std::string& dbName) {
    std::cout << "Deleting an existing database...";
    driver.databases.get(dbName).deleteDatabase(); // Delete the database if it exists already
    std::cout << "OK" << std::endl;
    if (!createDatabase(driver, dbName)) {
        std::cout << "Failed to create a new database. Terminating..." << std::endl;
        exit(EXIT_FAILURE);
    }
    return true;
}

void dbSchemaSetup(TypeDB::Session& schemaSession, const std::string& schemaFile = "iam-schema.tql") {
    std::string defineQuery;
    std::ifstream newfile;
    newfile.open(schemaFile, std::ios::in);
    if (newfile.is_open()){
        std::string tmp;
        while(getline(newfile, tmp)){
            defineQuery = defineQuery + tmp + "\n";
        }
        newfile.close();
    } else {
        std::cerr << "Failed to open a file. Terminating..." << std::endl;
        exit(EXIT_FAILURE);
    }
    TypeDB::Options options;
    TypeDB::Transaction tx = schemaSession.transaction(TypeDB::TransactionType::WRITE, options);
    std::cout << "Defining schema...";
    tx.query.define(defineQuery).get();
    tx.commit();
    std::cout << "OK" << std::endl;
}

void dbDatasetSetup(TypeDB::Session& dataSession, const std::string& dataFile = "iam-data-single-query.tql") {
    std::ifstream file(dataFile);
    std::string insertQuery((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    TypeDB::Options options;
    TypeDB::Transaction tx = dataSession.transaction(TypeDB::TransactionType::WRITE, options);
    std::cout << "Loading data...";
    TypeDB::ConceptMapIterable response = tx.query.insert(insertQuery);
    int16_t count = 0;
    for (TypeDB::ConceptMap& conceptMap : response) { count += 1; }
    tx.commit();
    std::cout << "OK" << std::endl;
}

bool dbCheck(TypeDB::Session& dataSession) {
    TypeDB::Options options;
    TypeDB::Transaction tx = dataSession.transaction(TypeDB::TransactionType::READ, options);
    std::string testQuery = "match $u isa user; get $u; count;";
    std::cout << "Testing the database...";
    TypeDB::AggregateFuture response = tx.query.getAggregate(testQuery);
    int16_t result = response.get().value().get()->asLong();
    if (result == 3) {
        std::cout << "Passed" << std::endl;
        return true;
    } else {
        std::cout << "Failed with the result: " << result << "\nExpected result: 3." << std::endl;
        return false;
    }
}

std::vector<TypeDB::JSON> fetchAllUsers(TypeDB::Driver& driver, const std::string& dbName) {
    std::vector<TypeDB::JSON> users;
    TypeDB::Options options;
    {
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        TypeDB::Transaction tx = session.transaction(TypeDB::TransactionType::READ, options);
        TypeDB::JSONIterable queryResult = tx.query.fetch("match $u isa user; fetch $u: full-name, email;");
        int16_t c = 1;
        for (TypeDB::JSON user : queryResult) {
                users.push_back(user);
                std::cout << "User #" << c++ << " ";
                printJSON(user);
                std::cout << std::endl;
        }
    }
    return users;
}

void printJSON(TypeDB::JSON json) {
    if (json.isString()) {
        std::cout << "'" << json.asString() << "'" << std::endl;
    }
    if (json.isMap()) {
        TypeDB::JSONMap jsonMap = json.asMap();
        for (const auto& p : jsonMap ) {
                std::cout << p.first << ": "<< std::endl;
                printJSON(p.second);
            }
    }
}

std::vector<TypeDB::ConceptMap> insertNewUser(TypeDB::Driver& driver, const std::string& dbName, const std::string& name, const std::string& email) {
    std::vector<TypeDB::ConceptMap> response;
    TypeDB::Options options;
    {
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        TypeDB::Transaction tx = session.transaction(TypeDB::TransactionType::WRITE, options);
        TypeDB::ConceptMapIterable response = tx.query.insert("insert $p isa person, has full-name $fn, has email $e; $fn == '" + name + "'; $e == '" + email + "';");
        for (TypeDB::ConceptMap& conceptMap : response) {
            std::string name = conceptMap.get("fn")->asAttribute()->getValue()->asString();
            std::string email = conceptMap.get("e")->asAttribute()->getValue()->asString();
            std::cout << "Added new user. Name: " << name << ", E-mail: " << email << std::endl;
        }
        tx.commit();
    }
    return response;
}

std::vector<std::string> getFilesByUser(TypeDB::Driver& driver, const std::string& dbName, const std::string& name, bool inference = false) {
    TypeDB::Options options;
    options.infer(inference);
    {
        std::vector<std::string> files;
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        TypeDB::Transaction tx = session.transaction(TypeDB::TransactionType::READ, options);
        TypeDB::ConceptMapIterable users = tx.query.get("match $u isa user, has full-name '" + name + "'; get;");
        int16_t userCount = 0;
        for (TypeDB::ConceptMap& user : users) {userCount += 1;}
        if (userCount > 1) {
            std::cout << "Error: Found more than one user with that name." << std::endl;
            return files;
        } else if (userCount == 1) {
            TypeDB::ConceptMapIterable response = tx.query.get(R"(
                                                                match
                                                                $fn == ')" + name + R"(';
                                                                $u isa user, has full-name $fn;
                                                                $p($u, $pa) isa permission;
                                                                $o isa object, has path $fp;
                                                                $pa($o, $va) isa access;
                                                                $va isa action, has name 'view_file';
                                                                get $fp; sort $fp asc;
                                                                )");
            int16_t resultCounter = 0;
            for (TypeDB::ConceptMap& cm : response) {
                resultCounter += 1;
                files.push_back(cm.get("fp") ->asAttribute()->getValue()->asString());
                std::cout << "File #" << std::to_string(resultCounter) << ": " << cm.get("fp") ->asAttribute()->getValue()->asString() << std::endl;
            }
            if (resultCounter == 0) {
                std::cout << "No files found. Try enabling inference." << std::endl;
            }
            return files;
        } else {
            std::cout << "Error: No users found with that name." << std::endl;
            return files;
        }
    }
}

int16_t updateFilePath(TypeDB::Driver& driver, const std::string& dbName, const std::string& oldPath, const std::string& newPath) {
    std::vector<TypeDB::ConceptMap> response;
    int16_t count = 0;
    {
        TypeDB::Options options;
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        TypeDB::Transaction tx = session.transaction(TypeDB::TransactionType::WRITE, options);
        TypeDB::ConceptMapIterable response = tx.query.update(R"(
                                                                match
                                                                $f isa file, has path $old_path;
                                                                $old_path = ')" + oldPath + R"(';
                                                                delete
                                                                $f has $old_path;
                                                                insert
                                                                $f has path $new_path;
                                                                $new_path = ')" + newPath + R"(';
                                                                )");
        for (TypeDB::ConceptMap& conceptMap : response) { count += 1; }
        if (count > 0) {
            tx.commit();
            std::cout << "Total number of paths updated: " << count << "." << std::endl;
        } else {
            std::cout << "No matched paths: nothing to update." << std::endl;
        }
    }
    return count;
}

bool deleteFile(TypeDB::Driver& driver, const std::string& dbName, const std::string& path) {
    {
        TypeDB::Options options;
        TypeDB::Session session = driver.session(dbName, TypeDB::SessionType::DATA, options);
        TypeDB::Transaction tx = session.transaction(TypeDB::TransactionType::WRITE, options);
        TypeDB::ConceptMapIterable response = tx.query.get(R"(
                                                            match
                                                            $f isa file, has path ')" + path + R"(';
                                                            get;
                                                            )");
        int16_t count = 0;
        for (TypeDB::ConceptMap& conceptMap : response) { count += 1; }
        if (count == 1) {
            tx.query.matchDelete(R"(
                                match
                                $f isa file, has path ')" + path + R"(';
                                delete
                                $f isa file;
                                )").get();
            tx.commit();
            std::cout << "The file has been deleted." << std::endl;
            return true;
        } else if (count > 1) {
            std::cout << "Matched more than one file with the same path." << std::endl;
            std::cout << "No files were deleted." << std::endl;
            return false;
        } else {
            std::cout << "No files matched in the database." << std::endl;
            std::cout << "No files were deleted." << std::endl;
            return false;
        }
    }
}