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

import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.group.api.assignor.ConsumerGroupPartitionAssignor;
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.MemberSubscription;
import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException;
import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl;

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

/**
 * A range assignor assigns contiguous partition ranges to members of a consumer group such that:
 * <ol>
 *   <li>Each subscribed member receives at least one partition from that topic.</li>
 *   <li>Each member receives the same partition number from every subscribed topic when co-partitioning is possible.</li>
 * </ol>
 *
 * Co-partitioning is possible when the below conditions are satisfied:
 * <ol>
 *   <li>All the members are subscribed to the same set of topics.</li>
 *   <li>All the topics have the same number of partitions.</li>
 * </ol>
 *
 * Co-partitioning is useful in performing joins on data streams.
 *
 * <p>For example, suppose there are two members M0 and M1, two topics T1 and T2, and each topic has 3 partitions.
 *
 * <p>The co-partitioned assignment will be:
 * <ul>
 * <li<code>    M0: [T1P0, T1P1, T2P0, T2P1]    </code></li>
 * <li><code>   M1: [T1P2, T2P2]                </code></li>
 * </ul>
 *
 * Since the introduction of static membership, we could leverage <code>member.instance.id</code> to make the
 * assignment behavior more sticky.
 * For the above example, after one rolling bounce, the group coordinator will attempt to assign new member Ids towards
 * members, for example if <code>M0</code> -&gt; <code>M3</code> <code>M1</code> -&gt; <code>M2</code>.
 *
 * <p>The assignment could be completely shuffled to:
 * <ul>
 * <li><code>   M3 (was M0): [T1P2, T2P2]               (before it was [T1P0, T1P1, T2P0, T2P1])  </code>
 * <li><code>   M2 (was M1): [T1P0, T1P1, T2P0, T2P1]   (before it was [T1P2, T2P2])  </code>
 * </ul>
 *
 * The assignment change was caused by the change of <code>member.id</code> relative order, and
 * can be avoided by setting the instance.id.
 * Members will have individual instance Ids <code>I0</code>, <code>I1</code>. As long as
 * 1. Number of members remain the same.
 * 2. Topic metadata doesn't change.
 * 3. Subscription pattern doesn't change for any member.
 *
 * <p>The assignment will always be:
 * <ul>
 * <li><code>   I0: [T1P0, T1P1, T2P0, T2P1]    </code>
 * <li><code>   I1: [T1P2, T2P2]                </code>
 * </ul>
 * <p>
 */
public class RangeAssignor implements ConsumerGroupPartitionAssignor {
    public static final String NAME = "range";

    @Override
    public String name() {
        return NAME;
    }

    /**
     * Metadata for a topic including partition and subscription details.
     */
    private static class TopicMetadata {
        private final Uuid topicId;
        private final int numPartitions;
        private int numMembers;
        private int minQuota = -1;
        private int extraPartitions = -1;
        private int nextRange = 0;

        /**
         * Constructs a new TopicMetadata instance.
         *
         * @param topicId           The topic Id.
         * @param numPartitions     The number of partitions.
         * @param numMembers        The number of subscribed members.
         */
        private TopicMetadata(Uuid topicId, int numPartitions, int numMembers) {
            this.topicId = topicId;
            this.numPartitions = numPartitions;
            this.numMembers = numMembers;
        }

        /**
         * Computes the minimum partition quota per member and the extra partitions, if not already computed.
         */
        private void maybeComputeQuota() {
            if (minQuota != -1) return;

            // The minimum number of partitions each member should receive for a balanced assignment.
            minQuota = numPartitions / numMembers;

            // Extra partitions to be distributed one to each member.
            extraPartitions = numPartitions % numMembers;
        }

        @Override
        public String toString() {
            return "TopicMetadata(topicId=" + topicId +
                ", numPartitions=" + numPartitions +
                ", numMembers=" + numMembers +
                ", minQuota=" + minQuota +
                ", extraPartitions=" + extraPartitions +
                ", nextRange=" + nextRange +
                ')';
        }
    }

    /**
     * Assigns partitions to members of a homogeneous group. All members are subscribed to the same set of topics.
     * Assignment will be co-partitioned when all the topics have an equal number of partitions.
     */
    private GroupAssignment assignHomogeneousGroup(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) throws PartitionAssignorException {
        List<String> memberIds = sortMemberIds(groupSpec);
        int numMembers = groupSpec.memberIds().size();

        MemberSubscription subs = groupSpec.memberSubscription(memberIds.get(0));
        List<TopicMetadata> topics = new ArrayList<>(subs.subscribedTopicIds().size());

        for (Uuid topicId : subs.subscribedTopicIds()) {
            int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
            if (numPartitions == -1) {
                throw new PartitionAssignorException("Member is subscribed to a non-existent topic");
            }
            TopicMetadata m = new TopicMetadata(
                topicId,
                numPartitions,
                numMembers
            );
            topics.add(m);
        }

        Map<String, MemberAssignment> assignments = new HashMap<>((int) ((groupSpec.memberIds().size() / 0.75f) + 1));
        int memberAssignmentInitialCapacity = (int) ((topics.size() / 0.75f) + 1);

        for (String memberId : memberIds) {
            Map<Uuid, Set<Integer>> assignment = new HashMap<>(memberAssignmentInitialCapacity);
            for (TopicMetadata topicMetadata : topics) {
                topicMetadata.maybeComputeQuota();
                addPartitionsToAssignment(topicMetadata, assignment);
            }
            assignments.put(memberId, new MemberAssignmentImpl(assignment));
        }

        return new GroupAssignment(assignments);
    }

    /**
     * Assigns partitions to members of a heterogeneous group. Not all members are subscribed to the same topics.
     */
    private GroupAssignment assignHeterogeneousGroup(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) throws PartitionAssignorException {
        List<String> memberIds = sortMemberIds(groupSpec);

        Map<Uuid, TopicMetadata> topics = new HashMap<>();

        for (String memberId : memberIds) {
            MemberSubscription subs = groupSpec.memberSubscription(memberId);
            for (Uuid topicId : subs.subscribedTopicIds()) {
                TopicMetadata topicMetadata = topics.computeIfAbsent(topicId, __ -> {
                    int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
                    if (numPartitions == -1) {
                        throw new PartitionAssignorException("Member is subscribed to a non-existent topic");
                    }

                    return new TopicMetadata(
                        topicId,
                        numPartitions,
                        0
                    );
                });
                topicMetadata.numMembers++;
            }
        }

        Map<String, MemberAssignment> assignments = new HashMap<>((int) ((groupSpec.memberIds().size() / 0.75f) + 1));

        for (String memberId : memberIds) {
            MemberSubscription subs = groupSpec.memberSubscription(memberId);
            Map<Uuid, Set<Integer>> assignment = new HashMap<>((int) ((subs.subscribedTopicIds().size() / 0.75f) + 1));
            for (Uuid topicId : subs.subscribedTopicIds()) {
                TopicMetadata metadata = topics.get(topicId);
                metadata.maybeComputeQuota();
                addPartitionsToAssignment(metadata, assignment);
            }
            assignments.put(memberId, new MemberAssignmentImpl(assignment));
        }

        return new GroupAssignment(assignments);
    }

    /**
     * Sorts members based on their instance Ids if available or by member Ids if not.
     *
     * Static members are placed first and non-static members follow.
     *
     * Prioritizing static members helps them retain their partitions, enhancing stickiness
     * and stability. Non-static members, which do not have guaranteed rejoining Ids, are placed
     * later, allowing for more dynamic and flexible partition assignments.
     *
     * @param groupSpec     The group specification containing the member information.
     * @return A sorted list of member Ids.
     */
    private List<String> sortMemberIds(
        GroupSpec groupSpec
    ) {
        List<String> sortedMemberIds = new ArrayList<>(groupSpec.memberIds());

        sortedMemberIds.sort((memberId1, memberId2) -> {
            Optional<String> instanceId1 = groupSpec.memberSubscription(memberId1).instanceId();
            Optional<String> instanceId2 = groupSpec.memberSubscription(memberId2).instanceId();

            if (instanceId1.isPresent() && instanceId2.isPresent()) {
                return instanceId1.get().compareTo(instanceId2.get());
            } else if (instanceId1.isPresent()) {
                return -1;
            } else if (instanceId2.isPresent()) {
                return 1;
            } else {
                return memberId1.compareTo(memberId2);
            }
        });
        return sortedMemberIds;
    }

    /**
     * Assigns a range of partitions to the specified topic based on the provided metadata.
     *
     * @param topicMetadata         Metadata containing the topic details, including the number of partitions,
     *                              the next range to assign, minQuota, and extra partitions.
     * @param memberAssignment      Map from topic Id to the set of assigned partition Ids.
     */
    private void addPartitionsToAssignment(
        TopicMetadata topicMetadata,
        Map<Uuid, Set<Integer>> memberAssignment
    ) {
        int start = topicMetadata.nextRange;
        int quota = topicMetadata.minQuota;

        // Adjust quota to account for extra partitions if available.
        if (topicMetadata.extraPartitions > 0) {
            quota++;
            topicMetadata.extraPartitions--;
        }

        // Calculate the end using the quota.
        int end = Math.min(start + quota, topicMetadata.numPartitions);

        topicMetadata.nextRange = end;

        if (start < end) {
            memberAssignment.put(topicMetadata.topicId, new RangeSet(start, end));
        }
    }

    /**
     * Assigns partitions to members based on their topic subscriptions and the properties of a range assignor:
     *
     * @param groupSpec                     The group specification containing the member information.
     * @param subscribedTopicDescriber      The describer for subscribed topics to get the number of partitions.
     * @return The group's assignment with the partition assignments for each member.
     * @throws PartitionAssignorException if any member is subscribed to a non-existent topic.
     */
    @Override
    public GroupAssignment assign(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) throws PartitionAssignorException {
        if (groupSpec.memberIds().isEmpty()) {
            return new GroupAssignment(Map.of());
        } else if (groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS) {
            return assignHomogeneousGroup(groupSpec, subscribedTopicDescriber);
        } else {
            return assignHeterogeneousGroup(groupSpec, subscribedTopicDescriber);
        }
    }
}
