/*
 * 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.processor.internals;

import org.apache.kafka.clients.producer.MockProducer;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.metrics.KafkaMetric;
import org.apache.kafka.common.metrics.Measurable;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.TopologyConfig;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.TimestampExtractor;
import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
import org.apache.kafka.streams.state.internals.ThreadCache;
import org.apache.kafka.test.MockClientSupplier;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;

import java.io.File;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;

import static java.util.Collections.emptySet;
import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.closeTo;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.STRICT_STUBS)
public class ActiveTaskCreatorTest {

    @Mock
    private InternalTopologyBuilder builder;
    @Mock
    private StateDirectory stateDirectory;

    private final MockClientSupplier mockClientSupplier = new MockClientSupplier();
    private final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(new Metrics(), "clientId", new MockTime());
    private final Map<String, Object> properties = mkMap(
        mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
        mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "dummy:1234")
    );
    final UUID uuid = UUID.randomUUID();

    private ActiveTaskCreator activeTaskCreator;

    // non-EOS test

    // functional test

    @Test
    public void shouldConstructProducerMetricsWithEosDisabled() {
        shouldConstructStreamsProducerMetric();
    }

    @Test
    public void shouldConstructClientIdWithEosDisabled() {
        createTasks();

        final String clientIds = activeTaskCreator.producerClientIds();

        assertThat(clientIds, is("clientId-StreamThread-0-producer"));
    }

    @Test
    public void shouldCloseIfEosDisabled() {
        createTasks();

        activeTaskCreator.close();

        assertThat(mockClientSupplier.producers.get(0).closed(), is(true));
    }

    @Test
    public void shouldReturnBlockedTimeWhenStreamsProducer() {
        final double blockedTime = 123.0;
        createTasks();
        final MockProducer<?, ?> producer = mockClientSupplier.producers.get(0);
        addMetric(producer, "flush-time-ns-total", blockedTime);

        assertThat(activeTaskCreator.totalProducerBlockedTime(), closeTo(blockedTime, 0.01));
    }

    // error handling

    @Test
    public void shouldReturnStreamsProducerIfAtLeastOnceIsEnabled() {
        createTasks();

        final StreamsProducer threadProducer = activeTaskCreator.streamsProducer();

        assertThat(mockClientSupplier.producers.size(), is(1));
        assertThat(threadProducer.kafkaProducer(), is(mockClientSupplier.producers.get(0)));
    }

    @Test
    public void shouldThrowStreamsExceptionOnErrorCloseIfEosDisabled() {
        createTasks();
        mockClientSupplier.producers.get(0).closeException = new RuntimeException("KABOOM!");

        final StreamsException thrown = assertThrows(
            StreamsException.class,
            activeTaskCreator::close
        );

        assertThat(thrown.getMessage(), is("Thread producer encounter error trying to close."));
        assertThat(thrown.getCause().getMessage(), is("KABOOM!"));
    }



    // eos-v2 test

    // functional test

    @Test
    public void shouldReturnStreamsProducerIfEosV2Enabled() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");

        createTasks();

        final StreamsProducer threadProducer = activeTaskCreator.streamsProducer();

        assertThat(mockClientSupplier.producers.size(), is(1));
        assertThat(threadProducer.kafkaProducer(), is(mockClientSupplier.producers.get(0)));
    }

    @Test
    public void shouldConstructProducerMetricsWithEosV2Enabled() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");

        shouldConstructStreamsProducerMetric();
    }

    @Test
    public void shouldConstructClientIdWithEosV2Enabled() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");
        createTasks();

        final String clientIds = activeTaskCreator.producerClientIds();

        assertThat(clientIds, is("clientId-StreamThread-0-producer"));
    }

    @Test
    public void shouldCloseIfEosV2Enabled() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");
        createTasks();

        activeTaskCreator.close();

        assertThat(activeTaskCreator.isClosed(), is(true));
        assertThat(mockClientSupplier.producers.get(0).closed(), is(true));
    }

    @Test
    public void shouldNotResetProducerAfterDisableRest() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");
        createTasks();
        assertThat(mockClientSupplier.producers.size(), is(1));

        activeTaskCreator.close();
        activeTaskCreator.reInitializeProducer();
        // Verifies that disableReset() prevents reInitializeProducer() from creating a new producer instance
        // Without disabling reset, the producers collection would contain more than one producer
        assertThat("Producer should not be recreated after disabling reset",
            mockClientSupplier.producers.size(),
            is(1));
    }

    // error handling

    @Test
    public void shouldThrowStreamsExceptionOnErrorCloseIfEosV2Enabled() {
        properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2);
        mockClientSupplier.setApplicationIdForProducer("appId");
        createTasks();
        mockClientSupplier.producers.get(0).closeException = new RuntimeException("KABOOM!");

        final StreamsException thrown = assertThrows(
            StreamsException.class,
            activeTaskCreator::close
        );

        assertThat(thrown.getMessage(), is("Thread producer encounter error trying to close."));
        assertThat(thrown.getCause().getMessage(), is("KABOOM!"));
    }

    private void shouldConstructStreamsProducerMetric() {
        createTasks();

        final MetricName testMetricName = new MetricName("test_metric", "", "", new HashMap<>());
        final Metric testMetric = new KafkaMetric(
            new Object(),
            testMetricName,
            (Measurable) (config, now) -> 0,
            null,
            new MockTime());
        mockClientSupplier.producers.get(0).setMockMetrics(testMetricName, testMetric);
        assertThat(mockClientSupplier.producers.size(), is(1));

        final Map<MetricName, Metric> producerMetrics = activeTaskCreator.producerMetrics();

        assertThat(producerMetrics.size(), is(1));
        assertThat(producerMetrics.get(testMetricName), is(testMetric));
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private void createTasks() {
        final TaskId task00 = new TaskId(0, 0);
        final TaskId task01 = new TaskId(0, 1);

        final ProcessorTopology topology = mock(ProcessorTopology.class);
        final SourceNode sourceNode = mock(SourceNode.class);

        when(builder.topologyConfigs()).thenReturn(new TopologyConfig(new StreamsConfig(properties)));
        when(builder.buildSubtopology(0)).thenReturn(topology);
        when(topology.sinkTopics()).thenReturn(emptySet());
        when(stateDirectory.getOrCreateDirectoryForTask(task00)).thenReturn(mock(File.class));
        when(stateDirectory.getOrCreateDirectoryForTask(task01)).thenReturn(mock(File.class));
        when(topology.source("topic")).thenReturn(sourceNode);
        when(sourceNode.timestampExtractor()).thenReturn(mock(TimestampExtractor.class));
        when(topology.sources()).thenReturn(Collections.singleton(sourceNode));

        final StreamsConfig config = new StreamsConfig(properties);
        activeTaskCreator = new ActiveTaskCreator(
            new TopologyMetadata(builder, config),
            config,
            streamsMetrics,
            stateDirectory,
            new ThreadCache(new LogContext(), 0L, streamsMetrics),
            new MockTime(),
            mockClientSupplier,
            "clientId-StreamThread-0",
            0,
            uuid,
            new LogContext(),
            false);

        assertThat(
            activeTaskCreator.createTasks(
                mockClientSupplier.consumer,
                mkMap(
                    mkEntry(task00, Collections.singleton(new TopicPartition("topic", 0))),
                    mkEntry(task01, Collections.singleton(new TopicPartition("topic", 1)))
                )
            ).stream().map(Task::id).collect(Collectors.toSet()),
            equalTo(Set.of(task00, task01))
        );
    }

    private void addMetric(
        final MockProducer<?, ?> producer,
        final String name,
        final double value) {
        final MetricName metricName = metricName(name);
        producer.setMockMetrics(metricName, new Metric() {
            @Override
            public MetricName metricName() {
                return metricName;
            }

            @Override
            public Object metricValue() {
                return value;
            }
        });
    }

    private MetricName metricName(final String name) {
        return new MetricName(name, "", "", Collections.emptyMap());
    }
}
