package ru.yandex.qe.jetty.metrics;

import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import com.google.common.base.Stopwatch;
import org.eclipse.jetty.io.Connection;
import org.eclipse.jetty.io.EndPoint;
import org.eclipse.jetty.server.ConnectionFactory;
import org.eclipse.jetty.server.Connector;
import org.eclipse.jetty.util.component.ContainerLifeCycle;

import ru.yandex.monlib.metrics.histogram.Histograms;
import ru.yandex.monlib.metrics.labels.Labels;
import ru.yandex.monlib.metrics.primitives.Histogram;
import ru.yandex.monlib.metrics.primitives.Rate;
import ru.yandex.monlib.metrics.registry.MetricRegistry;
import ru.yandex.qe.dispenser.solomon.SolomonHolder;

public class InstrumentedConnectionFactory extends ContainerLifeCycle implements ConnectionFactory {

    private static final AtomicLong FACTORY_INDEX = new AtomicLong(0);

    private final ConnectionFactory connectionFactory;
    private final Rate connectionRate;
    private final Histogram connectionDuration;

    public InstrumentedConnectionFactory(
            ConnectionFactory connectionFactory,
            SolomonHolder solomonHolder
    ) {
        this.connectionFactory = connectionFactory;
        MetricRegistry metricRegistry = solomonHolder.getRootRegistry();
        String factoryIndex = String.valueOf(FACTORY_INDEX.incrementAndGet());
        this.connectionRate = metricRegistry.rate("jetty.connections.rate", Labels.of("factory_index", factoryIndex));
        this.connectionDuration = metricRegistry.histogramRate("jetty.connections.duration_millis",
                Labels.of("factory_index", factoryIndex), Histograms.exponential(22, 2.0d, 1.0d));
        addBean(connectionFactory);
    }

    @Override
    public String getProtocol() {
        return connectionFactory.getProtocol();
    }

    @Override
    public List<String> getProtocols() {
        return connectionFactory.getProtocols();
    }

    @Override
    public Connection newConnection(Connector connector, EndPoint endPoint) {
        Connection connection = connectionFactory.newConnection(connector, endPoint);
        connection.addListener(new Connection.Listener() {
            private Stopwatch stopwatch;

            @Override
            public void onOpened(Connection connection) {
                this.stopwatch = Stopwatch.createStarted();
                connectionRate.inc();
            }

            @Override
            public void onClosed(Connection connection) {
                if (stopwatch != null) {
                    connectionDuration.record(stopwatch.elapsed(TimeUnit.MILLISECONDS));
                }
            }
        });
        return connection;
    }
}
