AltitudeWeb/database/src/main/java/com/alttd/altitudeweb/setup/Connection.java

package com.alttd.altitudeweb.setup;

import com.alttd.altitudeweb.database.Databases;
import com.alttd.altitudeweb.database.web_db.DatabaseSettings;
import com.alttd.altitudeweb.database.web_db.SettingsMapper;
import com.alttd.altitudeweb.type_handler.UUIDTypeHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.datasource.pooled.PooledDataSource;
import org.apache.ibatis.mapping.Environment;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.ibatis.session.SqlSessionFactoryBuilder;
import org.apache.ibatis.transaction.jdbc.JdbcTransactionFactory;

import java.util.HashMap;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

@Slf4j
public class Connection {

    private static final HashMap<Databases, Connection> connections = new HashMap<>();
    private SqlSessionFactory sqlSessionFactory;
    private final DatabaseSettings settings;
    private final AddMappers addMappers;

    private Connection(DatabaseSettings settings, AddMappers addMappers) {
        this.settings = settings;
        this.addMappers = addMappers;
    }

    public static void initDatabases() {
        InitializeWebDb.init();
        InitializeLiteBans.init();
        InitializeLuckPerms.init();
        InitializeDiscord.init();
    }

    @FunctionalInterface
    public interface AddMappers {
        void apply(Configuration configuration);
    }

    public static Connection getConnection(Databases database) {
        if (connections.containsKey(database)) {
            return connections.get(database);
        }
        throw new RuntimeException("Database " + database + " has not been initialized");
    }

    protected static CompletableFuture<Connection> getConnection(Databases database, AddMappers addMappers) {
        if (connections.containsKey(database)) {
            return CompletableFuture.completedFuture(connections.get(database));
        }
        if (database == Databases.DEFAULT) {
            return loadDefaultDatabase(addMappers);
        }
        CompletableFuture<DatabaseSettings> settingsFuture = new CompletableFuture<>();
        getConnection(Databases.DEFAULT, (mapper -> mapper.addMapper(SettingsMapper.class))).thenApply(connection -> {
            log.debug("Loading settings for database {}", database.getInternalName());
            connection.runQuery(session -> {
                log.debug("Running query to load settings for database");
                DatabaseSettings loadedSettings = session.getMapper(SettingsMapper.class).getSettings(database.getInternalName());
                if (loadedSettings == null) {
                    log.error("Failed to load settings for database {}", database.getInternalName());
                }
                log.debug("Loaded settings {}", loadedSettings);
                settingsFuture.complete(loadedSettings);
            });
            return null;
        });
        return settingsFuture.thenApply(loadedSettings -> {
            log.debug("Storing connection for database {}", database.getInternalName());
            Connection connection = new Connection(loadedSettings, addMappers);
            connections.put(database, connection);
            return connection;
        });
    }

    private static CompletableFuture<Connection> loadDefaultDatabase(AddMappers addMappers) {
        DatabaseSettings databaseSettings = new DatabaseSettings(
                System.getenv("DB_HOST"),
                Integer.parseInt(System.getenv("DB_PORT")),
                System.getenv("DB_NAME"),
                System.getenv("DB_USER"),
                System.getenv("DB_PASS")
        );
        log.debug("Loaded default database settings {}", databaseSettings);
        Connection connection = new Connection(databaseSettings, addMappers);
        log.debug("Created default database connection {}", connection);
        connections.put(Databases.DEFAULT, connection);
        return CompletableFuture.completedFuture(connection);
    }

    public void runQuery(Consumer<SqlSession> consumer) {
        new Thread(() -> {
            if (sqlSessionFactory == null) {
                sqlSessionFactory = createSqlSessionFactory(settings, addMappers);
            }

            SqlSession session = null;
            try {
                session = sqlSessionFactory.openSession();
                consumer.accept(session);
                session.commit();
            } catch (Exception e) {
                if (session != null) {
                    session.rollback();
                }
                log.error("Failed to run query", e);
            } finally {
                if (session != null) {
                    session.close();
                }
            }
        }).start();
    }

    private SqlSessionFactory createSqlSessionFactory(DatabaseSettings settings, AddMappers addMappers) {
        PooledDataSource dataSource = new PooledDataSource();
        dataSource.setDriver("com.mysql.cj.jdbc.Driver");
        dataSource.setUrl(String.format("jdbc:mysql://%s:%d/%s", settings.host(),
                                        settings.port(), settings.name()));
        dataSource.setUsername(settings.username());
        dataSource.setPassword(settings.password());
        Environment environment = new Environment("production", new JdbcTransactionFactory(), dataSource);
        Configuration configuration = new Configuration(environment);
        configuration.getTypeHandlerRegistry().register(UUID.class, UUIDTypeHandler.class);
        addMappers.apply(configuration);

        return new SqlSessionFactoryBuilder().build(configuration);
    }

}