/*
 * 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.clients.consumer;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.KafkaClient;
import org.apache.kafka.clients.MockClient;
import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy;
import org.apache.kafka.clients.consumer.internals.ShareConsumerImpl;
import org.apache.kafka.clients.consumer.internals.ShareConsumerMetadata;
import org.apache.kafka.clients.consumer.internals.SubscriptionState;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.metrics.KafkaMetric;
import org.apache.kafka.common.metrics.Measurable;
import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.requests.RequestTestUtils;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.telemetry.internals.ClientTelemetryReporter;
import org.apache.kafka.common.utils.LogCaptureAppender;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;

import org.apache.logging.log4j.Level;
import org.junit.jupiter.api.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.internal.stubbing.answers.CallsRealMethods;

import java.time.Duration;
import java.util.AbstractMap;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_SHARE_METRIC_GROUP_PREFIX;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.atMostOnce;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.never;

public class KafkaShareConsumerMetricsTest {
    private final String topic = "test";
    private final Uuid topicId = Uuid.randomUuid();
    private final Map<String, Uuid> topicIds = Stream.of(
                    new AbstractMap.SimpleEntry<>(topic, topicId))
            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    private final Time time = new MockTime();
    private final SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.EARLIEST);
    private final String groupId = "mock-group";

    @Test
    public void testPollTimeMetrics() {
        ShareConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
        consumer.subscribe(Set.of(topic));
        // MetricName objects to check
        Metrics metrics = consumer.metricsRegistry();
        MetricName lastPollSecondsAgoName = metrics.metricName("last-poll-seconds-ago", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
        MetricName timeBetweenPollAvgName = metrics.metricName("time-between-poll-avg", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
        MetricName timeBetweenPollMaxName = metrics.metricName("time-between-poll-max", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
        // Test default values
        assertEquals(-1.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Call first poll
        consumer.poll(Duration.ZERO);
        assertEquals(0.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 5,000 (total time = 5,000)
        time.sleep(5 * 1000L);
        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call second poll
        consumer.poll(Duration.ZERO);
        assertEquals(2.5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 10,000 (total time = 15,000)
        time.sleep(10 * 1000L);
        assertEquals(10.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call third poll
        consumer.poll(Duration.ZERO);
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 5,000 (total time = 20,000)
        time.sleep(5 * 1000L);
        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call fourth poll
        consumer.poll(Duration.ZERO);
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
    }

    @Test
    public void testPollIdleRatio() {
        ShareConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
        // MetricName object to check
        Metrics metrics = consumer.metricsRegistry();
        MetricName pollIdleRatio = metrics.metricName("poll-idle-ratio-avg", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
        // Test default value
        assertEquals(Double.NaN, consumer.metrics().get(pollIdleRatio).metricValue());

        // 1st poll
        // Spend 50ms in poll so value = 1.0
        consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
        time.sleep(50);
        consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());

        assertEquals(1.0d, consumer.metrics().get(pollIdleRatio).metricValue());

        // 2nd poll
        // Spend 50m outside poll and 0ms in poll so value = 0.0
        time.sleep(50);
        consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
        consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());

        // Avg of first two data points
        assertEquals((1.0d + 0.0d) / 2, consumer.metrics().get(pollIdleRatio).metricValue());

        // 3rd poll
        // Spend 25ms outside poll and 25ms in poll so value = 0.5
        time.sleep(25);
        consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
        time.sleep(25);
        consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());

        // Avg of three data points
        assertEquals((1.0d + 0.0d + 0.5d) / 3, consumer.metrics().get(pollIdleRatio).metricValue());
    }

    private static boolean consumerMetricPresent(KafkaShareConsumer<String, String> consumer, String name) {
        MetricName metricName = new MetricName(name, CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics", "", Map.of());
        return consumer.metricsRegistry().metrics().containsKey(metricName);
    }

    @Test
    public void testClosingConsumerUnregistersConsumerMetrics() {
        Time time = new MockTime(1L);
        ShareConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
        consumer.subscribe(List.of(topic));
        assertTrue(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
        assertTrue(consumerMetricPresent(consumer, "time-between-poll-avg"));
        assertTrue(consumerMetricPresent(consumer, "time-between-poll-max"));
        consumer.close();
        assertFalse(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
        assertFalse(consumerMetricPresent(consumer, "time-between-poll-avg"));
        assertFalse(consumerMetricPresent(consumer, "time-between-poll-max"));
    }

    @Test
    public void testRegisteringCustomMetricsDoesntAffectConsumerMetrics() {
        Time time = new MockTime(1L);
        ShareConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
        Map<MetricName, KafkaMetric> customMetrics = customMetrics();
        customMetrics.forEach((name, metric) -> consumer.registerMetricForSubscription(metric));

        Map<MetricName, ? extends Metric> consumerMetrics = consumer.metrics();
        customMetrics.forEach((name, metric) -> assertFalse(consumerMetrics.containsKey(name)));
    }

    @Test
    public void testRegisteringCustomMetricsWithSameNameDoesntAffectConsumerMetrics() {
        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
            appender.setClassLogger(ShareConsumerImpl.class, Level.DEBUG);
            Time time = new MockTime(1L);
            ShareConsumerMetadata metadata = createMetadata(subscription);
            MockClient client = new MockClient(time, metadata);
            initMetadata(client, Map.of(topic, 1));

            KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
            KafkaMetric existingMetricToAdd = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.registerMetricForSubscription(existingMetricToAdd);
            final String expectedMessage = String.format("Skipping registration for metric %s. Existing consumer metrics cannot be overwritten.", existingMetricToAdd.metricName());
            assertTrue(appender.getMessages().stream().anyMatch(m -> m.contains(expectedMessage)));
        }
    }

    @Test
    public void testUnregisteringCustomMetricsWithSameNameDoesntAffectConsumerMetrics() {
        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
            appender.setClassLogger(ShareConsumerImpl.class, Level.DEBUG);
            Time time = new MockTime(1L);
            ShareConsumerMetadata metadata = createMetadata(subscription);
            MockClient client = new MockClient(time, metadata);
            initMetadata(client, Map.of(topic, 1));

            KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
            KafkaMetric existingMetricToRemove = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.unregisterMetricFromSubscription(existingMetricToRemove);
            final String expectedMessage = String.format("Skipping unregistration for metric %s. Existing consumer metrics cannot be removed.", existingMetricToRemove.metricName());
            assertTrue(appender.getMessages().stream().anyMatch(m -> m.contains(expectedMessage)));
        }
    }

    @Test
    public void testShouldOnlyCallMetricReporterMetricChangeOnceWithExistingConsumerMetric() {
        try (MockedStatic<CommonClientConfigs> mockedCommonClientConfigs = mockStatic(CommonClientConfigs.class, new CallsRealMethods())) {
            ClientTelemetryReporter clientTelemetryReporter = mock(ClientTelemetryReporter.class);
            clientTelemetryReporter.configure(any());
            mockedCommonClientConfigs.when(() -> CommonClientConfigs.telemetryReporter(anyString(), any())).thenReturn(Optional.of(clientTelemetryReporter));

            Time time = new MockTime(1L);
            ShareConsumerMetadata metadata = createMetadata(subscription);
            MockClient client = new MockClient(time, metadata);
            initMetadata(client, Map.of(topic, 1));

            KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);

            KafkaMetric existingMetric = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.registerMetricForSubscription(existingMetric);
            // This test would fail without the check as the existing metric is registered in the consumer on startup
            Mockito.verify(clientTelemetryReporter, atMostOnce()).metricChange(existingMetric);
        }
    }

    @Test
    public void testShouldNotCallMetricReporterMetricRemovalWithExistingConsumerMetric() {
        try (MockedStatic<CommonClientConfigs> mockedCommonClientConfigs = mockStatic(CommonClientConfigs.class, new CallsRealMethods())) {
            ClientTelemetryReporter clientTelemetryReporter = mock(ClientTelemetryReporter.class);
            clientTelemetryReporter.configure(any());
            mockedCommonClientConfigs.when(() -> CommonClientConfigs.telemetryReporter(anyString(), any())).thenReturn(Optional.of(clientTelemetryReporter));

            Time time = new MockTime(1L);
            ShareConsumerMetadata metadata = createMetadata(subscription);
            MockClient client = new MockClient(time, metadata);
            initMetadata(client, Map.of(topic, 1));

            KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);

            KafkaMetric existingMetric = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.unregisterMetricFromSubscription(existingMetric);
            Mockito.verify(clientTelemetryReporter, never()).metricRemoval(existingMetric);
        }
    }

    @Test
    public void testUnregisteringNonexistingMetricsDoesntCauseError() {
        Time time = new MockTime(1L);
        ShareConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);

        Map<MetricName, KafkaMetric> customMetrics = customMetrics();
        // Metrics never registered but removed should not cause an error
        customMetrics.forEach((name, metric) -> assertDoesNotThrow(() -> consumer.unregisterMetricFromSubscription(metric)));
    }

    private ShareConsumerMetadata createMetadata(SubscriptionState subscription) {
        return new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,
                subscription, new LogContext(), new ClusterResourceListeners());
    }

    private KafkaShareConsumer<String, String> newShareConsumer(Time time,
                                                                KafkaClient client,
                                                                SubscriptionState subscription,
                                                                ShareConsumerMetadata metadata) {
        return newShareConsumer(
                time,
                client,
                subscription,
                metadata,
                groupId,
                Optional.of(new StringDeserializer())
        );
    }

    private KafkaShareConsumer<String, String> newShareConsumer(Time time,
                                                                KafkaClient client,
                                                                SubscriptionState subscriptions,
                                                                ShareConsumerMetadata metadata,
                                                                String groupId,
                                                                Optional<Deserializer<String>> valueDeserializerOpt) {
        String clientId = "mock-consumer";
        Deserializer<String> keyDeserializer = new StringDeserializer();
        Deserializer<String> valueDeserializer = valueDeserializerOpt.orElse(new StringDeserializer());
        LogContext logContext = new LogContext();
        ShareConsumerConfig config = newConsumerConfig(groupId, valueDeserializer);
        return new KafkaShareConsumer<>(
                logContext,
                clientId,
                groupId,
                config,
                keyDeserializer,
                valueDeserializer,
                time,
                client,
                subscriptions,
                metadata
        );
    }

    private ShareConsumerConfig newConsumerConfig(String groupId,
                                             Deserializer<String> valueDeserializer) {
        String clientId = "mock-consumer";
        long retryBackoffMs = 100;
        long retryBackoffMaxMs = 1000;
        int minBytes = 1;
        int maxBytes = Integer.MAX_VALUE;
        int maxWaitMs = 500;
        int fetchSize = 1024 * 1024;
        int maxPollRecords = Integer.MAX_VALUE;
        boolean checkCrcs = true;
        int rebalanceTimeoutMs = 60000;
        int defaultApiTimeoutMs = 60000;
        int requestTimeoutMs = defaultApiTimeoutMs / 2;

        Map<String, Object> configs = new HashMap<>();
        configs.put(ConsumerConfig.CHECK_CRCS_CONFIG, checkCrcs);
        configs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
        configs.put(ConsumerConfig.CLIENT_RACK_CONFIG, CommonClientConfigs.DEFAULT_CLIENT_RACK);
        configs.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, defaultApiTimeoutMs);
        configs.put(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxBytes);
        configs.put(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, maxWaitMs);
        configs.put(ConsumerConfig.FETCH_MIN_BYTES_CONFIG, minBytes);
        configs.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        configs.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, fetchSize);
        configs.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, rebalanceTimeoutMs);
        configs.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords);
        configs.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs);
        configs.put(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG, retryBackoffMaxMs);
        configs.put(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, retryBackoffMs);
        configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer.getClass());
        configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");

        return new ShareConsumerConfig(configs);
    }
    private void initMetadata(MockClient mockClient, Map<String, Integer> partitionCounts) {
        Map<String, Uuid> metadataIds = new HashMap<>();
        for (String name : partitionCounts.keySet()) {
            metadataIds.put(name, topicIds.get(name));
        }
        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, metadataIds);

        mockClient.updateMetadata(initialMetadata);
    }

    private Map<MetricName, KafkaMetric> customMetrics() {
        MetricConfig metricConfig = new MetricConfig();
        Object lock = new Object();
        MetricName metricNameOne = new MetricName("metricOne", "stream-metrics", "description for metric one", new HashMap<>());
        MetricName metricNameTwo = new MetricName("metricTwo", "stream-metrics", "description for metric two", new HashMap<>());

        KafkaMetric streamClientMetricOne = new KafkaMetric(lock, metricNameOne, (Measurable) (m, now) -> 1.0, metricConfig, Time.SYSTEM);
        KafkaMetric streamClientMetricTwo = new KafkaMetric(lock, metricNameTwo, (Measurable) (m, now) -> 2.0, metricConfig, Time.SYSTEM);
        return Map.of(metricNameOne, streamClientMetricOne, metricNameTwo, streamClientMetricTwo);
    }
}
