/*
 * 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;

import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.GroupProtocol;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.config.TopicConfig;
import org.apache.kafka.common.test.ClusterInstance;
import org.apache.kafka.common.test.TestUtils;
import org.apache.kafka.common.test.api.ClusterConfigProperty;
import org.apache.kafka.common.test.api.ClusterTest;
import org.apache.kafka.common.test.api.Type;
import org.apache.kafka.coordinator.group.GroupCoordinatorConfig;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;

public class ClientRebootstrapTest {
    private static final String TOPIC = "topic";
    private static final int PARTITIONS = 1;
    private static final int REPLICAS = 2;

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testAdminRebootstrap(ClusterInstance clusterInstance) {
        var broker0 = 0;
        var broker1 = 1;
        var timeout = 60;

        clusterInstance.shutdownBroker(broker0);

        try (var admin = clusterInstance.admin()) {
            admin.createTopics(List.of(new NewTopic(TOPIC, PARTITIONS, (short) REPLICAS)));

            // Only the broker 1 is available for the admin client during the bootstrap.
            assertDoesNotThrow(() -> admin.listTopics().names().get(timeout, TimeUnit.SECONDS).contains(TOPIC));

            clusterInstance.shutdownBroker(broker1);
            clusterInstance.startBroker(broker0);

            // The broker 1, originally cached during the bootstrap, is offline.
            // However, the broker 0 from the bootstrap list is online.
            // Should be able to list topics again.
            assertDoesNotThrow(() -> admin.listTopics().names().get(timeout, TimeUnit.SECONDS).contains(TOPIC));
        }
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testAdminRebootstrapDisabled(ClusterInstance clusterInstance) {
        var broker0 = 0;
        var broker1 = 1;

        clusterInstance.shutdownBroker(broker0);

        var admin = clusterInstance.admin(Map.of(CommonClientConfigs.METADATA_RECOVERY_STRATEGY_CONFIG, "none"));
        admin.createTopics(List.of(new NewTopic(TOPIC, PARTITIONS, (short) REPLICAS)));

        // Only the broker 1 is available for the admin client during the bootstrap.
        assertDoesNotThrow(() -> admin.listTopics().names().get(60, TimeUnit.SECONDS).contains(TOPIC));

        clusterInstance.shutdownBroker(broker1);
        clusterInstance.startBroker(broker0);

        // The broker 1, originally cached during the bootstrap, is offline.
        // As a result, the admin client will throw a TimeoutException when trying to get list of the topics.
        assertThrows(TimeoutException.class, () -> admin.listTopics().names().get(5, TimeUnit.SECONDS));
        // Since the brokers cached during the bootstrap are offline, the admin client needs to wait the default timeout for other threads.
        admin.close(Duration.ZERO);
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testProducerRebootstrap(ClusterInstance clusterInstance) throws ExecutionException, InterruptedException {
        try (var admin = clusterInstance.admin()) {
            admin.createTopics(List.of(new NewTopic(TOPIC, PARTITIONS, (short) REPLICAS)));
        }

        var broker0 = 0;
        var broker1 = 1;

        // It's ok to shut the leader down, cause the reelection is small enough to the producer timeout.
        clusterInstance.shutdownBroker(broker0);

        try (var producer = clusterInstance.producer()) {
            // Only the broker 1 is available for the producer during the bootstrap.
            var recordMetadata0 = producer.send(new ProducerRecord<>(TOPIC, "value 0".getBytes())).get();
            assertEquals(0, recordMetadata0.offset());

            clusterInstance.shutdownBroker(broker1);
            clusterInstance.startBroker(broker0);

            // Current broker 1 is offline.
            // However, the broker 0 from the bootstrap list is online.
            // Should be able to produce records.
            var recordMetadata1 = producer.send(new ProducerRecord<>(TOPIC, "value 1".getBytes())).get();
            assertEquals(0, recordMetadata1.offset());
        }
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testProducerRebootstrapDisabled(ClusterInstance clusterInstance) throws ExecutionException, InterruptedException {
        try (var admin = clusterInstance.admin()) {
            admin.createTopics(List.of(new NewTopic(TOPIC, PARTITIONS, (short) REPLICAS)));
        }

        var broker0 = 0;
        var broker1 = 1;

        // It's ok to shut the leader down, cause the reelection is small enough to the producer timeout.
        clusterInstance.shutdownBroker(broker0);

        var producer = clusterInstance.producer(Map.of(CommonClientConfigs.METADATA_RECOVERY_STRATEGY_CONFIG, "none"));

        // Only the broker 1 is available for the producer during the bootstrap.
        var recordMetadata0 = producer.send(new ProducerRecord<>(TOPIC, "value 0".getBytes())).get();
        assertEquals(0, recordMetadata0.offset());

        clusterInstance.shutdownBroker(broker1);
        clusterInstance.startBroker(broker0);

        // The broker 1, originally cached during the bootstrap, is offline.
        // As a result, the producer will throw a TimeoutException when trying to send a message.
        assertThrows(TimeoutException.class, () -> producer.send(new ProducerRecord<>(TOPIC, "value 1".getBytes())).get(5, TimeUnit.SECONDS));
        // Since the brokers cached during the bootstrap are offline, the producer needs to wait the default timeout for other threads.
        producer.close(Duration.ZERO);
    }

    public void consumerRebootstrap(ClusterInstance clusterInstance, GroupProtocol groupProtocol) throws InterruptedException, ExecutionException {
        clusterInstance.createTopic(TOPIC, PARTITIONS, (short) REPLICAS);

        var broker0 = 0;
        var broker1 = 1;
        var partitions = List.of(new TopicPartition(TOPIC, 0));

        try (var producer = clusterInstance.producer(Map.of(ProducerConfig.ACKS_CONFIG, "-1"))) {
            var recordMetadata = producer.send(new ProducerRecord<>(TOPIC, "value 0".getBytes())).get();
            assertEquals(0, recordMetadata.offset());
        }

        clusterInstance.shutdownBroker(broker0);

        try (var consumer = clusterInstance.consumer(Map.of(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name))) {
            // Only the server 1 is available for the consumer during the bootstrap.
            consumer.assign(partitions);
            consumer.seekToBeginning(partitions);
            TestUtils.waitForCondition(() -> consumer.poll(Duration.ofMillis(100)).count() == 1, 10 * 1000, "Failed to poll data.");

            // Bring back the server 0 and shut down 1.
            clusterInstance.shutdownBroker(broker1);
            clusterInstance.startBroker(broker0);

            try (var producer = clusterInstance.producer(Map.of(ProducerConfig.ACKS_CONFIG, "-1"))) {
                var recordMetadata = producer.send(new ProducerRecord<>(TOPIC, "value 1".getBytes())).get();
                assertEquals(1, recordMetadata.offset());
            }

            // The server 1 originally cached during the bootstrap, is offline.
            // However, the server 0 from the bootstrap list is online.
            TestUtils.waitForCondition(() -> consumer.poll(Duration.ofMillis(100)).count() == 1, 10 * 1000, "Failed to poll data.");
        }
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2"),
        })
    public void testClassicConsumerRebootstrap(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException {
        consumerRebootstrap(clusterInstance, GroupProtocol.CLASSIC);
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2"),
        })
    public void testConsumerRebootstrap(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException {
        consumerRebootstrap(clusterInstance, GroupProtocol.CONSUMER);
    }

    public void consumerRebootstrapDisabled(ClusterInstance clusterInstance, GroupProtocol groupProtocol) throws InterruptedException, ExecutionException {
        clusterInstance.createTopic(TOPIC, PARTITIONS, (short) REPLICAS);

        var broker0 = 0;
        var broker1 = 1;
        var tp = new TopicPartition(TOPIC, 0);

        try (var producer = clusterInstance.producer(Map.of(ProducerConfig.ACKS_CONFIG, "-1"))) {
            var recordMetadata = producer.send(new ProducerRecord<>(TOPIC, "value 0".getBytes())).get();
            assertEquals(0, recordMetadata.offset());
        }

        clusterInstance.shutdownBroker(broker0);

        try (var consumer = clusterInstance.consumer(Map.of(
            CommonClientConfigs.METADATA_RECOVERY_STRATEGY_CONFIG, "none",
            ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name)
        )) {
            // Only the server 1 is available for the consumer during the bootstrap.
            consumer.assign(List.of(tp));
            consumer.seekToBeginning(List.of(tp));
            TestUtils.waitForCondition(() -> consumer.poll(Duration.ofMillis(100)).count() == 1, 10 * 1000, "Failed to poll data.");

            // Bring back the server 0 and shut down 1.
            clusterInstance.shutdownBroker(broker1);
            clusterInstance.startBroker(broker0);

            try (var producer = clusterInstance.producer(Map.of(ProducerConfig.ACKS_CONFIG, "-1"))) {
                var recordMetadata = producer.send(new ProducerRecord<>(TOPIC, "value 1".getBytes())).get();
                assertEquals(1, recordMetadata.offset());
            }

            // The server 1 originally cached during the bootstrap, is offline.
            // However, the server 0 from the bootstrap list is online.
            assertEquals(0, consumer.poll(Duration.ofMillis(100)).count());
        }
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testClassicConsumerRebootstrapDisabled(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException {
        consumerRebootstrapDisabled(clusterInstance, GroupProtocol.CLASSIC);
    }

    @ClusterTest(
        brokers = REPLICAS,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG, value = "true"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "2")
        }
    )
    public void testConsumerRebootstrapDisabled(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException {
        consumerRebootstrapDisabled(clusterInstance, GroupProtocol.CONSUMER);
    }
}
