package ru.yandex.qe.dispenser.datasources;

import java.util.concurrent.atomic.AtomicLong;

import org.jetbrains.annotations.NotNull;

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;

public class DataSourceSensors {

    private static final String CONNECTION_LEASE_TIME = "db.connection.lease_time_millis";
    private static final String CONNECTION_LEASE_WAIT_TIME = "db.connection.lease_wait_time_millis";
    private static final String CONNECTION_LEASE_RATE = "db.connection.lease_rate";
    private static final String CONNECTION_RELEASE_RATE = "db.connection.release_rate";
    private static final String CONNECTION_LEASE_FAILURE_RATE = "db.connection.lease_failure_rate";
    private static final String COMMIT_RATE = "db.commit_rate";
    private static final String COMMIT_FAILURE_RATE = "db.commit_failure_rate";
    private static final String COMMIT_TIME = "db.commit_time_millis";
    private static final String ROLLBACK_RATE = "db.rollback_rate";
    private static final String ROLLBACK_FAILURE_RATE = "db.rollback_failure_rate";
    private static final String ROLLBACK_TIME = "db.rollback_time_millis";
    private static final String REQUEST_RATE = "db.request_rate";
    private static final String REQUEST_FAILURE_RATE = "db.request_failure_rate";
    private static final String REQUEST_TIME = "db.request_time";
    private static final String IN_FLIGHT_REQUESTS = "db.in_flight_requests";
    private static final String IN_FLIGHT_REQUESTS_HISTOGRAM = "db.in_flight_requests_histogram";
    private static final String STATEMENT_ALLOCATION_RATE = "db.statement.allocation_rate";
    private static final String RESULT_SET_ALLOCATION_RATE = "db.result_set.allocation_rate";
    private static final String STATEMENT_CLOSE_RATE = "db.statement.close_rate";
    private static final String RESULT_SET_CLOSE_RATE = "db.result_set.close_rate";
    private static final String STATEMENT_LEASE_TIME = "db.statement.lease_time_millis";
    private static final String RESULT_SET_LEASE_TIME = "db.result_set.lease_time_millis";
    private static final String RESULT_SET_SIZE = "db.result_set.size";
    private static final String DB_DATA_SOURCE = "db_data_source";

    @NotNull
    private final AtomicLong inFlightRequests = new AtomicLong(0L);
    @NotNull
    private final Histogram connectionLeaseTime;
    @NotNull
    private final Histogram connectionLeaseWaitTime;
    @NotNull
    private final Rate connectionLeaseRate;
    @NotNull
    private final Rate connectionLeaseFailureRate;
    @NotNull
    private final Rate connectionReleaseRate;
    @NotNull
    private final Rate commitRate;
    @NotNull
    private final Rate commitFailureRate;
    @NotNull
    private final Rate rollbackRate;
    @NotNull
    private final Rate rollbackFailureRate;
    @NotNull
    private final Histogram commitTime;
    @NotNull
    private final Histogram rollbackTime;
    @NotNull
    private final Rate requestRate;
    @NotNull
    private final Rate requestFailureRate;
    @NotNull
    private final Histogram requestTime;
    @NotNull
    private final Histogram inFlightRequestsHistogram;
    @NotNull
    private final Rate statementAllocationRate;
    @NotNull
    private final Rate resultSetAllocationRate;
    @NotNull
    private final Rate statementCloseRate;
    @NotNull
    private final Rate resultSetCloseRate;
    @NotNull
    private final Histogram statementLeaseTime;
    @NotNull
    private final Histogram resultSetLeaseTime;
    @NotNull
    private final Histogram resultSetSize;

    public DataSourceSensors(@NotNull final MetricRegistry registry, @NotNull final String dataSourceKey) {
        this.connectionLeaseTime = registry.histogramRate(CONNECTION_LEASE_TIME,
                Labels.of(DB_DATA_SOURCE, dataSourceKey), Histograms.exponential(22, 2.0d, 1.0d));
        this.connectionLeaseWaitTime = registry.histogramRate(CONNECTION_LEASE_WAIT_TIME,
                Labels.of(DB_DATA_SOURCE, dataSourceKey), Histograms.exponential(22, 2.0d, 1.0d));
        this.connectionLeaseRate = registry.rate(CONNECTION_LEASE_RATE,
                Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.connectionLeaseFailureRate = registry.rate(CONNECTION_LEASE_FAILURE_RATE,
                Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.connectionReleaseRate = registry.rate(CONNECTION_RELEASE_RATE,
                Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.commitRate = registry.rate(COMMIT_RATE,
                Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.commitFailureRate = registry.rate(COMMIT_FAILURE_RATE,
                Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.commitTime = registry.histogramRate(COMMIT_TIME, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
        this.rollbackRate = registry.rate(ROLLBACK_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.rollbackFailureRate = registry.rate(ROLLBACK_FAILURE_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.rollbackTime = registry.histogramRate(ROLLBACK_TIME, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
        this.requestRate = registry.rate(REQUEST_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.requestFailureRate = registry.rate(REQUEST_FAILURE_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.requestTime = registry.histogramRate(REQUEST_TIME, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
        registry.lazyGaugeInt64(IN_FLIGHT_REQUESTS, Labels.of(DB_DATA_SOURCE, dataSourceKey), inFlightRequests::get);
        this.inFlightRequestsHistogram = registry.histogramRate(IN_FLIGHT_REQUESTS_HISTOGRAM,
                Labels.of(DB_DATA_SOURCE, dataSourceKey), Histograms.exponential(22, 2.0d, 1.0d));
        this.statementAllocationRate = registry.rate(STATEMENT_ALLOCATION_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.resultSetAllocationRate = registry.rate(RESULT_SET_ALLOCATION_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.statementCloseRate = registry.rate(STATEMENT_CLOSE_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.resultSetCloseRate = registry.rate(RESULT_SET_CLOSE_RATE, Labels.of(DB_DATA_SOURCE, dataSourceKey));
        this.statementLeaseTime = registry.histogramRate(STATEMENT_LEASE_TIME, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
        this.resultSetLeaseTime = registry.histogramRate(RESULT_SET_LEASE_TIME, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
        this.resultSetSize = registry.histogramRate(RESULT_SET_SIZE, Labels.of(DB_DATA_SOURCE, dataSourceKey),
                Histograms.exponential(22, 2.0d, 1.0d));
    }

    public void onBeforeConnectionLease() {
        connectionLeaseRate.inc();
    }

    public void onAfterConnectionLease(final long connectionLeaseWaitDurationMillis, final boolean success) {
        if (!success) {
            connectionLeaseFailureRate.inc();
        }
        connectionLeaseWaitTime.record(connectionLeaseWaitDurationMillis);
    }

    public void onConnectionRelease(final long connectionLeaseDurationMillis) {
        connectionLeaseTime.record(connectionLeaseDurationMillis);
        connectionReleaseRate.inc();
    }

    public void onBeforeCommit() {
        commitRate.inc();
    }

    public void onAfterCommit(final long commitDurationMillis, final boolean success) {
        if (!success) {
            commitFailureRate.inc();
        }
        commitTime.record(commitDurationMillis);
    }

    public void onBeforeRollback() {
        rollbackRate.inc();
    }

    public void onAfterRollback(final long rollbackDurationMillis, final boolean success) {
        if (!success) {
            rollbackFailureRate.inc();
        }
        rollbackTime.record(rollbackDurationMillis);
    }

    public void onBeforeRequest() {
        final long value = inFlightRequests.incrementAndGet();
        inFlightRequestsHistogram.record(value);
        requestRate.inc();
    }

    public void onAfterRequest(final long requestDurationMillis, final boolean success) {
        final long value = inFlightRequests.decrementAndGet();
        inFlightRequestsHistogram.record(value);
        if (!success) {
            requestFailureRate.inc();
        }
        requestTime.record(requestDurationMillis);
    }

    public void onStatementAllocation() {
        statementAllocationRate.inc();
    }

    public void onStatementClose(final long statementLeaseDurationMillis) {
        statementLeaseTime.record(statementLeaseDurationMillis);
        statementCloseRate.inc();
    }

    public void onResultSetAllocation() {
        resultSetAllocationRate.inc();
    }

    public void onResultSetClose(final long resultSetLeaseDurationMillis, final long totalScrolledResults) {
        resultSetLeaseTime.record(resultSetLeaseDurationMillis);
        resultSetCloseRate.inc();
        resultSetSize.record(totalScrolledResults);
    }

}
