/*
 * 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.integration.utils;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.AlterConfigOp;
import org.apache.kafka.clients.admin.Config;
import org.apache.kafka.clients.admin.ConfigEntry;
import org.apache.kafka.clients.admin.ListTopicsOptions;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.GroupProtocol;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.config.ConfigResource;
import org.apache.kafka.common.config.SslConfigs;
import org.apache.kafka.common.config.types.Password;
import org.apache.kafka.common.errors.InvalidReplicationFactorException;
import org.apache.kafka.common.errors.TopicExistsException;
import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.test.KafkaClusterTestKit;
import org.apache.kafka.common.test.TestKitNodes;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.coordinator.group.GroupConfig;
import org.apache.kafka.coordinator.group.GroupCoordinatorConfig;
import org.apache.kafka.coordinator.transaction.TransactionLogConfig;
import org.apache.kafka.network.SocketServerConfigs;
import org.apache.kafka.server.config.ServerConfigs;
import org.apache.kafka.server.config.ServerLogConfigs;
import org.apache.kafka.server.util.MockTime;
import org.apache.kafka.storage.internals.log.CleanerConfig;
import org.apache.kafka.test.TestCondition;
import org.apache.kafka.test.TestUtils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_PROTOCOL_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
import static org.apache.kafka.common.utils.Utils.mkProperties;

/**
 * Setup an embedded Kafka KRaft cluster for integration tests (using {@link org.apache.kafka.common.test.KafkaClusterTestKit} internally) with the
 * specified number of brokers and the specified broker properties.
 * Additional Kafka client properties can also be supplied if required.
 * This class also provides various utility methods to easily create Kafka topics, produce data, consume data etc.
 */
public class EmbeddedKafkaCluster {

    private static final Logger log = LoggerFactory.getLogger(EmbeddedKafkaCluster.class);
    private final KafkaClusterTestKit cluster;
    private final Properties brokerConfig;
    public final MockTime time;
    public EmbeddedKafkaCluster(final int numBrokers) {
        this(numBrokers, new Properties());
    }

    public EmbeddedKafkaCluster(final int numBrokers, final Properties brokerConfig) {
        this(numBrokers, brokerConfig, Collections.emptyMap());
    }

    public EmbeddedKafkaCluster(final int numBrokers,
                                final Properties brokerConfig,
                                final long mockTimeMillisStart) {
        this(numBrokers, brokerConfig, Collections.emptyMap(), mockTimeMillisStart, System.nanoTime());
    }
    public EmbeddedKafkaCluster(final int numBrokers,
                                final Properties brokerConfig,
                                final Map<Integer, Map<String, String>> brokerConfigOverrides) {
        this(numBrokers, brokerConfig, brokerConfigOverrides, System.currentTimeMillis(), System.nanoTime());
    }
    public EmbeddedKafkaCluster(final int numBrokers,
                                final Properties brokerConfig,
                                final Map<Integer, Map<String, String>> brokerConfigOverrides,
                                final long mockTimeMillisStart,
                                final long mockTimeNanoStart) {
        addDefaultBrokerPropsIfAbsent(brokerConfig);

        if (!brokerConfigOverrides.isEmpty() && brokerConfigOverrides.size() != numBrokers) {
            throw new IllegalArgumentException("Size of brokerConfigOverrides " + brokerConfigOverrides.size()
                    + " must match broker number " + numBrokers);
        }
        try {
            final KafkaClusterTestKit.Builder clusterBuilder = new KafkaClusterTestKit.Builder(
                    new TestKitNodes.Builder()
                            .setCombined(true)
                            .setNumBrokerNodes(numBrokers)
                            .setPerServerProperties(brokerConfigOverrides)
                            // Reduce number of controllers for faster startup
                            // We may make this configurable in the future if there's a use case for it
                            .setNumControllerNodes(1)
                            .build()
            );

            brokerConfig.forEach((k, v) -> clusterBuilder.setConfigProp((String) k, v));
            cluster = clusterBuilder.build();
            cluster.nonFatalFaultHandler().setIgnore(true);
        } catch (final Exception e) {
            throw new KafkaException("Failed to create test Kafka cluster", e);
        }
        this.brokerConfig = brokerConfig;
        this.time = new MockTime(mockTimeMillisStart, mockTimeNanoStart);
    }

    public void start() {
        try {
            cluster.format();
            cluster.startup();
            cluster.waitForReadyBrokers();
        } catch (final Exception e) {
            throw new KafkaException("Failed to start test Kafka cluster", e);
        }

        verifyClusterReadiness();
    }

    /**
     * Perform an extended check to ensure that the primary APIs of the cluster are available, including:
     * <ul>
     *     <li>Ability to create a topic</li>
     *     <li>Ability to produce to a topic</li>
     *     <li>Ability to form a consumer group</li>
     *     <li>Ability to consume from a topic</li>
     * </ul>
     * If this method completes successfully, all resources created to verify the cluster health
     * (such as topics and consumer groups) will be cleaned up before it returns.
     * <p>
     * This provides extra guarantees compared to other cluster readiness checks such as
     * {@link KafkaClusterTestKit#waitForReadyBrokers()}, which verify that brokers have
     * completed startup and joined the cluster, but do not verify that the internal consumer
     * offsets topic has been created or that it's actually possible for users to create and
     * interact with topics.
     */
    public void verifyClusterReadiness() {
        final UUID uuid = UUID.randomUUID();
        final String consumerGroupId = "group-warmup-" + uuid;
        final Map<String, Object> consumerConfig = Map.of(
            GROUP_ID_CONFIG, consumerGroupId,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name()
        );
        final String topic = "topic-warmup-" + uuid;

        createTopic(topic);
        final Map<String, Object> producerProps = new HashMap<>(clientDefaultConfig());
        producerProps.put(ProducerConfig.CLIENT_ID_CONFIG, "warmup-producer");
        produce(producerProps, topic, null, "warmup message key", "warmup message value");

        try (Consumer<?, ?> consumer = createConsumerAndSubscribeTo(consumerConfig, topic)) {
            final ConsumerRecords<?, ?> records = consumer.poll(Duration.ofMillis(TimeUnit.MINUTES.toMillis(2)));
            if (records.isEmpty()) {
                throw new AssertionError("Failed to verify availability of group coordinator and produce/consume APIs on Kafka cluster in time");
            }
        }

        try (Admin admin = createAdminClient()) {
            admin.deleteConsumerGroups(Collections.singleton(consumerGroupId)).all().get(30, TimeUnit.SECONDS);
            admin.deleteTopics(Collections.singleton(topic)).all().get(30, TimeUnit.SECONDS);
        } catch (final InterruptedException | ExecutionException | TimeoutException e) {
            throw new AssertionError("Failed to clean up cluster health check resource(s)", e);
        }
    }

    /**
     * Stop the Kafka cluster.
     */
    public void stop() {
        final AtomicReference<Throwable> shutdownFailure = new AtomicReference<>();
        Utils.closeQuietly(cluster, "embedded Kafka cluster", shutdownFailure);
        if (shutdownFailure.get() != null) {
            throw new KafkaException("Failed to shut down producer / embedded Kafka cluster", shutdownFailure.get());
        }
    }

    public String bootstrapServers() {
        return cluster.bootstrapServers();
    }

    public boolean sslEnabled() {
        final String listenerSecurityProtocolMap = brokerConfig.getProperty(SocketServerConfigs.LISTENER_SECURITY_PROTOCOL_MAP_CONFIG);
        if (listenerSecurityProtocolMap == null)
            return false;
        return listenerSecurityProtocolMap.contains(":SSL") || listenerSecurityProtocolMap.contains(":SASL_SSL");
    }

    /**
     * Create multiple Kafka topics each with 1 partition and a replication factor of 1.
     *
     * @param topics The name of the topics.
     */
    public void createTopics(final String... topics) throws InterruptedException {
        for (final String topic : topics) {
            createTopic(topic, 1, 1, Collections.emptyMap());
        }
    }

    /**
     * Create a Kafka topic with 1 partition and a replication factor of 1.
     *
     * @param topic The name of the topic.
     */
    public void createTopic(final String topic) {
        createTopic(topic, 1);
    }

    /**
     * Create a Kafka topic with given partition and a replication factor of 1.
     *
     * @param topic The name of the topic.
     * @param partitions  The number of partitions for this topic.
     */
    public void createTopic(final String topic, final int partitions) {
        createTopic(topic, partitions, 1, Collections.emptyMap());
    }

    /**
     * Create a Kafka topic with the given parameters.
     *
     * @param topic       The name of the topic.
     * @param partitions  The number of partitions for this topic.
     * @param replication The replication factor for (the partitions of) this topic.
     */
    public void createTopic(final String topic, final int partitions, final int replication) throws InterruptedException {
        createTopic(topic, partitions, replication, Collections.emptyMap());
    }

    /**
     * Create a Kafka topic with given partition, replication factor, and topic config.
     *
     * @param topic The name of the topic.
     * @param partitions  The number of partitions for this topic.
     * @param replication The replication factor for (partitions of) this topic.
     * @param topicConfig Additional topic-level configuration settings.
     */
    public void createTopic(final String topic, final int partitions, final int replication, final Map<String, String> topicConfig) {
        if (replication > cluster.brokers().size()) {
            throw new InvalidReplicationFactorException("Insufficient brokers ("
                    + cluster.brokers().size() + ") for desired replication (" + replication + ")");
        }

        log.info("Creating topic { name: {}, partitions: {}, replication: {}, config: {} }",
                topic, partitions, replication, topicConfig);
        final NewTopic newTopic = new NewTopic(topic, partitions, (short) replication);
        newTopic.configs(topicConfig);

        try (final Admin adminClient = createAdminClient()) {
            adminClient.createTopics(Collections.singletonList(newTopic)).all().get();
            TestUtils.waitForCondition(() -> adminClient.listTopics().names().get().contains(topic),
                    "Wait for topic " + topic + " to get created.");
        } catch (final TopicExistsException ignored) {
        } catch (final InterruptedException | ExecutionException e) {
            if (!(e.getCause() instanceof TopicExistsException)) {
                throw new RuntimeException(e);
            }
        }
    }

    public void deleteTopics(final String... topics) {
        for (final String topic : topics) {
            deleteTopic(topic);
        }
    }


    /**
     * Delete a Kafka topic.
     *
     * @param topic the topic to delete; may not be null
     */
    public void deleteTopic(final String topic) {
        try (final Admin adminClient = createAdminClient()) {
            adminClient.deleteTopics(Collections.singleton(topic)).all().get();
        } catch (final InterruptedException | ExecutionException e) {
            if (!(e.getCause() instanceof UnknownTopicOrPartitionException)) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * Delete all topics except internal topics.
     */
    public void deleteAllTopics() {
        try (final Admin adminClient = createAdminClient()) {
            final Set<String> topics = adminClient.listTopics().names().get();
            adminClient.deleteTopics(topics).all().get();
        } catch (final UnknownTopicOrPartitionException ignored) {
        } catch (final ExecutionException | InterruptedException e) {
            if (!(e.getCause() instanceof UnknownTopicOrPartitionException)) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * Produce given key and value to topic partition.
     * @param topic the topic to produce to; may not be null.
     * @param partition the topic partition to produce to.
     * @param key the record key.
     * @param value the record value.
     */
    public void produce(final Map<String, Object> producerProps, final String topic, final Integer partition, final String key, final String value) {
        try (KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(producerProps, new ByteArraySerializer(), new ByteArraySerializer())) {
            final ProducerRecord<byte[], byte[]> msg = new ProducerRecord<>(topic, partition, key == null ? null : key.getBytes(), value == null ? null : value.getBytes());
            try {
                producer.send(msg).get(TimeUnit.SECONDS.toMillis(120), TimeUnit.MILLISECONDS);
                producer.flush();
            } catch (final Exception e) {
                throw new KafkaException("Could not produce message: " + msg, e);
            }
        }
    }

    public Admin createAdminClient() {
        return Admin.create(mkProperties(clientDefaultConfig()));
    }

    public Map<String, String> clientDefaultConfig() {
        final Map<String, String> props = new HashMap<>();
        props.putIfAbsent(BOOTSTRAP_SERVERS_CONFIG, bootstrapServers());
        if (sslEnabled()) {
            props.putIfAbsent(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, brokerConfig.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG).toString());
            props.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, ((Password) brokerConfig.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG)).value());
            props.putIfAbsent(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SSL");
        }
        return props;
    }

    public KafkaConsumer<byte[], byte[]> createConsumer(final Map<String, Object> consumerProps) {
        final Map<String, Object> props = new HashMap<>(clientDefaultConfig());
        props.putAll(consumerProps);

        props.putIfAbsent(GROUP_ID_CONFIG, UUID.randomUUID().toString());
        props.putIfAbsent(ENABLE_AUTO_COMMIT_CONFIG, "false");
        props.putIfAbsent(AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.putIfAbsent(KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
        props.putIfAbsent(VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");

        final KafkaConsumer<byte[], byte[]> consumer;
        try {
            consumer = new KafkaConsumer<>(props);
        } catch (final Throwable t) {
            throw new KafkaException("Failed to create consumer", t);
        }
        return consumer;
    }

    public KafkaConsumer<byte[], byte[]> createConsumerAndSubscribeTo(final Map<String, Object> consumerProps, final String... topics) {
        return createConsumerAndSubscribeTo(consumerProps, null, topics);
    }

    public KafkaConsumer<byte[], byte[]> createConsumerAndSubscribeTo(final Map<String, Object> consumerProps, final ConsumerRebalanceListener rebalanceListener, final String... topics) {
        final KafkaConsumer<byte[], byte[]> consumer = createConsumer(consumerProps);
        if (rebalanceListener != null) {
            consumer.subscribe(Arrays.asList(topics), rebalanceListener);
        } else {
            consumer.subscribe(Arrays.asList(topics));
        }
        return consumer;
    }

    private void addDefaultBrokerPropsIfAbsent(final Properties brokerConfig) {
        brokerConfig.putIfAbsent(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, 2 * 1024 * 1024L);
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.STREAMS_GROUP_MIN_SESSION_TIMEOUT_MS_CONFIG, "100");
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.STREAMS_GROUP_MIN_HEARTBEAT_INTERVAL_MS_CONFIG, "100");
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.GROUP_MIN_SESSION_TIMEOUT_MS_CONFIG, "0");
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.GROUP_INITIAL_REBALANCE_DELAY_MS_CONFIG, "0");
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.OFFSETS_TOPIC_PARTITIONS_CONFIG, "5");
        brokerConfig.putIfAbsent(GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
        brokerConfig.putIfAbsent(TransactionLogConfig.TRANSACTIONS_TOPIC_PARTITIONS_CONFIG, "5");
        brokerConfig.putIfAbsent(TransactionLogConfig.TRANSACTIONS_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
        brokerConfig.putIfAbsent(ServerLogConfigs.AUTO_CREATE_TOPICS_ENABLE_CONFIG, true);
        brokerConfig.putIfAbsent(ServerConfigs.DELETE_TOPIC_ENABLE_CONFIG, true);
    }

    public void waitForRemainingTopics(final long timeoutMs, final String... topics) throws InterruptedException {
        TestUtils.waitForCondition(new TopicsRemainingCondition(topics), timeoutMs, "Topics are not expected after " + timeoutMs + " milliseconds.");
    }

    public Set<String> getAllTopicsInCluster() {
        try (final Admin adminClient = createAdminClient()) {
            return adminClient.listTopics(new ListTopicsOptions().listInternal(true)).names().get();
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public Properties getLogConfig(final String topic) {
        try (final Admin adminClient = createAdminClient()) {
            final ConfigResource configResource = new ConfigResource(ConfigResource.Type.TOPIC, topic);
            final Config config = adminClient.describeConfigs(Collections.singleton(configResource)).values().get(configResource).get();
            final Properties properties = new Properties();
            for (final ConfigEntry configEntry : config.entries()) {
                if (configEntry.source() == ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG) {
                    properties.put(configEntry.name(), configEntry.value());
                }
            }
            return properties;
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public void setGroupSessionTimeout(final String groupId, final int sessionTimeoutMs) {
        try (final Admin adminClient = createAdminClient()) {
            adminClient.incrementalAlterConfigs(
                Map.of(
                    new ConfigResource(ConfigResource.Type.GROUP, groupId),
                    List.of(new AlterConfigOp(new ConfigEntry(GroupConfig.STREAMS_SESSION_TIMEOUT_MS_CONFIG, String.valueOf(sessionTimeoutMs)), AlterConfigOp.OpType.SET))
                )
            ).all().get();
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public void setGroupHeartbeatTimeout(final String groupId, final int heartbeatTimeoutMs) {
        try (final Admin adminClient = createAdminClient()) {
            adminClient.incrementalAlterConfigs(
                Map.of(
                    new ConfigResource(ConfigResource.Type.GROUP, groupId),
                    List.of(new AlterConfigOp(new ConfigEntry(GroupConfig.STREAMS_HEARTBEAT_INTERVAL_MS_CONFIG, String.valueOf(heartbeatTimeoutMs)), AlterConfigOp.OpType.SET))
                )
            ).all().get();
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public void setGroupStandbyReplicas(final String groupId, final int numStandbyReplicas) {
        try (final Admin adminClient = createAdminClient()) {
            adminClient.incrementalAlterConfigs(
                Map.of(
                    new ConfigResource(ConfigResource.Type.GROUP, groupId),
                    List.of(new AlterConfigOp(new ConfigEntry(GroupConfig.STREAMS_NUM_STANDBY_REPLICAS_CONFIG, String.valueOf(numStandbyReplicas)), AlterConfigOp.OpType.SET))
                )
            ).all().get();
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public void setGroupStreamsInitialRebalanceDelay(final String groupId, final int initialRebalanceDelayMs) {
        try (final Admin adminClient = createAdminClient()) {
            adminClient.incrementalAlterConfigs(
                Map.of(
                    new ConfigResource(ConfigResource.Type.GROUP, groupId),
                    List.of(new AlterConfigOp(new ConfigEntry(GroupConfig.STREAMS_INITIAL_REBALANCE_DELAY_MS_CONFIG, String.valueOf(initialRebalanceDelayMs)), AlterConfigOp.OpType.SET))
                )
            ).all().get();
        } catch (final InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    private final class TopicsRemainingCondition implements TestCondition {
        final Set<String> remainingTopics = new HashSet<>();

        private TopicsRemainingCondition(final String... topics) {
            Collections.addAll(remainingTopics, topics);
        }

        @Override
        public boolean conditionMet() {
            final Set<String> allTopics = getAllTopicsInCluster();
            return allTopics.equals(remainingTopics);
        }
    }
}
