package io.digdag.core.database;

import javax.annotation.PreDestroy;
import javax.sql.DataSource;
import java.sql.SQLException;
import com.google.common.base.Throwables;
import com.google.inject.Inject;
import com.google.inject.Provider;
import org.h2.jdbcx.JdbcDataSource;
import com.zaxxer.hikari.HikariDataSource;
import com.zaxxer.hikari.HikariConfig;
import io.digdag.core.database.DatabaseMigrator;
import io.digdag.core.database.DatabaseConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class DataSourceProvider
        implements Provider<DataSource>, AutoCloseable
{
    private final Logger logger = LoggerFactory.getLogger(getClass());

    private final DatabaseConfig config;
    private DataSource ds;
    private AutoCloseable closer;

    @Inject
    public DataSourceProvider(DatabaseConfig config)
    {
        this.config = config;
    }

    public synchronized DataSource get()
    {
        if (ds == null) {
            switch (config.getType()) {
            case "h2":
                // h2 database doesn't need connection pool
                createSimpleDataSource();
                break;
            default:
                createPooledDataSource();
                break;
            }
        }
        return ds;
    }

    private void createSimpleDataSource()
    {
        String url = DatabaseConfig.buildJdbcUrl(config);

        // By default, H2 closes database when all of the connections are closed. When database
        // is closed, data is gone with in-memory mode. It's unexpected. However, if here disables
        // that such behavior using DB_CLOSE_DELAY=-1 option, there're no methods to close the
        // database explicitly. Only way to close is to not disable shutdown hook of H2 database
        // (DB_CLOSE_ON_EXIT=TRUE). But this also causes unexpected behavior when PreDestroy is
        // triggered in a shutdown hook. Therefore, here needs to rely on injector to take care of
        // dependencies so that the database is closed after calling all other PreDestroy methods
        // that depend on this DataSourceProvider.
        // To solve this issue, here holds one Connection until PreDestroy.
        JdbcDataSource ds = new JdbcDataSource();
        ds.setUrl(url + ";DB_CLOSE_ON_EXIT=FALSE");

        logger.debug("Using database URL {}", url);

        try {
            this.closer = ds.getConnection();
        }
        catch (SQLException ex) {
            throw Throwables.propagate(ex);
        }
        this.ds = ds;
    }

    private void createPooledDataSource()
    {
        String url = DatabaseConfig.buildJdbcUrl(config);

        HikariConfig hikari = new HikariConfig();
        hikari.setJdbcUrl(url);
        hikari.setDriverClassName(DatabaseMigrator.getDriverClassName(config.getType()));
        hikari.setDataSourceProperties(DatabaseConfig.buildJdbcProperties(config));

        hikari.setConnectionTimeout(config.getConnectionTimeout() * 1000);
        hikari.setIdleTimeout(config.getIdleTimeout() * 1000);
        hikari.setValidationTimeout(config.getValidationTimeout() * 1000);
        hikari.setMaximumPoolSize(config.getMaximumPoolSize());
        hikari.setMinimumIdle(config.getMinimumPoolSize());
        hikari.setRegisterMbeans(config.getEnableJMX());
        hikari.setLeakDetectionThreshold(config.getLeakDetectionThreshold());

        // Here should not set connectionTestQuery (that overrides isValid) because
        // ThreadLocalTransactionManager.commit assumes that Connection.isValid returns
        // false when an error happened during a transaction.

        logger.debug("Using database URL {}", hikari.getJdbcUrl());

        HikariDataSource ds = new HikariDataSource(hikari);
        this.ds = ds;
        this.closer = ds;
    }

    @PreDestroy
    public synchronized void close()
    {
        if (ds != null) {
            try {
                closer.close();
            }
            catch (Exception ex) {
                throw Throwables.propagate(ex);
            }
            ds = null;
            closer = null;
        }
    }
}