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

import kafka.test.ClusterConfig;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.server.common.Features;
import org.apache.kafka.server.common.GroupVersion;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static kafka.test.annotation.Type.CO_KRAFT;
import static kafka.test.annotation.Type.KRAFT;
import static kafka.test.annotation.Type.ZK;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.GROUP_COORDINATOR_REBALANCE_PROTOCOLS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.NEW_GROUP_COORDINATOR_ENABLE_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.OFFSETS_TOPIC_PARTITIONS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG;

/**
 * The old test framework {@link kafka.api.BaseConsumerTest#getTestQuorumAndGroupProtocolParametersAll} test for the following cases:
 * <ul>
 *     <li>(ZK / KRAFT servers) with (group.coordinator.new.enable=false) with (classic group protocol) = 2 cases</li>
 *     <li>(KRAFT server) with (group.coordinator.new.enable=true) with (classic group protocol) = 1 case</li>
 *     <li>(KRAFT server) with (group.coordinator.new.enable=true) with (consumer group protocol) = 1 case</li>
 * </ul>
 * <p>
 * The new test framework run seven cases for the following cases:
 * <ul>
 *     <li>(ZK / KRAFT / CO_KRAFT servers) with (group.coordinator.new.enable=false) with (classic group protocol) = 3 cases</li>
 *     <li>(KRAFT / CO_KRAFT servers) with (group.coordinator.new.enable=true) with (classic group protocol) = 2 cases</li>
 *     <li>(KRAFT / CO_KRAFT servers) with (group.coordinator.new.enable=true) with (consumer group protocol) = 2 cases</li>
 * </ul>
 * <p>
 * We can reduce the number of cases as same as the old test framework by using the following methods:
 * <ul>
 *     <li>{@link #forConsumerGroupCoordinator} for the case of (consumer group protocol)</li>
 *     <li>(CO_KRAFT servers) with (group.coordinator.new.enable=true) with (classic / consumer group protocols) = 2 cases</li>
 * </ul>
 * <ul>
 *     <li>{@link #forClassicGroupCoordinator} for the case of (classic group protocol)</li>
 *     <li>(ZK / KRAFT servers) with (group.coordinator.new.enable=false) with (classic group protocol) = 2 cases</li>
 * </ul>
 */
class ConsumerGroupCommandTestUtils {

    private ConsumerGroupCommandTestUtils() {
    }

    static List<ClusterConfig> generator() {
        return Stream.concat(forConsumerGroupCoordinator().stream(), forClassicGroupCoordinator().stream())
                .collect(Collectors.toList());
    }

    static List<ClusterConfig> forConsumerGroupCoordinator() {
        Map<String, String> serverProperties = new HashMap<>();
        serverProperties.put(OFFSETS_TOPIC_PARTITIONS_CONFIG, "1");
        serverProperties.put(OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
        serverProperties.put(NEW_GROUP_COORDINATOR_ENABLE_CONFIG, "true");
        serverProperties.put(GROUP_COORDINATOR_REBALANCE_PROTOCOLS_CONFIG, "classic,consumer");

        return Collections.singletonList(ClusterConfig.defaultBuilder()
                .setFeatures(Collections.singletonMap(Features.GROUP_VERSION, GroupVersion.GV_1.featureLevel()))
                .setTypes(Collections.singleton(CO_KRAFT))
                .setServerProperties(serverProperties)
                .setTags(Collections.singletonList("consumerGroupCoordinator"))
                .build());
    }

    static List<ClusterConfig> forClassicGroupCoordinator() {
        Map<String, String> serverProperties = new HashMap<>();
        serverProperties.put(OFFSETS_TOPIC_PARTITIONS_CONFIG, "1");
        serverProperties.put(OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
        serverProperties.put(NEW_GROUP_COORDINATOR_ENABLE_CONFIG, "false");

        return Collections.singletonList(ClusterConfig.defaultBuilder()
                .setTypes(Stream.of(ZK, KRAFT).collect(Collectors.toSet()))
                .setServerProperties(serverProperties)
                .setTags(Collections.singletonList("classicGroupCoordinator"))
                .build());
    }

    static <T> AutoCloseable buildConsumers(int numberOfConsumers,
                                            Set<TopicPartition> partitions,
                                            Supplier<KafkaConsumer<T, T>> consumerSupplier) {
        return buildConsumers(numberOfConsumers, false, consumerSupplier,
                consumer -> consumer.assign(partitions));
    }

    static <T> AutoCloseable buildConsumers(int numberOfConsumers,
                                            boolean syncCommit,
                                            String topic,
                                            Supplier<KafkaConsumer<T, T>> consumerSupplier) {
        return buildConsumers(numberOfConsumers, syncCommit, consumerSupplier,
                consumer -> consumer.subscribe(Collections.singleton(topic)));
    }

    private static <T> AutoCloseable buildConsumers(int numberOfConsumers,
                                                    boolean syncCommit,
                                                    Supplier<KafkaConsumer<T, T>> consumerSupplier,
                                                    Consumer<KafkaConsumer<T, T>> setPartitions) {
        List<KafkaConsumer<T, T>> consumers = new ArrayList<>(numberOfConsumers);
        ExecutorService executor = Executors.newFixedThreadPool(numberOfConsumers);
        AtomicBoolean closed = new AtomicBoolean(false);
        final AutoCloseable closeable = () -> releaseConsumers(closed, consumers, executor);
        try {
            for (int i = 0; i < numberOfConsumers; i++) {
                KafkaConsumer<T, T> consumer = consumerSupplier.get();
                consumers.add(consumer);
                executor.execute(() -> initConsumer(syncCommit, () -> {
                    setPartitions.accept(consumer);
                    return consumer;
                }, closed));
            }
            return closeable;
        } catch (Throwable e) {
            Utils.closeQuietly(closeable, "Release Consumer");
            throw e;
        }
    }

    private static <T> void releaseConsumers(AtomicBoolean closed, List<KafkaConsumer<T, T>> consumers, ExecutorService executor) throws InterruptedException {
        closed.set(true);
        consumers.forEach(KafkaConsumer::wakeup);
        executor.shutdown();
        executor.awaitTermination(1, TimeUnit.MINUTES);
    }

    private static <T> void initConsumer(boolean syncCommit,
                                         Supplier<KafkaConsumer<T, T>> consumerSupplier,
                                         AtomicBoolean closed) {
        try (KafkaConsumer<T, T> kafkaConsumer = consumerSupplier.get()) {
            while (!closed.get()) {
                kafkaConsumer.poll(Duration.ofMillis(Long.MAX_VALUE));
                if (syncCommit)
                    kafkaConsumer.commitSync();
            }
        } catch (WakeupException e) {
            // OK
        }
    }
}
