/*
 * 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.connect.runtime.errors;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.errors.RetriableException;
import org.apache.kafka.connect.runtime.ConnectMetrics;
import org.apache.kafka.connect.runtime.ConnectorConfig;
import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.runtime.isolation.Plugins;
import org.apache.kafka.connect.runtime.isolation.PluginsTest.TestConverter;
import org.apache.kafka.connect.runtime.isolation.PluginsTest.TestableWorkerConfig;
import org.apache.kafka.connect.sink.SinkConnector;
import org.apache.kafka.connect.sink.SinkTask;
import org.apache.kafka.connect.source.SourceRecord;
import org.apache.kafka.connect.util.ConnectorTaskId;

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 org.mockito.stubbing.OngoingStubbing;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;

import static org.apache.kafka.common.utils.Time.SYSTEM;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_MAX_DELAY_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_MAX_DELAY_DEFAULT;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_TIMEOUT_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_TIMEOUT_DEFAULT;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_DEFAULT;
import static org.apache.kafka.connect.runtime.errors.ToleranceType.ALL;
import static org.apache.kafka.connect.runtime.errors.ToleranceType.NONE;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

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

    private static final Map<String, String> PROPERTIES = new HashMap<>() {{
            put(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG, Objects.toString(2));
            put(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG, Objects.toString(3000));
            put(CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG, Sensor.RecordingLevel.INFO.toString());

            // define required properties
            put(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
            put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
            put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
        }};

    public static <T> RetryWithToleranceOperator<T> noneOperator() {
        return genericOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, NONE, new ErrorHandlingMetrics(
                new ConnectorTaskId("errors-none-tolerate-connector", -1),
                new ConnectMetrics("errors-none-tolerate-worker", new TestableWorkerConfig(PROPERTIES),
                        Time.SYSTEM, "test-cluster")));
    }

    public static <T> RetryWithToleranceOperator<T> allOperator() {
        return genericOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, ALL, new ErrorHandlingMetrics(
                new ConnectorTaskId("errors-all-tolerate-connector", -1),
                new ConnectMetrics("errors-all-tolerate-worker", new TestableWorkerConfig(PROPERTIES),
                        Time.SYSTEM, "test-cluster")));
    }

    private static <T> RetryWithToleranceOperator<T> genericOperator(int retryTimeout, ToleranceType toleranceType, ErrorHandlingMetrics metrics) {
        return new RetryWithToleranceOperator<>(retryTimeout, ERRORS_RETRY_MAX_DELAY_DEFAULT, toleranceType, SYSTEM, metrics);
    }

    @Mock
    private Operation<String> mockOperation;

    @Mock
    private ConsumerRecord<byte[], byte[]> consumerRecord;

    @Mock
    ErrorHandlingMetrics errorHandlingMetrics;

    @Mock
    Plugins plugins;

    @Test
    public void testExecuteFailed() {
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = genericOperator(0, ALL, errorHandlingMetrics);

        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        retryWithToleranceOperator.executeFailed(context, Stage.TASK_PUT,
            SinkTask.class, new Throwable());
    }

    @Test
    public void testExecuteFailedNoTolerance() {
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = genericOperator(0, NONE, errorHandlingMetrics);

        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        assertThrows(ConnectException.class, () -> retryWithToleranceOperator.executeFailed(context, Stage.TASK_PUT,
            SinkTask.class, new Throwable()));
    }

    @Test
    public void testHandleExceptionInTransformations() {
        testHandleExceptionInStage(Stage.TRANSFORMATION, new Exception(), ALL);
    }

    @Test
    public void testHandleRetriableExceptionInTransformationsToleranceNone() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.TRANSFORMATION, new RetriableException("Test"), NONE));
    }


    @Test
    public void testHandleExceptionInHeaderConverter() {
        testHandleExceptionInStage(Stage.HEADER_CONVERTER, new Exception(), ALL);
    }

    @Test
    public void testHandleRetriableExceptionInHeaderConverterToleranceNone() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.HEADER_CONVERTER, new RetriableException("Test"), NONE));
    }

    @Test
    public void testHandleExceptionInValueConverter() {
        testHandleExceptionInStage(Stage.VALUE_CONVERTER, new Exception(), ALL);
    }

    @Test
    public void testHandleRetriableExceptionInValueConverterToleranceNone() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.VALUE_CONVERTER, new RetriableException("Test"), NONE));
    }

    @Test
    public void testHandleExceptionInKeyConverter() {
        testHandleExceptionInStage(Stage.KEY_CONVERTER, new Exception(), ALL);
    }

    @Test
    public void testHandleRetriableExceptionInKeyConverterToleranceNone() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.KEY_CONVERTER, new RetriableException("Test"), NONE));
    }

    @Test
    public void testHandleExceptionInTaskPut() {
        testHandleExceptionInStage(Stage.TASK_PUT, new org.apache.kafka.connect.errors.RetriableException("Test"), ALL);
    }

    @Test
    public void testHandleExceptionInTaskPoll() {
        testHandleExceptionInStage(Stage.TASK_POLL, new org.apache.kafka.connect.errors.RetriableException("Test"), ALL);
    }

    @Test
    public void testThrowExceptionInTaskPut() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.TASK_PUT, new Exception(), ALL));
    }

    @Test
    public void testThrowExceptionInTaskPoll() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.TASK_POLL, new Exception(), ALL));
    }

    @Test
    public void testThrowExceptionInKafkaConsume() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.KAFKA_CONSUME, new Exception(), ALL));
    }

    @Test
    public void testThrowExceptionInKafkaProduce() {
        assertThrows(ConnectException.class, () -> testHandleExceptionInStage(Stage.KAFKA_PRODUCE, new Exception(), ALL));
    }

    private void testHandleExceptionInStage(Stage type, Exception ex, ToleranceType toleranceType) {
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = setupExecutor(toleranceType);
        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        Operation<?> exceptionThrower = () -> {
            throw ex;
        };
        retryWithToleranceOperator.execute(context, exceptionThrower, type, RetryWithToleranceOperator.class);
        assertTrue(context.failed());
    }

    private <T> RetryWithToleranceOperator<T> setupExecutor(ToleranceType toleranceType) {
        return genericOperator(0, toleranceType, errorHandlingMetrics);
    }

    @Test
    public void testExecAndHandleRetriableErrorOnce() throws Exception {
        execAndHandleRetriableError(6000, 1, List.of(300L), new RetriableException("Test"), true);
    }

    @Test
    public void testExecAndHandleRetriableErrorThrice() throws Exception {
        execAndHandleRetriableError(6000, 3, List.of(300L, 600L, 1200L), new RetriableException("Test"), true);
    }

    @Test
    public void testExecAndHandleRetriableErrorWithInfiniteRetries() throws Exception {
        execAndHandleRetriableError(-1, 8, List.of(300L, 600L, 1200L, 2400L, 4800L, 9600L, 19200L, 38400L), new RetriableException("Test"), true);
    }

    @Test
    public void testExecAndHandleRetriableErrorWithMaxRetriesExceeded() throws Exception {
        execAndHandleRetriableError(6000, 6, List.of(300L, 600L, 1200L, 2400L, 1500L), new RetriableException("Test"), false);
    }

    public void execAndHandleRetriableError(long errorRetryTimeout, int numRetriableExceptionsThrown, List<Long> expectedWaits, Exception e, boolean successExpected) throws Exception {
        MockTime time = new MockTime(0, 0, 0);
        CountDownLatch exitLatch = mock(CountDownLatch.class);
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = new RetryWithToleranceOperator<>(errorRetryTimeout, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time, errorHandlingMetrics, exitLatch);

        OngoingStubbing<String> mockOperationCall = when(mockOperation.call());
        for (int i = 0; i < numRetriableExceptionsThrown; i++) {
            mockOperationCall = mockOperationCall.thenThrow(e);
        }
        if (successExpected) {
            mockOperationCall.thenReturn("Success");
        }

        for (Long expectedWait : expectedWaits) {
            when(exitLatch.await(expectedWait, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
                time.sleep(expectedWait);
                return false;
            });
        }

        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        String result = retryWithToleranceOperator.execAndHandleError(context, mockOperation, Exception.class);

        if (successExpected) {
            assertFalse(context.failed());
            assertEquals("Success", result);
        } else {
            assertTrue(context.failed());
        }

        verifyNoMoreInteractions(exitLatch);
        verify(mockOperation, times(successExpected ? numRetriableExceptionsThrown + 1 : numRetriableExceptionsThrown)).call();
    }

    @Test
    public void testExecAndHandleNonRetriableError() throws Exception {
        MockTime time = new MockTime(0, 0, 0);
        CountDownLatch exitLatch = mock(CountDownLatch.class);
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = new RetryWithToleranceOperator<>(6000, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time, errorHandlingMetrics, exitLatch);

        when(mockOperation.call()).thenThrow(new Exception("Test"));

        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        String result = retryWithToleranceOperator.execAndHandleError(context, mockOperation, Exception.class);
        assertTrue(context.failed());
        assertNull(result);

        // expect no call to exitLatch.await() which is only called during the retry backoff
        verify(mockOperation).call();
        verify(exitLatch, never()).await(anyLong(), any());
    }

    @Test
    public void testExitLatch() throws Exception {
        MockTime time = new MockTime(0, 0, 0);
        CountDownLatch exitLatch = mock(CountDownLatch.class);
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = new RetryWithToleranceOperator<>(-1, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time, errorHandlingMetrics, exitLatch);
        when(mockOperation.call()).thenThrow(new RetriableException("test"));

        when(exitLatch.await(300L, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(300);
            return false;
        });
        when(exitLatch.await(600L, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(600);
            return false;
        });
        when(exitLatch.await(1200L, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(1200);
            return false;
        });
        when(exitLatch.await(2400L, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(2400);
            retryWithToleranceOperator.triggerStop();
            return false;
        });

        // expect no more calls to exitLatch.await() after retryWithToleranceOperator.triggerStop() is called

        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        retryWithToleranceOperator.execAndHandleError(context, mockOperation, Exception.class);
        assertTrue(context.failed());
        assertEquals(4500L, time.milliseconds());
        verify(exitLatch).countDown();
        verifyNoMoreInteractions(exitLatch);
    }

    @Test
    public void testBackoffLimit() throws Exception {
        MockTime time = new MockTime(0, 0, 0);
        CountDownLatch exitLatch = mock(CountDownLatch.class);
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = new RetryWithToleranceOperator<>(5, 5000, NONE, time, errorHandlingMetrics, exitLatch);

        when(exitLatch.await(300, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(300);
            return false;
        });
        when(exitLatch.await(600, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(600);
            return false;
        });
        when(exitLatch.await(1200, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(1200);
            return false;
        });
        when(exitLatch.await(2400, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(2400);
            return false;
        });
        when(exitLatch.await(500, TimeUnit.MILLISECONDS)).thenAnswer(i -> {
            time.sleep(500);
            return false;
        });
        when(exitLatch.await(0, TimeUnit.MILLISECONDS)).thenReturn(false);

        retryWithToleranceOperator.backoff(1, 5000);
        retryWithToleranceOperator.backoff(2, 5000);
        retryWithToleranceOperator.backoff(3, 5000);
        retryWithToleranceOperator.backoff(4, 5000);
        retryWithToleranceOperator.backoff(5, 5000);
        // Simulate a small delay between calculating the deadline, and backing off
        time.sleep(1);
        // We may try to begin backing off after the deadline has already passed; make sure
        // that we don't wait with a negative timeout
        retryWithToleranceOperator.backoff(6, 5000);

        verifyNoMoreInteractions(exitLatch);
    }

    @Test
    public void testToleranceLimit() {
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = genericOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, NONE, errorHandlingMetrics);
        retryWithToleranceOperator.markAsFailed();
        assertFalse(retryWithToleranceOperator.withinToleranceLimits(), "should not tolerate any errors");

        retryWithToleranceOperator = genericOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, ALL, errorHandlingMetrics);
        retryWithToleranceOperator.markAsFailed();
        retryWithToleranceOperator.markAsFailed();
        assertTrue(retryWithToleranceOperator.withinToleranceLimits(), "should tolerate all errors");

        retryWithToleranceOperator = genericOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, NONE, errorHandlingMetrics);
        assertTrue(retryWithToleranceOperator.withinToleranceLimits(), "no tolerance is within limits if no failures");
    }

    @Test
    public void testDefaultConfigs() {
        ConnectorConfig configuration = config(Map.of());
        assertEquals(ERRORS_RETRY_TIMEOUT_DEFAULT, configuration.errorRetryTimeout());
        assertEquals(ERRORS_RETRY_MAX_DELAY_DEFAULT, configuration.errorMaxDelayInMillis());
        assertEquals(ERRORS_TOLERANCE_DEFAULT, configuration.errorToleranceType());
    }

    ConnectorConfig config(Map<String, String> connProps) {
        Map<String, String> props = new HashMap<>();
        props.put(ConnectorConfig.NAME_CONFIG, "test");
        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, SinkConnector.class.getName());
        props.putAll(connProps);
        return new ConnectorConfig(plugins, props);
    }

    @Test
    public void testSetConfigs() {
        ConnectorConfig configuration;
        configuration = config(Map.of(ERRORS_RETRY_TIMEOUT_CONFIG, "100"));
        assertEquals(100, configuration.errorRetryTimeout());

        configuration = config(Map.of(ERRORS_RETRY_MAX_DELAY_CONFIG, "100"));
        assertEquals(100, configuration.errorMaxDelayInMillis());

        configuration = config(Map.of(ERRORS_TOLERANCE_CONFIG, "none"));
        assertEquals(ToleranceType.NONE, configuration.errorToleranceType());
    }

    @Test
    public void testReportWithSingleReporter() {
        testReport(1);
    }

    @Test
    public void testReportWithMultipleReporters() {
        testReport(2);
    }

    private void testReport(int numberOfReports) {
        MockTime time = new MockTime(0, 0, 0);
        CountDownLatch exitLatch = mock(CountDownLatch.class);
        RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator = new RetryWithToleranceOperator<>(-1, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time, errorHandlingMetrics, exitLatch);
        ConsumerRecord<byte[], byte[]> consumerRecord = new ConsumerRecord<>("t", 0, 0, null, null);
        List<CompletableFuture<RecordMetadata>> fs = IntStream.range(0, numberOfReports).mapToObj(i -> new CompletableFuture<RecordMetadata>()).toList();
        List<ErrorReporter<ConsumerRecord<byte[], byte[]>>> reporters = IntStream.range(0, numberOfReports).mapToObj(i -> (ErrorReporter<ConsumerRecord<byte[], byte[]>>) c -> fs.get(i)).toList();
        retryWithToleranceOperator.reporters(reporters);
        ProcessingContext<ConsumerRecord<byte[], byte[]>> context = new ProcessingContext<>(consumerRecord);
        Future<Void> result = retryWithToleranceOperator.report(context);
        fs.forEach(f -> {
            assertFalse(result.isDone());
            f.complete(new RecordMetadata(new TopicPartition("t", 0), 0, 0, 0, 0, 0));
        });
        assertTrue(result.isDone());
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testCloseErrorReporters() {
        ErrorReporter<SourceRecord> reporterA = mock(ErrorReporter.class);
        ErrorReporter<SourceRecord> reporterB = mock(ErrorReporter.class);

        RetryWithToleranceOperator<SourceRecord> retryWithToleranceOperator = allOperator();

        retryWithToleranceOperator.reporters(List.of(reporterA, reporterB));

        // Even though the reporters throw exceptions, they should both still be closed.

        retryWithToleranceOperator.close();
        verify(reporterA).close();
        verify(reporterB).close();
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testCloseErrorReportersExceptionPropagation() {
        ErrorReporter<SourceRecord> reporterA = mock(ErrorReporter.class);
        ErrorReporter<SourceRecord> reporterB = mock(ErrorReporter.class);

        RetryWithToleranceOperator<SourceRecord> retryWithToleranceOperator = allOperator();

        retryWithToleranceOperator.reporters(List.of(reporterA, reporterB));

        // Even though the reporters throw exceptions, they should both still be closed.
        doThrow(new RuntimeException()).when(reporterA).close();
        doThrow(new RuntimeException()).when(reporterB).close();

        assertThrows(ConnectException.class, retryWithToleranceOperator::close);
        verify(reporterA).close();
        verify(reporterB).close();
    }
}
