/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.streams.state.internals.metrics;

import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.metrics.Sensor.RecordingLevel;
import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;

import org.junit.jupiter.api.Test;
import org.mockito.MockedStatic;

import java.util.Collections;
import java.util.Map;
import java.util.function.Supplier;

import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;

public class StateStoreMetricsTest {

    private static final String TASK_ID = "test-task";
    private static final String STORE_NAME = "test-store";
    private static final String STORE_TYPE = "test-type";
    private static final String STORE_LEVEL_GROUP = "stream-state-metrics";
    private static final String BUFFER_NAME = "test-buffer";

    private final Sensor expectedSensor = mock(Sensor.class);
    private final StreamsMetricsImpl streamsMetrics = mock(StreamsMetricsImpl.class);
    private final Map<String, String> storeTagMap = Collections.singletonMap("hello", "world");

    @Test
    public void shouldGetPutSensor() {
        final String metricName = "put";
        final String descriptionOfRate = "The average number of calls to put per second";
        final String descriptionOfAvg = "The average latency of calls to put";
        final String descriptionOfMax = "The maximum latency of calls to put";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.putSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetPutIfAbsentSensor() {
        final String metricName = "put-if-absent";
        final String descriptionOfRate = "The average number of calls to put-if-absent per second";
        final String descriptionOfAvg = "The average latency of calls to put-if-absent";
        final String descriptionOfMax = "The maximum latency of calls to put-if-absent";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.putIfAbsentSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetPutAllSensor() {
        final String metricName = "put-all";
        final String descriptionOfRate = "The average number of calls to put-all per second";
        final String descriptionOfAvg = "The average latency of calls to put-all";
        final String descriptionOfMax = "The maximum latency of calls to put-all";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.putAllSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetFetchSensor() {
        final String metricName = "fetch";
        final String descriptionOfRate = "The average number of calls to fetch per second";
        final String descriptionOfAvg = "The average latency of calls to fetch";
        final String descriptionOfMax = "The maximum latency of calls to fetch";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.fetchSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetGetSensor() {
        final String metricName = "get";
        final String descriptionOfRate = "The average number of calls to get per second";
        final String descriptionOfAvg = "The average latency of calls to get";
        final String descriptionOfMax = "The maximum latency of calls to get";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.getSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetAllSensor() {
        final String metricName = "all";
        final String descriptionOfRate = "The average number of calls to all per second";
        final String descriptionOfAvg = "The average latency of calls to all";
        final String descriptionOfMax = "The maximum latency of calls to all";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.allSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetRangeSensor() {
        final String metricName = "range";
        final String descriptionOfRate = "The average number of calls to range per second";
        final String descriptionOfAvg = "The average latency of calls to range";
        final String descriptionOfMax = "The maximum latency of calls to range";
        setupStreamsMetrics(
            metricName
        );

        getAndVerifySensor(
            () -> StateStoreMetrics.rangeSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetFlushSensor() {
        final String metricName = "flush";
        final String descriptionOfRate = "The average number of calls to flush per second";
        final String descriptionOfAvg = "The average latency of calls to flush";
        final String descriptionOfMax = "The maximum latency of calls to flush";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.flushSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetRemoveSensor() {
        final String metricName = "remove";
        final String descriptionOfRate = "The average number of calls to remove per second";
        final String descriptionOfAvg = "The average latency of calls to remove";
        final String descriptionOfMax = "The maximum latency of calls to remove";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.removeSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetDeleteSensor() {
        final String metricName = "delete";
        final String descriptionOfRate = "The average number of calls to delete per second";
        final String descriptionOfAvg = "The average latency of calls to delete";
        final String descriptionOfMax = "The maximum latency of calls to delete";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.deleteSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetRestoreSensor() {
        final String metricName = "restore";
        final String descriptionOfRate = "The average number of restorations per second";
        final String descriptionOfAvg = "The average latency of restorations";
        final String descriptionOfMax = "The maximum latency of restorations";
        setupStreamsMetrics(metricName);

        getAndVerifySensor(
            () -> StateStoreMetrics.restoreSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax,
            descriptionOfRate
        );
    }

    @Test
    public void shouldGetPrefixScanSensor() {
        final String metricName = "prefix-scan";
        final String descriptionOfRate = "The average number of calls to prefix-scan per second";
        final String descriptionOfAvg = "The average latency of calls to prefix-scan";
        final String descriptionOfMax = "The maximum latency of calls to prefix-scan";
        when(streamsMetrics.storeLevelSensor(TASK_ID, STORE_NAME, metricName, RecordingLevel.DEBUG))
            .thenReturn(expectedSensor);
        when(streamsMetrics.storeLevelTagMap(TASK_ID, STORE_TYPE, STORE_NAME)).thenReturn(storeTagMap);

        try (final MockedStatic<StreamsMetricsImpl> streamsMetricsStaticMock = mockStatic(StreamsMetricsImpl.class)) {
            final Sensor sensor = StateStoreMetrics.prefixScanSensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics);

            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addInvocationRateToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    metricName,
                    descriptionOfRate
                )
            );
            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addAvgAndMaxToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    latencyMetricName(metricName),
                    descriptionOfAvg,
                    descriptionOfMax
                )
            );
            assertThat(sensor, is(expectedSensor));
        }
    }

    @Test
    public void shouldGetSuppressionBufferCountSensor() {
        final String metricName = "suppression-buffer-count";
        final String descriptionOfAvg = "The average count of buffered records";
        final String descriptionOfMax = "The maximum count of buffered records";
        setupStreamsMetricsForSuppressionBufferSensor(metricName);

        verifySensorSuppressionBufferSensor(
            () -> StateStoreMetrics.suppressionBufferCountSensor(TASK_ID, STORE_TYPE, BUFFER_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax
        );
    }

    @Test
    public void shouldGetSuppressionBufferSizeSensor() {
        final String metricName = "suppression-buffer-size";
        final String descriptionOfAvg = "The average size of buffered records";
        final String descriptionOfMax = "The maximum size of buffered records";
        setupStreamsMetricsForSuppressionBufferSensor(metricName);

        verifySensorSuppressionBufferSensor(
            () -> StateStoreMetrics.suppressionBufferSizeSensor(TASK_ID, STORE_TYPE, BUFFER_NAME, streamsMetrics),
            metricName,
            descriptionOfAvg,
            descriptionOfMax
        );
    }

    @Test
    public void shouldGetRecordE2ELatencySensor() {
        final String metricName = "record-e2e-latency";
        final String e2eLatencyDescription =
            "end-to-end latency of a record, measuring by comparing the record timestamp with the "
                + "system time when it has been fully processed by the node";
        final String descriptionOfAvg = "The average " + e2eLatencyDescription;
        final String descriptionOfMin = "The minimum " + e2eLatencyDescription;
        final String descriptionOfMax = "The maximum " + e2eLatencyDescription;
        when(streamsMetrics.storeLevelSensor(TASK_ID, STORE_NAME, metricName, RecordingLevel.TRACE))
            .thenReturn(expectedSensor);
        when(streamsMetrics.storeLevelTagMap(TASK_ID, STORE_TYPE, STORE_NAME)).thenReturn(storeTagMap);

        try (final MockedStatic<StreamsMetricsImpl> streamsMetricsStaticMock = mockStatic(StreamsMetricsImpl.class)) {
            final Sensor sensor =
                StateStoreMetrics.e2ELatencySensor(TASK_ID, STORE_TYPE, STORE_NAME, streamsMetrics);
            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addAvgAndMinAndMaxToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    metricName,
                    descriptionOfAvg,
                    descriptionOfMin,
                    descriptionOfMax
                )
            );
            assertThat(sensor, is(expectedSensor));
        }
    }

    private void setupStreamsMetrics(final String metricName) {
        when(streamsMetrics.storeLevelSensor(TASK_ID, STORE_NAME, metricName, RecordingLevel.DEBUG))
            .thenReturn(expectedSensor);
        when(streamsMetrics.storeLevelTagMap(TASK_ID, STORE_TYPE, STORE_NAME)).thenReturn(storeTagMap);
    }

    private void getAndVerifySensor(final Supplier<Sensor> sensorSupplier,
                                    final String metricName,
                                    final String descriptionOfAvg,
                                    final String descriptionOfMax,
                                    final String descriptionOfRate) {
        try (final MockedStatic<StreamsMetricsImpl> streamsMetricsStaticMock = mockStatic(StreamsMetricsImpl.class)) {
            final Sensor sensor = sensorSupplier.get();
            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addInvocationRateToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    metricName,
                    descriptionOfRate
                )
            );
            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addAvgAndMaxToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    latencyMetricName(metricName),
                    descriptionOfAvg,
                    descriptionOfMax
                )
            );
            assertThat(sensor, is(expectedSensor));
        }
    }

    private String latencyMetricName(final String metricName) {
        return metricName + StreamsMetricsImpl.LATENCY_SUFFIX;
    }

    private void setupStreamsMetricsForSuppressionBufferSensor(final String metricName) {
        when(streamsMetrics.storeLevelSensor(
            TASK_ID,
            BUFFER_NAME,
            metricName,
            RecordingLevel.DEBUG
        )).thenReturn(expectedSensor);
        when(streamsMetrics.storeLevelTagMap(TASK_ID, STORE_TYPE, BUFFER_NAME)).thenReturn(storeTagMap);
    }

    private void verifySensorSuppressionBufferSensor(final Supplier<Sensor> sensorSupplier,
                                                     final String metricName,
                                                     final String descriptionOfAvg,
                                                     final String descriptionOfMax) {
        try (final MockedStatic<StreamsMetricsImpl> streamsMetricsStaticMock = mockStatic(StreamsMetricsImpl.class)) {
            final Sensor sensor = sensorSupplier.get();
            streamsMetricsStaticMock.verify(
                () -> StreamsMetricsImpl.addAvgAndMaxToSensor(
                    expectedSensor,
                    STORE_LEVEL_GROUP,
                    storeTagMap,
                    metricName,
                    descriptionOfAvg,
                    descriptionOfMax
                )
            );
            assertThat(sensor, is(expectedSensor));
        }
    }
}
