Untitled

package com.trexmine.bedwars.database.impl.mongodb;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.mongodb.MongoClientSettings;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import com.mongodb.client.*;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Indexes;
import com.mongodb.client.model.UpdateOptions;
import com.mongodb.client.model.Updates;
import com.trexmine.bedwars.constant.holder.HotbarItemHolder;
import com.trexmine.bedwars.constant.holder.StatsRowHolder;
import com.trexmine.bedwars.database.Database;
import com.trexmine.bedwars.database.content.HotBarData;
import com.trexmine.bedwars.database.content.QuickbuyData;
import com.trexmine.bedwars.database.content.StatisticsData;
import com.trexmine.bedwars.manager.party.PrivateGameOptions;
import com.trexmine.bedwars.menu.ItemType;
import com.trexmine.bedwars.utils.functional.LogUtils;
import lombok.RequiredArgsConstructor;
import org.bson.Document;
import org.bson.UuidRepresentation;
import org.bson.conversions.Bson;

import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.logging.Level;
import java.util.logging.Logger;

@RequiredArgsConstructor
public final class MongoDB implements Database {

    private static final String STATS_COLLECTION = "stats";
    private static final String QUICKBUY_COLLECTION = "quickbuy";
    private static final String HOTBAR_COLLECTION = "hotbar";
    private static final String PRIVATEGAME_COLLECTION = "privategame";

    static {
        Logger.getLogger("org.mongodb.driver").setLevel(Level.SEVERE);
    }

    private final String database;
    private final MongoClientSettings settings;
    private MongoClient client;

    public MongoDB(String host, int port, String database, String databaseAuth, String username, String password) {
        this.database = database;
        this.settings = MongoClientSettings.builder()
                .applyToClusterSettings(builder -> builder.hosts(Collections.singletonList(new ServerAddress(host, port))))
                .credential(MongoCredential.createCredential(username, databaseAuth, password.toCharArray()))
                .uuidRepresentation(UuidRepresentation.STANDARD)
                .build();
    }

    @Override
    public void connect() {
        try {
            client = MongoClients.create(settings);
            Set<String> collections = Sets.newHashSet(
                    STATS_COLLECTION,
                    QUICKBUY_COLLECTION,
                    HOTBAR_COLLECTION,
                    PRIVATEGAME_COLLECTION
            );

            MongoDatabase db = client.getDatabase(database);
            db.listCollectionNames().forEach(collections::remove);

            for (String name : collections) {
                db.createCollection(name);

                if (name.equals(STATS_COLLECTION)) {
                    db.getCollection(name).createIndex(Indexes.text("username"));
                }
            }
        } catch (Exception e) {
            LogUtils.warn("Couldn't connect to Mongo database!");
        }
    }

    @Override
    public void disconnect() {
        if (isConnected()) {
            client.close();
            client = null;
        }
    }

    private Bson findQuery(String username) {
        return new Document("$text", new Document("$search", "\"" + username + "\"")).append("username", username);
    }

    private CompletableFuture<StatisticsData> getStats(Bson filter, Document insertion) {
        if (!isConnected()) {
            return CompletableFuture.completedFuture(new StatisticsData().defaults().temporary());
        }
        return CompletableFuture.supplyAsync(() -> {
            MongoCollection<Document> collection = getConnection().getCollection(STATS_COLLECTION);

            try (MongoCursor<Document> iterator = collection.find(filter).iterator()) {
                if (iterator.hasNext()) {
                    List<Integer> data = iterator.next().getList("data", Integer.class);
                    StatisticsData stats = new StatisticsData();

                    for (StatsRowHolder holder : StatsRowHolder.values()) {
                        stats.register(holder, data.get(holder.ordinal()));
                    }
                    return stats;
                } else {
                    if (insertion != null) {
                        collection.insertOne(insertion);
                    }
                    return new StatisticsData().defaults();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
            return null;
        });
    }

    @Override
    public CompletableFuture<StatisticsData> getStats(String username) {
        return getStats(findQuery(username.toLowerCase()), null);
    }

    @Override
    public CompletableFuture<StatisticsData> getStats(UUID uniqueId) {
        return getStats(Filters.eq("_id", uniqueId), null);
    }

    @Override
    public CompletableFuture<StatisticsData> getStatsOrInsert(UUID uniqueId, String username) {
        return getStats(Filters.eq("_id", uniqueId), new Document()
                .append("_id", uniqueId)
                .append("username_case_sensitive", username)
                .append("username", username.toLowerCase())
                .append("data", StatsRowHolder.DEFAULTS));
    }

    @Override
    public CompletableFuture<QuickbuyData> getQuickbuyOrInsert(UUID uniqueId) {
        if (!isConnected()) {
            return CompletableFuture.completedFuture(new QuickbuyData().defaults().temporary());
        }

        return CompletableFuture.supplyAsync(() -> {
            MongoCollection<Document> collection = getConnection().getCollection(QUICKBUY_COLLECTION);

            try (MongoCursor<Document> iterator = collection.find(Filters.eq("_id", uniqueId)).iterator()) {
                if (iterator.hasNext()) {
                    List<Integer> data = iterator.next().getList("data", Integer.class);
                    QuickbuyData quickbuyData = new QuickbuyData();

                    for (int i = 0; i < 21; i++) {
                        quickbuyData.register(i, data.get(i));
                    }
                    return quickbuyData;
                } else {
                    collection.insertOne(new Document()
                            .append("_id", uniqueId)
                            .append("data", ItemType.DEFAULTS));

                    return new QuickbuyData().defaults();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
            return null;
        });
    }

    @Override
    public CompletableFuture<HotBarData> getHotBarOrInsert(UUID uniqueId) {
        if (!isConnected()) {
            return CompletableFuture.completedFuture(new HotBarData().defaults().temporary());
        }

        return CompletableFuture.supplyAsync(() -> {
            MongoCollection<Document> collection = getConnection().getCollection(HOTBAR_COLLECTION);

            try (MongoCursor<Document> iterator = collection.find(Filters.eq("_id", uniqueId)).iterator()) {
                if (iterator.hasNext()) {
                    List<Integer> data = iterator.next().getList("data", Integer.class);
                    HotBarData hotBarData = new HotBarData();

                    for (int i = 0; i < data.size(); i++) {
                        hotBarData.register(i, data.get(i));
                    }
                    return hotBarData;
                } else {
                    collection.insertOne(new Document()
                            .append("_id", uniqueId)
                            .append("data", HotbarItemHolder.DEFAULTS));
                    return new HotBarData().defaults();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
            return null;
        });
    }

    @Override
    public CompletableFuture<PrivateGameOptions> getPrivateGameOptions(UUID leaderId) {
        return CompletableFuture.supplyAsync(() -> {
            try (MongoCursor<Document> iterator = getConnection()
                    .getCollection("privategame")
                    .find(Filters.eq("_id", leaderId))
                    .iterator()) {
                PrivateGameOptions options = new PrivateGameOptions();

                if (iterator.hasNext()) {
                    List<Integer> data = iterator.next().getList("data", Integer.class);
                    options.byMongoData(data);
                    return options;
                }
            }
            return new PrivateGameOptions();
        });
    }

    @Override
    public CompletableFuture<Void> savePrivateGameSettings(UUID uniqueId, PrivateGameOptions data) {
        if (!isConnected()) {
            return CompletableFuture.completedFuture(null);
        }
        return CompletableFuture.runAsync(() -> {
            updateRow(uniqueId, "privategame", Updates.set("data", data.getMongoData()));
        });
    }

    @Override
    public CompletableFuture<Void> saveStats(UUID uniqueId, StatisticsData data) {
        if (!isConnected() || data.isTemporary()) {
            return CompletableFuture.completedFuture(null);
        }

        return CompletableFuture.runAsync(() -> {
            List<Bson> updates = Lists.newArrayList();

            boolean[] sets = data.getSetChanges();

            int[] changes = data.getChanges();
            int[] base = data.getBase();

            for (int i = 0; i < changes.length; i++) {
                if (sets[i] || changes[i] != 0) {
                    updates.add(Updates.set("data." + i, sets[i]
                            ? changes[i]
                            : base[i] + changes[i]));
                }
            }
            if (!updates.isEmpty()) {
                updateRow(uniqueId, STATS_COLLECTION, Updates.combine(updates));
            }
        });
    }

    @Override
    public CompletableFuture<Void> saveQuickbuy(UUID uniqueId, QuickbuyData data) {
        if (!isConnected() || data.isTemporary()) {
            return CompletableFuture.completedFuture(null);
        }

        return CompletableFuture.runAsync(() -> {
            List<Bson> updates = Lists.newArrayList();

            int[] changes = data.getChanges();
            int[] base = data.getBase();

            for (int i = 0; i < changes.length; i++) {
                if (changes[i] != base[i]) {
                    updates.add(Updates.set("data." + i, changes[i]));
                }
            }
            if (!updates.isEmpty()) {
                updateRow(uniqueId, QUICKBUY_COLLECTION, Updates.combine(updates));
            }
        });
    }

    @Override
    public CompletableFuture<Void> saveHotBar(UUID uniqueId, HotBarData data) {
        if (!isConnected() || data.isTemporary()) {
            return CompletableFuture.completedFuture(null);
        }
        return CompletableFuture.runAsync(() -> {

            List<Bson> updates = Lists.newArrayList();

            int[] changes = data.getChanges();
            int[] base = data.getBase();

            for (int i = 0; i < changes.length; i++) {
                if (changes[i] != base[i]) {
                    updates.add(Updates.set("data." + i, changes[i]));
                }
            }
            if (!updates.isEmpty()) {
                updateRow(uniqueId, HOTBAR_COLLECTION, Updates.combine(updates));
            }
        });
    }

    private boolean isConnected() {
        return client != null;
    }

    private MongoDatabase getConnection() {
        if (isConnected()) {
            return client.getDatabase(database);
        } else {
            throw new IllegalStateException("Database is not connected");
        }
    }

    private void updateRow(UUID uniqueId, String table, Bson update) {
        UpdateOptions options = new UpdateOptions().upsert(false);

        getConnection()
                .getCollection(table)
                .updateOne(Filters.eq(uniqueId), update, options);
    }
}