package ru.yandex.qe.http.telemetry;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

import org.apache.http.HttpClientConnection;
import org.apache.http.HttpConnectionMetrics;
import org.apache.http.conn.ConnectionRequest;
import org.apache.http.conn.HttpClientConnectionManager;
import org.apache.http.conn.routing.HttpRoute;
import org.apache.http.protocol.HttpContext;

import ru.yandex.monlib.metrics.histogram.Histograms;
import ru.yandex.monlib.metrics.labels.Labels;
import ru.yandex.monlib.metrics.primitives.GaugeInt64;
import ru.yandex.monlib.metrics.primitives.Histogram;
import ru.yandex.monlib.metrics.registry.MetricRegistry;


/**
 * Established by terry
 * on 07.02.16.
 */
public class InstrumentedHttpConnectionManager implements HttpClientConnectionManager {

    private final HttpClientConnectionManager delegate;

    private final GaugeInt64 bytesReceiveCounter;
    private final GaugeInt64 bytesSendCounter;

    private final Histogram bytesReceiveHistogram;
    private final Histogram bytesSendHistogram;

    public InstrumentedHttpConnectionManager(HttpClientConnectionManager delegate,
                                             MetricRegistry metricRegistry, String label) {
        this.delegate = delegate;
        bytesReceiveCounter = metricRegistry.gaugeInt64("http_client.received_bytes",
                Labels.of("http_client_label", label));
        bytesSendCounter = metricRegistry.gaugeInt64("http_client.sent_bytes",
                Labels.of("http_client_label", label));
        bytesReceiveHistogram = metricRegistry.histogramRate("http_client.responses_size",
                Labels.of("http_client_label", label), Histograms.exponential(12, 10.0d, 1.0d));
        bytesSendHistogram = metricRegistry.histogramRate("http_client.requests_size",
                Labels.of("http_client_label", label), Histograms.exponential(12, 10.0d, 1.0d));
    }

    @Override
    public ConnectionRequest requestConnection(HttpRoute route, Object state) {
        return delegate.requestConnection(route, state);
    }

    @Override
    public void releaseConnection(HttpClientConnection conn, Object newState, long validDuration, TimeUnit timeUnit) {
        final HttpConnectionMetrics metrics = conn.getMetrics();
        bytesReceiveCounter.add(metrics.getReceivedBytesCount());
        bytesReceiveHistogram.record(metrics.getReceivedBytesCount());
        bytesSendCounter.add(metrics.getSentBytesCount());
        bytesSendHistogram.record(metrics.getSentBytesCount());
        metrics.reset();
        delegate.releaseConnection(conn, newState, validDuration, timeUnit);
    }

    @Override
    public void connect(HttpClientConnection conn, HttpRoute route, int connectTimeout, HttpContext context) throws IOException {
        delegate.connect(conn, route, connectTimeout, context);
    }

    @Override
    public void upgrade(HttpClientConnection conn, HttpRoute route, HttpContext context) throws IOException {
        delegate.upgrade(conn, route, context);
    }

    @Override
    public void routeComplete(HttpClientConnection conn, HttpRoute route, HttpContext context) throws IOException {
        delegate.routeComplete(conn, route, context);
    }

    @Override
    public void closeIdleConnections(long idletime, TimeUnit tunit) {
        delegate.closeIdleConnections(idletime, tunit);
    }

    @Override
    public void closeExpiredConnections() {
        delegate.closeExpiredConnections();
    }

    @Override
    public void shutdown() {
        delegate.shutdown();
    }
}
