package ru.yandex.qe.dispenser.datasources;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.concurrent.TimeUnit;

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

public class ProxyStatement implements Statement {

    @NotNull
    protected final Statement delegate;
    @NotNull
    protected final ProxyConnection connection;
    @NotNull
    protected final DataSourceSensors dataSourceSensors;
    @NotNull
    protected final Stopwatch stopwatch;

    protected boolean releaseRegistered = false;

    public ProxyStatement(@NotNull final Statement delegate,
                          @NotNull final ProxyConnection connection,
                          @NotNull final DataSourceSensors dataSourceSensors) {
        this.delegate = delegate;
        this.connection = connection;
        this.dataSourceSensors = dataSourceSensors;
        this.stopwatch = Stopwatch.createStarted();
        dataSourceSensors.onStatementAllocation();
    }

    @Override
    public ResultSet executeQuery(final String sql) throws SQLException {
        return new ProxyResultSet(measure(() -> delegate.executeQuery(sql)), this, dataSourceSensors);
    }

    @Override
    public int executeUpdate(final String sql) throws SQLException {
        return measure(() -> delegate.executeUpdate(sql));
    }

    @Override
    public void close() throws SQLException {
        if (!releaseRegistered) {
            final long leaseDuration = stopwatch.elapsed(TimeUnit.MILLISECONDS);
            dataSourceSensors.onStatementClose(leaseDuration);
            releaseRegistered = true;
        }
        delegate.close();
    }

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

    @Override
    public void setMaxFieldSize(final int max) throws SQLException {
        delegate.setMaxFieldSize(max);
    }

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

    @Override
    public void setMaxRows(final int max) throws SQLException {
        delegate.setMaxRows(max);
    }

    @Override
    public void setEscapeProcessing(final boolean enable) throws SQLException {
        delegate.setEscapeProcessing(enable);
    }

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

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

    @Override
    public void cancel() throws SQLException {
        delegate.cancel();
    }

    @Override
    public SQLWarning getWarnings() throws SQLException {
        return delegate.getWarnings();
    }

    @Override
    public void clearWarnings() throws SQLException {
        delegate.clearWarnings();
    }

    @Override
    public void setCursorName(final String name) throws SQLException {
        delegate.setCursorName(name);
    }

    @Override
    public boolean execute(final String sql) throws SQLException {
        return measure(() -> delegate.execute(sql));
    }

    @Override
    public ResultSet getResultSet() throws SQLException {
        return new ProxyResultSet(delegate.getResultSet(), this, dataSourceSensors);
    }

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

    @Override
    public boolean getMoreResults() throws SQLException {
        return delegate.getMoreResults();
    }

    @Override
    public void setFetchDirection(final int direction) throws SQLException {
        delegate.setFetchDirection(direction);
    }

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

    @Override
    public void setFetchSize(final int rows) throws SQLException {
        delegate.setFetchSize(rows);
    }

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

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

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

    @Override
    public void addBatch(final String sql) throws SQLException {
        delegate.addBatch(sql);
    }

    @Override
    public void clearBatch() throws SQLException {
        delegate.clearBatch();
    }

    @Override
    public int[] executeBatch() throws SQLException {
        return measure(delegate::executeBatch);
    }

    @Override
    public Connection getConnection() throws SQLException {
        // Call delegate getConnection() to do any applicable validation,
        // then return our own proxy connection instance which wraps the same connection anyway
        delegate.getConnection();
        return connection;
    }

    @Override
    public boolean getMoreResults(final int current) throws SQLException {
        return delegate.getMoreResults(current);
    }

    @Override
    public ResultSet getGeneratedKeys() throws SQLException {
        return new ProxyResultSet(delegate.getGeneratedKeys(), this, dataSourceSensors);
    }

    @Override
    public int executeUpdate(final String sql, final int autoGeneratedKeys) throws SQLException {
        return measure(() -> delegate.executeUpdate(sql, autoGeneratedKeys));
    }

    @Override
    public int executeUpdate(final String sql, final int[] columnIndexes) throws SQLException {
        return measure(() -> delegate.executeUpdate(sql, columnIndexes));
    }

    @Override
    public int executeUpdate(final String sql, final String[] columnNames) throws SQLException {
        return measure(() -> delegate.executeUpdate(sql, columnNames));
    }

    @Override
    public boolean execute(final String sql, final int autoGeneratedKeys) throws SQLException {
        return measure(() -> delegate.execute(sql, autoGeneratedKeys));
    }

    @Override
    public boolean execute(final String sql, final int[] columnIndexes) throws SQLException {
        return measure(() -> delegate.execute(sql, columnIndexes));
    }

    @Override
    public boolean execute(final String sql, final String[] columnNames) throws SQLException {
        return measure(() -> delegate.execute(sql, columnNames));
    }

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

    @Override
    public boolean isClosed() throws SQLException {
        return delegate.isClosed();
    }

    @Override
    public void setPoolable(final boolean poolable) throws SQLException {
        delegate.setPoolable(poolable);
    }

    @Override
    public boolean isPoolable() throws SQLException {
        return delegate.isPoolable();
    }

    @Override
    public void closeOnCompletion() throws SQLException {
        delegate.closeOnCompletion();
    }

    @Override
    public boolean isCloseOnCompletion() throws SQLException {
        return delegate.isCloseOnCompletion();
    }

    @Override
    public long getLargeUpdateCount() throws SQLException {
        return delegate.getLargeUpdateCount();
    }

    @Override
    public void setLargeMaxRows(final long max) throws SQLException {
        delegate.setLargeMaxRows(max);
    }

    @Override
    public long getLargeMaxRows() throws SQLException {
        return delegate.getLargeMaxRows();
    }

    @Override
    public long[] executeLargeBatch() throws SQLException {
        return measure(delegate::executeLargeBatch);
    }

    @Override
    public long executeLargeUpdate(final String sql) throws SQLException {
        return measure(() -> delegate.executeLargeUpdate(sql));
    }

    @Override
    public long executeLargeUpdate(final String sql, final int autoGeneratedKeys) throws SQLException {
        return measure(() -> delegate.executeLargeUpdate(sql, autoGeneratedKeys));
    }

    @Override
    public long executeLargeUpdate(final String sql, final int[] columnIndexes) throws SQLException {
        return measure(() -> delegate.executeLargeUpdate(sql, columnIndexes));
    }

    @Override
    public long executeLargeUpdate(final String sql, final String[] columnNames) throws SQLException {
        return measure(() -> delegate.executeLargeUpdate(sql, columnNames));
    }

    @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);
    }

    protected <T> T measure(final SupplierWithSqlException<T> supplier) throws SQLException {
        final Stopwatch stopwatch = Stopwatch.createStarted();
        boolean success = false;
        try {
            dataSourceSensors.onBeforeRequest();
            final T result = supplier.get();
            success = true;
            return result;
        } finally {
            final long duration = stopwatch.elapsed(TimeUnit.MILLISECONDS);
            dataSourceSensors.onAfterRequest(duration, success);
        }
    }

}
