package ru.yandex.qe.dispenser.datasources;

import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

import javax.sql.DataSource;

import com.google.common.base.Stopwatch;
import org.jetbrains.annotations.NotNull;
import org.slf4j.LoggerFactory;

public class ProxyDataSource implements DataSource {

    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(ProxyDataSource.class);
    private static final long LONG_LEASE_WAIT_DURATION_THRESHOLD_MILLIS = 1000L;

    @NotNull
    private final DataSource delegate;
    @NotNull
    private final DataSourceSensors dataSourceSensors;

    public ProxyDataSource(@NotNull final DataSource delegate, @NotNull final DataSourceSensors dataSourceSensors) {
        this.delegate = delegate;
        this.dataSourceSensors = dataSourceSensors;
    }

    @Override
    public Connection getConnection() throws SQLException {
        return wrapConnection(delegate::getConnection);
    }

    @Override
    public Connection getConnection(final String username, final String password) throws SQLException {
        return wrapConnection(() -> delegate.getConnection(username, password));
    }

    @Override
    public <T> T unwrap(final Class<T> iface) throws SQLException {
        return delegate.unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(final Class<?> iface) throws SQLException {
        return delegate.isWrapperFor(iface);
    }

    @Override
    public PrintWriter getLogWriter() throws SQLException {
        return delegate.getLogWriter();
    }

    @Override
    public void setLogWriter(final PrintWriter out) throws SQLException {
        delegate.setLogWriter(out);
    }

    @Override
    public void setLoginTimeout(final int seconds) throws SQLException {
        delegate.setLoginTimeout(seconds);
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return delegate.getLoginTimeout();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        return delegate.getParentLogger();
    }

    private Connection wrapConnection(final SupplierWithSqlException<Connection> supplier) throws SQLException {
        final Stopwatch stopwatch = Stopwatch.createStarted();
        boolean success = false;
        try {
            dataSourceSensors.onBeforeConnectionLease();
            final Connection result = new ProxyConnection(supplier.get(), dataSourceSensors);
            success = true;
            return result;
        } finally {
            final long waitDuration = stopwatch.elapsed(TimeUnit.MILLISECONDS);
            if (waitDuration > LONG_LEASE_WAIT_DURATION_THRESHOLD_MILLIS) {
                LOG.warn("Long DB connection lease wait: {} ms, success: {}", waitDuration, success);
            }
            dataSourceSensors.onAfterConnectionLease(waitDuration, success);
        }
    }

}
