/*
 * 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.admin.AdminClient;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment;
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.GroupSubscription;
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.StreamsConfig.InternalConfig;
import org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo;
import org.apache.kafka.streams.processor.internals.assignment.FallbackPriorTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.HighAvailabilityTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.LegacyStickyTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.LegacyTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.ReferenceContainer;
import org.apache.kafka.test.MockApiProcessorSupplier;
import org.apache.kafka.test.MockClientSupplier;
import org.apache.kafka.test.MockInternalTopicManager;
import org.apache.kafka.test.MockKeyValueStoreBuilder;

import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static java.util.Arrays.asList;
import static java.util.Collections.emptySet;
import static java.util.Collections.singletonList;
import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_TASKS;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.createMockAdminClientForAssignor;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getInfo;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.processIdForInt;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.mockito.ArgumentMatchers.anySet;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

@Tag("integration")
@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.STRICT_STUBS)
public class StreamsAssignmentScaleTest {
    static final long MAX_ASSIGNMENT_DURATION = 120 * 1000L; // we should stay below `max.poll.interval.ms`
    static final String APPLICATION_ID = "streams-assignment-scale-test";

    private static final Logger log = LoggerFactory.getLogger(StreamsAssignmentScaleTest.class);

    /* HighAvailabilityTaskAssignor tests */

    @Test
    @Timeout(value = 300)
    public void testHighAvailabilityTaskAssignorLargePartitionCount() {
        completeLargeAssignment(6_000, 2, 1, 1, HighAvailabilityTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testHighAvailabilityTaskAssignorLargeNumConsumers() {
        completeLargeAssignment(1_000, 1_000, 1, 1, HighAvailabilityTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testHighAvailabilityTaskAssignorManyStandbys() {
        completeLargeAssignment(1_000, 100, 1, 50, HighAvailabilityTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testHighAvailabilityTaskAssignorManyThreadsPerClient() {
        completeLargeAssignment(1_000, 10, 1000, 1, HighAvailabilityTaskAssignor.class);
    }

    /* LegacyStickyTaskAssignor tests */

    @Test
    @Timeout(value = 300)
    public void testStickyTaskAssignorLargePartitionCount() {
        completeLargeAssignment(2_000, 2, 1, 1, LegacyStickyTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testStickyTaskAssignorLargeNumConsumers() {
        completeLargeAssignment(1_000, 1_000, 1, 1, LegacyStickyTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testStickyTaskAssignorManyStandbys() {
        completeLargeAssignment(1_000, 100, 1, 20, LegacyStickyTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testStickyTaskAssignorManyThreadsPerClient() {
        completeLargeAssignment(1_000, 10, 1000, 1, LegacyStickyTaskAssignor.class);
    }

    /* FallbackPriorTaskAssignor tests */

    @Test
    @Timeout(value = 300)
    public void testFallbackPriorTaskAssignorLargePartitionCount() {
        completeLargeAssignment(2_000, 2, 1, 1, FallbackPriorTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testFallbackPriorTaskAssignorLargeNumConsumers() {
        completeLargeAssignment(1_000, 1_000, 1, 1, FallbackPriorTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testFallbackPriorTaskAssignorManyStandbys() {
        completeLargeAssignment(1_000, 100, 1, 20, FallbackPriorTaskAssignor.class);
    }

    @Test
    @Timeout(value = 300)
    public void testFallbackPriorTaskAssignorManyThreadsPerClient() {
        completeLargeAssignment(1_000, 10, 1000, 1, FallbackPriorTaskAssignor.class);
    }

    private void completeLargeAssignment(final int numPartitions,
                                         final int numClients,
                                         final int numThreadsPerClient,
                                         final int numStandbys,
                                         final Class<? extends LegacyTaskAssignor> taskAssignor) {
        final List<String> topic = singletonList("topic");

        final Map<TopicPartition, Long> changelogEndOffsets = new HashMap<>();
        for (int p = 0; p < numPartitions; ++p) {
            changelogEndOffsets.put(new TopicPartition(APPLICATION_ID + "-store-changelog", p), 100_000L);
        }

        final List<PartitionInfo> partitionInfos = new ArrayList<>();
        for (int p = 0; p < numPartitions; ++p) {
            partitionInfos.add(new PartitionInfo("topic", p, Node.noNode(), new Node[0], new Node[0]));
        }

        final Cluster clusterMetadata = new Cluster(
            "cluster",
            Collections.singletonList(Node.noNode()),
            partitionInfos,
            emptySet(),
            emptySet()
        );
        final Map<String, Object> configMap = new HashMap<>();
        configMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
        configMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:8080");
        final InternalTopologyBuilder builder = new InternalTopologyBuilder();
        builder.addSource(null, "source", null, null, null, "topic");
        builder.addProcessor("processor", new MockApiProcessorSupplier<>(), "source");
        builder.addStateStore(new MockKeyValueStoreBuilder("store", false), "processor");
        final TopologyMetadata topologyMetadata = new TopologyMetadata(builder, new StreamsConfig(configMap));
        topologyMetadata.buildAndRewriteTopology();

        @SuppressWarnings("unchecked")
        final Consumer<byte[], byte[]> mainConsumer = mock(Consumer.class);
        final TaskManager taskManager = mock(TaskManager.class);
        when(taskManager.topologyMetadata()).thenReturn(topologyMetadata);

        final AdminClient adminClient = createMockAdminClientForAssignor(changelogEndOffsets, true);

        final ReferenceContainer referenceContainer = new ReferenceContainer();
        referenceContainer.mainConsumer = mainConsumer;
        referenceContainer.adminClient = adminClient;
        referenceContainer.taskManager = taskManager;
        referenceContainer.streamsMetadataState = mock(StreamsMetadataState.class);
        referenceContainer.time = new MockTime();
        configMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
        configMap.put(InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS, taskAssignor.getName());
        configMap.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, numStandbys);

        configMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_CONFIG,
            StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE);

        final MockInternalTopicManager mockInternalTopicManager = spy(new MockInternalTopicManager(
            new MockTime(),
            new StreamsConfig(configMap),
            new MockClientSupplier().restoreConsumer,
            false
        ));

        lenient().when(mockInternalTopicManager.getTopicPartitionInfo(anySet())).thenAnswer(
            invocation -> {
                final Map<String, List<TopicPartitionInfo>> answer = new HashMap<>();
                final Set<String> topics = invocation.getArgument(0);
                for (final String thisTopic : topics) {
                    final List<TopicPartitionInfo> topicPartitionInfos = new ArrayList<>();
                    partitionInfos.forEach(info -> {
                        final TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(
                            info.partition(),
                            info.leader(),
                            singletonList(new Node(1, "host1", 80, "rack-1")),
                            asList(info.inSyncReplicas())
                        );
                        topicPartitionInfos.add(topicPartitionInfo);
                    });
                    answer.put(thisTopic, topicPartitionInfos);
                }
                return answer;
            }
        );

        final StreamsPartitionAssignor partitionAssignor = new StreamsPartitionAssignor();
        partitionAssignor.configure(configMap);
        partitionAssignor.setInternalTopicManager(mockInternalTopicManager);

        final Map<String, Subscription> subscriptions = new HashMap<>();
        for (int client = 0; client < numClients; ++client) {
            for (int i = 0; i < numThreadsPerClient; ++i) {
                subscriptions.put(
                    getConsumerName(i, client),
                    new Subscription(
                        topic,
                        getInfo(processIdForInt(client), EMPTY_TASKS, EMPTY_TASKS).encode(),
                        Collections.emptyList(),
                        DEFAULT_GENERATION,
                        Optional.of(String.format("rack-%d", client % 31))
                    )
                );
            }
        }

        final long firstAssignmentStartMs = System.currentTimeMillis();
        final Map<String, Assignment> firstAssignments = partitionAssignor.assign(clusterMetadata, new GroupSubscription(subscriptions)).groupAssignment();
        final long firstAssignmentEndMs = System.currentTimeMillis();

        final long firstAssignmentDuration = firstAssignmentEndMs - firstAssignmentStartMs;
        if (firstAssignmentDuration > MAX_ASSIGNMENT_DURATION) {
            throw new AssertionError("The first assignment took too long to complete at " + firstAssignmentDuration + "ms.");
        } else {
            log.info("First assignment took {}ms.", firstAssignmentDuration);
        }

        // Use the assignment to generate the subscriptions' prev task data for the next rebalance
        for (int client = 0; client < numClients; ++client) {
            for (int i = 0; i < numThreadsPerClient; ++i) {
                final String consumer = getConsumerName(i, client);
                final Assignment assignment = firstAssignments.get(consumer);
                final AssignmentInfo info = AssignmentInfo.decode(assignment.userData());

                subscriptions.put(
                    consumer,
                    new Subscription(
                        topic,
                        getInfo(processIdForInt(client), new HashSet<>(info.activeTasks()), info.standbyTasks().keySet()).encode(),
                        assignment.partitions())
                );
            }
        }

        final long secondAssignmentStartMs = System.currentTimeMillis();
        final Map<String, Assignment> secondAssignments = partitionAssignor.assign(clusterMetadata, new GroupSubscription(subscriptions)).groupAssignment();
        final long secondAssignmentEndMs = System.currentTimeMillis();
        final long secondAssignmentDuration = secondAssignmentEndMs - secondAssignmentStartMs;
        if (secondAssignmentDuration > MAX_ASSIGNMENT_DURATION) {
            throw new AssertionError("The second assignment took too long to complete at " + secondAssignmentDuration + "ms.");
        } else {
            log.info("Second assignment took {}ms.", secondAssignmentDuration);
        }

        assertThat(secondAssignments.size(), is(numClients * numThreadsPerClient));
    }

    private String getConsumerName(final int consumerIndex, final int clientIndex) {
        return "consumer-" + clientIndex + "-" + consumerIndex;
    }
}
