package com.atlassian.stash.internal.jdbc;

import com.atlassian.bitbucket.util.ProxyUtils;

import javax.annotation.Nonnull;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.BiConsumer;

/**
 * A {@link ConnectionTracker} which wraps tracked connections in a Java {@code Proxy}. The proxy is used to
 * decorate {@code Connection.close()} to stop tracking the connection.
 *
 * @since 4.0
 */
public class ProxyingConnectionTracker implements ConnectionTracker {

    private final ConcurrentMap<Connection, Thread> leasesByConnection;

    public ProxyingConnectionTracker(int maxConnections) {
        leasesByConnection = new ConcurrentHashMap<>(maxConnections + (maxConnections / 2), 0.75f, 3);
    }

    @Override
    public void forEach(@Nonnull BiConsumer<Connection, Thread> consumer) {
        leasesByConnection.forEach(consumer);
    }

    @Override
    public int getConnectionCount() {
        return leasesByConnection.size();
    }

    @Nonnull
    @Override
    public Connection register(@Nonnull Connection connection) {
        return ProxyUtils.createProxy(Connection.class, new ConnectionInvocationHandler(connection));
    }

    private class ConnectionInvocationHandler implements InvocationHandler {

        private final Connection delegate;

        public ConnectionInvocationHandler(Connection delegate) {
            this.delegate = delegate;

            leasesByConnection.put(delegate, Thread.currentThread());
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            try {
                return method.invoke(delegate, args);
            } catch (InvocationTargetException e) {
                throw e.getCause();
            } finally {
                if ("close".equals(method.getName())) {
                    leasesByConnection.remove(delegate);
                }
            }
        }
    }
}
