/*
 * 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.jmh.assignor;

import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.metadata.PartitionRecord;
import org.apache.kafka.common.metadata.TopicRecord;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.group.Utils;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.TopicIds;
import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
import org.apache.kafka.coordinator.group.modern.share.ShareGroupMember;
import org.apache.kafka.image.MetadataDelta;
import org.apache.kafka.image.MetadataImage;
import org.apache.kafka.image.MetadataProvenance;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;


public class AssignorBenchmarkUtils {
    /**
     * Generate a reverse look up map of partition to member target assignments from the given member spec.
     *
     * @param groupAssignment       The group assignment.
     * @return Map of topic partition to member assignments.
     */
    public static Map<Uuid, Map<Integer, String>> computeInvertedTargetAssignment(
        GroupAssignment groupAssignment
    ) {
        Map<Uuid, Map<Integer, String>> invertedTargetAssignment = new HashMap<>();
        for (Map.Entry<String, MemberAssignment> memberEntry : groupAssignment.members().entrySet()) {
            String memberId = memberEntry.getKey();
            Map<Uuid, Set<Integer>> topicsAndPartitions = memberEntry.getValue().partitions();

            for (Map.Entry<Uuid, Set<Integer>> topicEntry : topicsAndPartitions.entrySet()) {
                Uuid topicId = topicEntry.getKey();
                Set<Integer> partitions = topicEntry.getValue();

                Map<Integer, String> partitionMap = invertedTargetAssignment.computeIfAbsent(topicId, k -> new HashMap<>());

                for (Integer partitionId : partitions) {
                    partitionMap.put(partitionId, memberId);
                }
            }
        }
        return invertedTargetAssignment;
    }

    /**
     * Generates a list of topic names for use in benchmarks.
     *
     * @param topicCount            The number of topic names to generate.
     * @return The list of topic names.
     */
    public static List<String> createTopicNames(int topicCount) {
        List<String> topicNames = new ArrayList<>();
        for (int i = 0; i < topicCount; i++) {
            topicNames.add("topic-" + i);
        }
        return topicNames;
    }

    /**
     * Creates a TopicsImage from the given topic names.
     *
     * @param allTopicNames      The topic names.
     * @param partitionsPerTopic Number of partitions per topic.
     * @return A TopicsImage containing the topic ids, names and partition counts from the
     *         subscription metadata.
     */
    public static CoordinatorMetadataImage createMetadataImage(
        List<String> allTopicNames,
        int partitionsPerTopic
    ) {
        MetadataDelta delta = new MetadataDelta.Builder()
            .setImage(MetadataImage.EMPTY)
            .build();

        for (String topicName : allTopicNames) {
            AssignorBenchmarkUtils.addTopic(
                delta,
                Uuid.randomUuid(),
                topicName,
                partitionsPerTopic
            );
        }

        return new KRaftCoordinatorMetadataImage(delta.apply(MetadataProvenance.EMPTY));
    }

    /**
     * Creates a GroupSpec from the given ConsumerGroupMembers.
     *
     * @param members               The ConsumerGroupMembers.
     * @param subscriptionType      The group's subscription type.
     * @param topicResolver         The TopicResolver to use.
     * @return The new GroupSpec.
     */
    public static GroupSpec createConsumerGroupSpec(
        Map<String, ConsumerGroupMember> members,
        SubscriptionType subscriptionType,
        TopicIds.TopicResolver topicResolver
    ) {
        Map<String, MemberSubscriptionAndAssignmentImpl> memberSpecs = new HashMap<>();

        for (Map.Entry<String, ConsumerGroupMember> memberEntry : members.entrySet()) {
            String memberId = memberEntry.getKey();
            ConsumerGroupMember member = memberEntry.getValue();

            memberSpecs.put(memberId, new MemberSubscriptionAndAssignmentImpl(
                Optional.ofNullable(member.rackId()),
                Optional.ofNullable(member.instanceId()),
                new TopicIds(member.subscribedTopicNames(), topicResolver),
                new Assignment(Utils.toAssignmentWithoutEpochs(member.assignedPartitions()))
            ));
        }

        return new GroupSpecImpl(
            memberSpecs,
            subscriptionType,
            Map.of()
        );
    }

    /**
     * Creates a GroupSpec from the given ShareGroupMembers.
     *
     * @param members               The ShareGroupMembers.
     * @param subscriptionType      The group's subscription type.
     * @param topicResolver         The TopicResolver to use.
     * @return The new GroupSpec.
     */
    public static GroupSpec createShareGroupSpec(
        Map<String, ShareGroupMember> members,
        SubscriptionType subscriptionType,
        TopicIds.TopicResolver topicResolver
    ) {
        Map<String, MemberSubscriptionAndAssignmentImpl> memberSpecs = new HashMap<>();

        for (Map.Entry<String, ShareGroupMember> memberEntry : members.entrySet()) {
            String memberId = memberEntry.getKey();
            ShareGroupMember member = memberEntry.getValue();

            memberSpecs.put(memberId, new MemberSubscriptionAndAssignmentImpl(
                Optional.ofNullable(member.rackId()),
                Optional.empty(),
                new TopicIds(member.subscribedTopicNames(), topicResolver),
                new Assignment(member.assignedPartitions())
            ));
        }

        return new GroupSpecImpl(
            memberSpecs,
            subscriptionType,
            Map.of()
        );
    }

    /**
     * Creates a ConsumerGroupMembers map where all members have the same topic subscriptions.
     *
     * @param memberCount           The number of members in the group.
     * @param getMemberId           A function to map member indices to member ids.
     * @param getMemberRackId       A function to map member indices to rack ids.
     * @param topicNames            The topics to subscribe to.
     * @return The new ConsumerGroupMembers map.
     */
    public static Map<String, ConsumerGroupMember> createHomogeneousMembers(
        int memberCount,
        Function<Integer, String> getMemberId,
        Function<Integer, Optional<String>> getMemberRackId,
        List<String> topicNames
    ) {
        Map<String, ConsumerGroupMember> members = new HashMap<>();

        for (int i = 0; i < memberCount; i++) {
            String memberId = getMemberId.apply(i);
            Optional<String> rackId = getMemberRackId.apply(i);

            members.put(memberId, new ConsumerGroupMember.Builder("member" + i)
                .setRackId(rackId.orElse(null))
                .setSubscribedTopicNames(topicNames)
                .build()
            );
        }

        return members;
    }

    /**
     * Creates a ConsumerGroupMembers map where members have different topic subscriptions.
     *
     * Divides members and topics into a given number of buckets. Within each bucket, members are
     * subscribed to the same topics.
     *
     * @param memberCount           The number of members in the group.
     * @param bucketCount           The number of buckets.
     * @param getMemberId           A function to map member indices to member ids.
     * @param getMemberRackId       A function to map member indices to rack ids.
     * @param topicNames            The topics to subscribe to.
     * @return The new ConsumerGroupMembers map.
     */
    public static Map<String, ConsumerGroupMember> createHeterogeneousBucketedMembers(
        int memberCount,
        int bucketCount,
        Function<Integer, String> getMemberId,
        Function<Integer, Optional<String>> getMemberRackId,
        List<String> topicNames
    ) {
        Map<String, ConsumerGroupMember> members = new HashMap<>();

        // Adjust bucket count based on member count when member count < max bucket count.
        bucketCount = Math.min(bucketCount, memberCount);

        // Check minimum topics requirement
        if (topicNames.size() < bucketCount) {
            throw new IllegalArgumentException("At least " + bucketCount + " topics are recommended for effective bucketing.");
        }

        int bucketSizeTopics = (int) Math.ceil((double) topicNames.size() / bucketCount);
        int bucketSizeMembers = (int) Math.ceil((double) memberCount / bucketCount);

        // Define buckets for each member and assign topics from the same bucket
        for (int bucket = 0; bucket < bucketCount; bucket++) {
            int memberStartIndex = bucket * bucketSizeMembers;
            int memberEndIndex = Math.min((bucket + 1) * bucketSizeMembers, memberCount);

            int topicStartIndex = bucket * bucketSizeTopics;
            int topicEndIndex = Math.min((bucket + 1) * bucketSizeTopics, topicNames.size());

            List<String> bucketTopicNames = topicNames.subList(topicStartIndex, topicEndIndex);

            // Assign topics to each member in the current bucket
            for (int i = memberStartIndex; i < memberEndIndex; i++) {
                String memberId = getMemberId.apply(i);
                Optional<String> rackId = getMemberRackId.apply(i);

                members.put(memberId, new ConsumerGroupMember.Builder("member" + i)
                    .setRackId(rackId.orElse(null))
                    .setSubscribedTopicNames(bucketTopicNames)
                    .build()
                );
            }
        }

        return members;
    }

    /**
     * Creates a ShareGroupMembers map where all members have the same topic subscriptions.
     *
     * @param memberCount           The number of members in the group.
     * @param getMemberId           A function to map member indices to member ids.
     * @param topicNames            The topics to subscribe to.
     * @return The new ShareGroupMembers map.
     */
    public static Map<String, ShareGroupMember> createHomogeneousShareGroupMembers(
        int memberCount,
        Function<Integer, String> getMemberId,
        List<String> topicNames
    ) {
        Map<String, ShareGroupMember> members = new HashMap<>();

        for (int i = 0; i < memberCount; i++) {
            String memberId = getMemberId.apply(i);

            members.put(memberId, new ShareGroupMember.Builder("member" + i)
                .setSubscribedTopicNames(topicNames)
                .build()
            );
        }

        return members;
    }

    /**
     * Creates a ShareGroupMembers map where members have different topic subscriptions.
     *
     * Divides members and topics into a given number of buckets. Within each bucket, members are
     * subscribed to the same topics.
     *
     * @param memberCount           The number of members in the group.
     * @param bucketCount           The number of buckets.
     * @param getMemberId           A function to map member indices to member ids.
     * @param topicNames            The topics to subscribe to.
     * @return The new ShareGroupMembers map.
     */
    public static Map<String, ShareGroupMember> createHeterogeneousBucketedShareGroupMembers(
        int memberCount,
        int bucketCount,
        Function<Integer, String> getMemberId,
        List<String> topicNames
    ) {
        Map<String, ShareGroupMember> members = new HashMap<>();

        // Adjust bucket count based on member count when member count < max bucket count.
        bucketCount = Math.min(bucketCount, memberCount);

        // Check minimum topics requirement
        if (topicNames.size() < bucketCount) {
            throw new IllegalArgumentException("At least " + bucketCount + " topics are recommended for effective bucketing.");
        }

        int bucketSizeTopics = (int) Math.ceil((double) topicNames.size() / bucketCount);
        int bucketSizeMembers = (int) Math.ceil((double) memberCount / bucketCount);

        // Define buckets for each member and assign topics from the same bucket
        for (int bucket = 0; bucket < bucketCount; bucket++) {
            int memberStartIndex = bucket * bucketSizeMembers;
            int memberEndIndex = Math.min((bucket + 1) * bucketSizeMembers, memberCount);

            int topicStartIndex = bucket * bucketSizeTopics;
            int topicEndIndex = Math.min((bucket + 1) * bucketSizeTopics, topicNames.size());

            List<String> bucketTopicNames = topicNames.subList(topicStartIndex, topicEndIndex);

            // Assign topics to each member in the current bucket
            for (int i = memberStartIndex; i < memberEndIndex; i++) {
                String memberId = getMemberId.apply(i);

                members.put(memberId, new ShareGroupMember.Builder("member" + i)
                    .setSubscribedTopicNames(bucketTopicNames)
                    .build()
                );
            }
        }

        return members;
    }

    public static void addTopic(
        MetadataDelta delta,
        Uuid topicId,
        String topicName,
        int numPartitions
    ) {
        // For testing purposes, the following criteria are used:
        // - Number of replicas for each partition: 2
        // - Number of brokers available in the cluster: 4
        delta.replay(new TopicRecord().setTopicId(topicId).setName(topicName));
        for (int i = 0; i < numPartitions; i++) {
            delta.replay(new PartitionRecord()
                .setTopicId(topicId)
                .setPartitionId(i)
                .setReplicas(List.of(i % 4, (i + 1) % 4)));
        }
    }
}
