/*
 * 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.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.ShareGroupPartitionAssignor;
import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl;
import org.apache.kafka.server.common.TopicIdPartition;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Set;

import static org.apache.kafka.coordinator.group.api.assignor.SubscriptionType.HOMOGENEOUS;

/**
 * A simple partition assignor for share groups that assigns partitions of the subscribed topics
 * based on the rules defined in KIP-932 to different members. It is not rack-aware.
 * <p>
 * Assignments are done according to the following principles:
 * <ol>
 *   <li>Balance:          Ensure partitions are distributed equally among all members.
 *                         The difference in assignments sizes between any two members
 *                         should not exceed one partition.</li>
 *   <li>Stickiness:       Minimize partition movements among members by retaining
 *                         as much of the existing assignment as possible.</li>
 * </ol>
 * <p>
 * Balance is prioritized above stickiness.
 */
public class SimpleAssignor implements ShareGroupPartitionAssignor {

    private static final String SIMPLE_ASSIGNOR_NAME = "simple";

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

    @Override
    public GroupAssignment assign(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) throws PartitionAssignorException {
        if (groupSpec.memberIds().isEmpty())
            return new GroupAssignment(Map.of());

        if (groupSpec.subscriptionType().equals(HOMOGENEOUS)) {
            return assignHomogeneous(groupSpec, subscribedTopicDescriber);
        } else {
            return assignHeterogeneous(groupSpec, subscribedTopicDescriber);
        }
    }

    private GroupAssignment assignHomogeneous(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        Set<Uuid> subscribedTopicIds = groupSpec.memberSubscription(groupSpec.memberIds().iterator().next())
            .subscribedTopicIds();
        if (subscribedTopicIds.isEmpty())
            return new GroupAssignment(Map.of());

        // Subscribed topic partitions for the share group.
        List<TopicIdPartition> targetPartitions = computeTargetPartitions(
            subscribedTopicIds, subscribedTopicDescriber);

        // The current assignment from topic partition to members.
        Map<TopicIdPartition, List<String>> currentAssignment = currentAssignment(groupSpec);
        return newAssignmentHomogeneous(groupSpec, subscribedTopicIds, targetPartitions, currentAssignment, subscribedTopicDescriber);
    }

    private GroupAssignment assignHeterogeneous(
        GroupSpec groupSpec,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        Map<String, List<TopicIdPartition>> memberToPartitionsSubscription = new HashMap<>();
        for (String memberId : groupSpec.memberIds()) {
            MemberSubscription spec = groupSpec.memberSubscription(memberId);
            if (spec.subscribedTopicIds().isEmpty())
                continue;

            // Subscribed topic partitions for the share group member.
            List<TopicIdPartition> targetPartitions = computeTargetPartitions(
                spec.subscribedTopicIds(), subscribedTopicDescriber);
            memberToPartitionsSubscription.put(memberId, targetPartitions);
        }

        // The current assignment from topic partition to members.
        Map<TopicIdPartition, List<String>> currentAssignment = currentAssignment(groupSpec);
        return newAssignmentHeterogeneous(groupSpec, memberToPartitionsSubscription, currentAssignment, subscribedTopicDescriber);
    }

    /**
     * Get the current assignment by topic partitions.
     * @param groupSpec The group metadata specifications.
     * @return the current assignment for subscribed topic partitions to memberIds.
     */
    private Map<TopicIdPartition, List<String>> currentAssignment(GroupSpec groupSpec) {
        Map<TopicIdPartition, List<String>> assignment = new HashMap<>();

        for (String member : groupSpec.memberIds()) {
            Map<Uuid, Set<Integer>> assignedTopicPartitions = groupSpec.memberAssignment(member).partitions();
            assignedTopicPartitions.forEach((topicId, partitions) -> partitions.forEach(
                partition -> assignment.computeIfAbsent(new TopicIdPartition(topicId, partition), k -> new ArrayList<>()).add(member)));
        }
        return assignment;
    }

    /**
     * This function computes the new assignment for a homogeneous group.
     * @param groupSpec           The group metadata specifications.
     * @param subscribedTopicIds  The set of all the subscribed topic ids for the group.
     * @param targetPartitions    The list of all topic partitions that need assignment.
     * @param currentAssignment   The current assignment for subscribed topic partitions to memberIds.
     * @return the new partition assignment for the members of the group.
     */
    private GroupAssignment newAssignmentHomogeneous(
        GroupSpec groupSpec,
        Set<Uuid> subscribedTopicIds,
        List<TopicIdPartition> targetPartitions,
        Map<TopicIdPartition, List<String>> currentAssignment,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        // For entirely balanced assignment, we would expect (numTargetPartitions / numGroupMembers) partitions per member, rounded upwards.
        // That can be expressed as         Math.ceil(numTargetPartitions / (double) numGroupMembers)
        // Using integer arithmetic, as     (numTargetPartitions + numGroupMembers - 1) / numGroupMembers
        int numGroupMembers = groupSpec.memberIds().size();
        int numTargetPartitions = targetPartitions.size();
        int desiredAssignmentCount = (numTargetPartitions + numGroupMembers - 1) / numGroupMembers;

        Map<TopicIdPartition, List<String>> newAssignment = newHashMap(numTargetPartitions);

        // Hash member IDs to topic partitions. Each member will be assigned one partition, but some partitions
        // might have been assigned to more than one member.
        memberHashAssignment(groupSpec.memberIds(), targetPartitions, newAssignment);

        // Combine current and new hashed assignments, sized to accommodate the expected number of mappings.
        Map<String, Set<TopicIdPartition>> finalAssignment = newHashMap(numGroupMembers);
        Map<TopicIdPartition, Set<String>> finalAssignmentByPartition = newHashMap(numTargetPartitions);

        // First, take the members assigned by hashing.
        newAssignment.forEach((targetPartition, members) -> members.forEach(member -> {
            finalAssignment.computeIfAbsent(member, k -> new HashSet<>()).add(targetPartition);
            finalAssignmentByPartition.computeIfAbsent(targetPartition, k -> new HashSet<>()).add(member);
        }));

        // Then, take the members from the current assignment, making sure that no member has too many assigned partitions.
        // When combining current assignment, we need to only consider the topics in current assignment that are also being
        // subscribed in the new assignment as well.
        currentAssignment.forEach((targetPartition, members) -> {
            if (subscribedTopicIds.contains(targetPartition.topicId())) {
                members.forEach(member -> {
                    if (groupSpec.memberIds().contains(member)) {
                        Set<TopicIdPartition> memberPartitions = finalAssignment.computeIfAbsent(member, k -> new HashSet<>());
                        if ((memberPartitions.size() < desiredAssignmentCount) && !newAssignment.containsKey(targetPartition)) {
                            memberPartitions.add(targetPartition);
                            finalAssignmentByPartition.computeIfAbsent(targetPartition, k -> new HashSet<>()).add(member);
                        }
                    }
                });
            }
        });

        // Finally, round-robin assignment for unassigned partitions which do not already have members assigned.
        // The order of steps differs slightly from KIP-932 because the desired assignment count has been taken into
        // account when copying partitions across from the current assignment, and this is more convenient.
        List<TopicIdPartition> unassignedPartitions = targetPartitions.stream()
            .filter(targetPartition -> !finalAssignmentByPartition.containsKey(targetPartition))
            .toList();

        roundRobinAssignmentWithCount(groupSpec.memberIds(), unassignedPartitions, finalAssignment, desiredAssignmentCount, subscribedTopicDescriber);

        return groupAssignment(finalAssignment, groupSpec.memberIds());
    }

    /**
     * This function computes the new assignment for a heterogeneous group.
     * @param groupSpec                       The group metadata specifications.
     * @param memberToPartitionsSubscription  The member to subscribed topic partitions map.
     * @param currentAssignment               The current assignment for subscribed topic partitions to memberIds.
     * @return the new partition assignment for the members of the group.
     */
    private GroupAssignment newAssignmentHeterogeneous(
        GroupSpec groupSpec,
        Map<String, List<TopicIdPartition>> memberToPartitionsSubscription,
        Map<TopicIdPartition, List<String>> currentAssignment,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        int numGroupMembers = groupSpec.memberIds().size();

        // Exhaustive set of all subscribed topic partitions.
        Set<TopicIdPartition> targetPartitions = new LinkedHashSet<>();
        memberToPartitionsSubscription.values().forEach(targetPartitions::addAll);

        // Create a map for topic to members subscription.
        Map<Uuid, Set<String>> topicToMemberSubscription = new HashMap<>();
        memberToPartitionsSubscription.forEach((member, partitions) ->
            partitions.forEach(partition -> topicToMemberSubscription.computeIfAbsent(partition.topicId(), k -> new LinkedHashSet<>()).add(member)));

        Map<TopicIdPartition, List<String>> newAssignment = new HashMap<>();

        // Step 1: Hash member IDs to partitions.
        memberToPartitionsSubscription.forEach((member, partitions) ->
            memberHashAssignment(List.of(member), partitions, newAssignment));

        // Step 2: Round-robin assignment for unassigned partitions which do not have members already assigned in the current assignment.
        Set<TopicIdPartition> assignedPartitions = new LinkedHashSet<>(newAssignment.keySet());
        Map<Uuid, List<TopicIdPartition>> unassignedPartitions = new HashMap<>();
        targetPartitions.forEach(targetPartition -> {
            if (!assignedPartitions.contains(targetPartition) && !currentAssignment.containsKey(targetPartition))
                unassignedPartitions.computeIfAbsent(targetPartition.topicId(), k -> new ArrayList<>()).add(targetPartition);
        });

        unassignedPartitions.keySet().forEach(unassignedTopic ->
            roundRobinAssignment(topicToMemberSubscription.get(unassignedTopic), unassignedPartitions.get(unassignedTopic), newAssignment, subscribedTopicDescriber));

        // Step 3: We combine current assignment and new assignment.
        Map<String, Set<TopicIdPartition>> finalAssignment = newHashMap(numGroupMembers);

        newAssignment.forEach((targetPartition, members) -> members.forEach(member ->
            finalAssignment.computeIfAbsent(member, k -> new HashSet<>()).add(targetPartition)));

        // When combining current assignment, we need to only consider the member topic subscription in current assignment
        // which is being subscribed in the new assignment as well.
        currentAssignment.forEach((topicIdPartition, members) -> members.forEach(member -> {
            if (topicToMemberSubscription.getOrDefault(topicIdPartition.topicId(), Set.of()).contains(member)
                && !newAssignment.containsKey(topicIdPartition))
                finalAssignment.computeIfAbsent(member, k -> new HashSet<>()).add(topicIdPartition);
        }));

        return groupAssignment(finalAssignment, groupSpec.memberIds());
    }

    /**
     * This function updates assignment by hashing the member IDs of the members and maps the partitions assigned to the
     * members based on the hash, one partition per member. This gives approximately even balance.
     * @param memberIds           The member ids to which the topic partitions need to be assigned.
     * @param partitionsToAssign  The subscribed topic partitions which needs assignment.
     * @param assignment          The existing assignment by topic partition. We need to pass it as a parameter because this
     *                            method can be called multiple times for heterogeneous assignment.
     */
    // Visible for testing
    void memberHashAssignment(
        Collection<String> memberIds,
        List<TopicIdPartition> partitionsToAssign,
        Map<TopicIdPartition, List<String>> assignment
    ) {
        if (!partitionsToAssign.isEmpty()) {
            for (String memberId : memberIds) {
                int topicPartitionIndex = Math.abs(memberId.hashCode() % partitionsToAssign.size());
                TopicIdPartition topicPartition = partitionsToAssign.get(topicPartitionIndex);
                assignment.computeIfAbsent(topicPartition, k -> new ArrayList<>()).add(memberId);
            }
        }
    }

    /**
     * This functions assigns topic partitions to members by a round-robin approach and updates the existing assignment.
     * @param memberIds                 The member ids to which the topic partitions need to be assigned, should be non-empty.
     * @param partitionsToAssign        The subscribed topic partitions which needs assignment.
     * @param assignment                The existing assignment by topic partition. We need to pass it as a parameter because this
     *                                  method can be called multiple times for heterogeneous assignment.
     * @param subscribedTopicDescriber  The topic describer to fetch assignable partitions from.
     */
    // Visible for testing
    void roundRobinAssignment(
        Collection<String> memberIds,
        List<TopicIdPartition> partitionsToAssign,
        Map<TopicIdPartition, List<String>> assignment,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        // We iterate through the target partitions and assign a memberId to them. In case we run out of members (members < targetPartitions),
        // we again start from the starting index of memberIds.
        Iterator<String> memberIdIterator = memberIds.iterator();
        for (TopicIdPartition topicPartition : partitionsToAssign) {
            if (!memberIdIterator.hasNext()) {
                memberIdIterator = memberIds.iterator();
            }
            String memberId = memberIdIterator.next();
            if (subscribedTopicDescriber.assignablePartitions(topicPartition.topicId()).contains(topicPartition.partitionId())) {
                assignment.computeIfAbsent(topicPartition, k -> new ArrayList<>()).add(memberId);
            }
        }
    }

    /**
     * This functions assigns topic partitions to members by a round-robin approach and updates the existing assignment.
     * @param memberIds                 The member ids to which the topic partitions need to be assigned, should be non-empty.
     * @param partitionsToAssign        The subscribed topic partitions which need assignment.
     * @param assignment                The existing assignment by topic partition. We need to pass it as a parameter because this
     *                                  method can be called multiple times for heterogeneous assignment.
     * @param desiredAssignmentCount    The number of partitions which can be assigned to each member to give even balance.
     *                                  Note that this number can be exceeded by one to allow for situations
     *                                  in which we have hashing collisions.
     * @param subscribedTopicDescriber  The topic describer to fetch assignable partitions from.
     */
    void roundRobinAssignmentWithCount(
        Collection<String> memberIds,
        List<TopicIdPartition> partitionsToAssign,
        Map<String, Set<TopicIdPartition>> assignment,
        int desiredAssignmentCount,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        Collection<String> memberIdsCopy = new LinkedHashSet<>(memberIds);

        // We iterate through the target partitions which are not in the assignment and assign a memberId to them.
        // In case we run out of members (memberIds < partitionsToAssign), we again start from the starting index of memberIds.
        Iterator<String> memberIdIterator = memberIdsCopy.iterator();
        ListIterator<TopicIdPartition> partitionListIterator = partitionsToAssign.listIterator();
        while (partitionListIterator.hasNext()) {
            TopicIdPartition partition = partitionListIterator.next();
            if (!subscribedTopicDescriber.assignablePartitions(partition.topicId()).contains(partition.partitionId())) {
                continue;
            }
            if (!memberIdIterator.hasNext()) {
                memberIdIterator = memberIdsCopy.iterator();
                if (memberIdsCopy.isEmpty()) {
                    // This should never happen, but guarding against an infinite loop
                    throw new PartitionAssignorException("Inconsistent number of member IDs");
                }
            }
            String memberId = memberIdIterator.next();
            Set<TopicIdPartition> memberPartitions = assignment.computeIfAbsent(memberId, k -> new HashSet<>());
            // We are prepared to add one more partition, even if the desired assignment count is already reached.
            if (memberPartitions.size() <= desiredAssignmentCount) {
                memberPartitions.add(partition);
            } else {
                memberIdIterator.remove();
                partitionListIterator.previous();
            }
        }
    }

    private List<TopicIdPartition> computeTargetPartitions(
        Set<Uuid> subscribedTopicIds,
        SubscribedTopicDescriber subscribedTopicDescriber
    ) {
        List<TopicIdPartition> targetPartitions = new ArrayList<>();
        subscribedTopicIds.forEach(topicId -> {
            int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
            if (numPartitions == -1) {
                throw new PartitionAssignorException(
                    "Members are subscribed to topic " + topicId
                        + " which doesn't exist in the topic metadata."
                );
            }

            // Since we are returning a list here, we can keep it sorted
            // to add determinism while testing and iterating.
            targetPartitions.addAll(subscribedTopicDescriber.assignablePartitions(topicId).stream()
                .sorted()
                .map(partition -> new TopicIdPartition(topicId, partition))
                .toList()
            );
        });
        return targetPartitions;
    }

    private GroupAssignment groupAssignment(
        Map<String, Set<TopicIdPartition>> assignmentByMember,
        Collection<String> allGroupMembers
    ) {
        Map<String, MemberAssignment> members = new HashMap<>();
        for (Map.Entry<String, Set<TopicIdPartition>> entry : assignmentByMember.entrySet()) {
            Map<Uuid, Set<Integer>> targetPartitions = new HashMap<>();
            entry.getValue().forEach(targetPartition ->
                targetPartitions.computeIfAbsent(targetPartition.topicId(), k -> new HashSet<>()).add(targetPartition.partitionId()));
            members.put(entry.getKey(), new MemberAssignmentImpl(targetPartitions));
        }
        allGroupMembers.forEach(member -> {
            if (!members.containsKey(member))
                members.put(member, new MemberAssignmentImpl(new HashMap<>()));
        });

        return new GroupAssignment(members);
    }

    private static <K, V> HashMap<K, V> newHashMap(int numMappings) {
        return new HashMap<>((int) (((numMappings + 1) / 0.75f) + 1));
    }
}
