import contextlib
import logging
from collections import defaultdict
from collections.abc import (
    Collection,
    Iterable,
    Mapping,
    MutableSequence,
    Sequence,
    Sized,
)
from copy import deepcopy
from typing import (
    Any,
    NamedTuple,
    Optional,
)

from aiokafka.cluster import ClusterMetadata
from aiokafka.coordinator.assignors.abstract import AbstractPartitionAssignor
from aiokafka.coordinator.assignors.sticky.partition_movements import PartitionMovements
from aiokafka.coordinator.assignors.sticky.sorted_set import SortedSet
from aiokafka.coordinator.protocol import (
    ConsumerProtocolMemberAssignment,
    ConsumerProtocolMemberMetadata,
    Schema,
)
from aiokafka.protocol.struct import Struct
from aiokafka.protocol.types import Array, Int32, String
from aiokafka.structs import TopicPartition

log = logging.getLogger(__name__)


class ConsumerGenerationPair(NamedTuple):
    consumer: str
    generation: int


class ConsumerSubscription(NamedTuple):
    consumer: str
    partitions: Sequence[TopicPartition]


def has_identical_list_elements(list_: Sequence[list[Any]]) -> bool:
    """Checks if all lists in the collection have the same members

    Arguments:
      list_: collection of lists

    Returns:
      true if all lists in the collection have the same members; false otherwise
    """
    if not list_:
        return True
    return all(list_[i] == list_[i - 1] for i in range(1, len(list_)))


def subscriptions_comparator_key(element: tuple[str, Sized]) -> tuple[int, str]:
    return len(element[1]), element[0]


def partitions_comparator_key(
    element: tuple[TopicPartition, Sized],
) -> tuple[int, str, int]:
    return len(element[1]), element[0].topic, element[0].partition


def remove_if_present(collection: MutableSequence[Any], element: Any) -> None:
    with contextlib.suppress(ValueError, KeyError):
        collection.remove(element)


class StickyAssignorMemberMetadataV1(NamedTuple):
    subscription: list[str]
    partitions: list[TopicPartition]
    generation: int


class StickyAssignorUserDataV1(Struct):
    """
    Used for preserving consumer's previously assigned partitions
    list and sending it as user data to the leader during a rebalance
    """

    class PreviousAssignment(NamedTuple):
        topic: str
        partitions: list[int]

    previous_assignment: list[PreviousAssignment]
    generation: int

    SCHEMA = Schema(
        (
            "previous_assignment",
            Array(("topic", String("utf-8")), ("partitions", Array(Int32))),
        ),
        ("generation", Int32),
    )


class StickyAssignmentExecutor:
    def __init__(
        self,
        cluster: ClusterMetadata,
        members: dict[str, StickyAssignorMemberMetadataV1],
    ) -> None:
        self.members = members
        # a mapping between consumers and their assigned partitions that is updated
        # during assignment procedure
        self.current_assignment: dict[str, list[TopicPartition]] = defaultdict(list)
        # an assignment from a previous generation
        self.previous_assignment: dict[TopicPartition, ConsumerGenerationPair] = {}
        # a mapping between partitions and their assigned consumers
        self.current_partition_consumer: dict[TopicPartition, str] = {}
        # a flag indicating that there were no previous assignments performed ever
        self.is_fresh_assignment = False
        # a mapping of all topic partitions to all consumers that can be assigned to
        # them
        self.partition_to_all_potential_consumers: dict[TopicPartition, list[str]] = {}
        # a mapping of all consumers to all potential topic partitions that can be
        # assigned to them
        self.consumer_to_all_potential_partitions: dict[str, list[TopicPartition]] = {}
        # an ascending sorted set of consumers based on how many topic partitions are
        # already assigned to them
        self.sorted_current_subscriptions: SortedSet[ConsumerSubscription] = SortedSet()
        # an ascending sorted list of topic partitions based on how many consumers can
        # potentially use them
        self.sorted_partitions: list[TopicPartition] = []
        # all partitions that need to be assigned
        self.unassigned_partitions: list[TopicPartition] = []
        # a flag indicating that a certain partition cannot remain assigned to its
        # current consumer because the consumer is no longer subscribed to its topic
        self.revocation_required = False

        self.partition_movements = PartitionMovements()
        self._initialize(cluster)

    def perform_initial_assignment(self) -> None:
        self._populate_sorted_partitions()
        self._populate_partitions_to_reassign()

    def balance(self) -> None:
        self._initialize_current_subscriptions()
        initializing = (
            len(self.current_assignment[self._get_consumer_with_most_subscriptions()])
            == 0
        )

        # assign all unassigned partitions
        for partition in self.unassigned_partitions:
            # skip if there is no potential consumer for the partition
            if not self.partition_to_all_potential_consumers[partition]:
                continue
            self._assign_partition(partition)

        # narrow down the reassignment scope to only those partitions that can actually
        # be reassigned
        fixed_partitions = set()
        for partition in self.partition_to_all_potential_consumers:
            if not self._can_partition_participate_in_reassignment(partition):
                fixed_partitions.add(partition)
        for fixed_partition in fixed_partitions:
            remove_if_present(self.sorted_partitions, fixed_partition)
            remove_if_present(self.unassigned_partitions, fixed_partition)

        # narrow down the reassignment scope to only those consumers that are subject to
        # reassignment
        fixed_assignments: dict[str, list[TopicPartition]] = {}
        for consumer in self.consumer_to_all_potential_partitions:
            if not self._can_consumer_participate_in_reassignment(consumer):
                self._remove_consumer_from_current_subscriptions_and_maintain_order(
                    consumer
                )
                fixed_assignments[consumer] = self.current_assignment[consumer]
                del self.current_assignment[consumer]

        # create a deep copy of the current assignment so we can revert to it
        # if we do not get a more balanced assignment later
        prebalance_assignment = deepcopy(self.current_assignment)
        prebalance_partition_consumers = deepcopy(self.current_partition_consumer)

        # if we don't already need to revoke something due to subscription changes,
        # first try to balance by only moving newly added partitions
        if not self.revocation_required:
            self._perform_reassignments(self.unassigned_partitions)
        reassignment_performed = self._perform_reassignments(self.sorted_partitions)

        # if we are not preserving existing assignments and we have made changes to the
        # current assignment make sure we are getting a more balanced assignment;
        # otherwise, revert to previous assignment
        if (
            not initializing
            and reassignment_performed
            and self._get_balance_score(self.current_assignment)
            >= self._get_balance_score(prebalance_assignment)
        ):
            self.current_assignment = prebalance_assignment
            self.current_partition_consumer.clear()
            self.current_partition_consumer.update(prebalance_partition_consumers)

        # add the fixed assignments (those that could not change) back
        for consumer, partitions in fixed_assignments.items():
            self.current_assignment[consumer] = partitions
            self._add_consumer_to_current_subscriptions_and_maintain_order(consumer)

    def get_final_assignment(self, member_id: str) -> Collection[tuple[str, list[int]]]:
        assignment: dict[str, list[int]] = defaultdict(list)
        for topic_partition in self.current_assignment[member_id]:
            assignment[topic_partition.topic].append(topic_partition.partition)
        assignment = {k: sorted(v) for k, v in assignment.items()}
        return assignment.items()

    def _initialize(self, cluster: ClusterMetadata) -> None:
        self._init_current_assignments(self.members)

        for topic in cluster.topics():
            partitions = cluster.partitions_for_topic(topic)
            if partitions is None:
                log.warning("No partition metadata for topic %s", topic)
                continue
            for p in partitions:
                partition = TopicPartition(topic=topic, partition=p)
                self.partition_to_all_potential_consumers[partition] = []
        for consumer_id, member_metadata in self.members.items():
            self.consumer_to_all_potential_partitions[consumer_id] = []
            for topic in member_metadata.subscription:
                partitions_for_topic = cluster.partitions_for_topic(topic)
                if partitions_for_topic is None:
                    log.warning("No partition metadata for topic %r", topic)
                    continue
                for p in partitions_for_topic:
                    partition = TopicPartition(topic=topic, partition=p)
                    self.consumer_to_all_potential_partitions[consumer_id].append(
                        partition
                    )
                    self.partition_to_all_potential_consumers[partition].append(
                        consumer_id
                    )
            if consumer_id not in self.current_assignment:
                self.current_assignment[consumer_id] = []

    def _init_current_assignments(
        self, members: dict[str, StickyAssignorMemberMetadataV1]
    ) -> None:
        # we need to process subscriptions' user data with each consumer's reported
        # generation in mind higher generations overwrite lower generations in case of
        # a conflict note that a conflict could exists only if user data is for
        # different generations

        # for each partition we create a map of its consumers by generation
        sorted_partition_consumers_by_generation: dict[
            TopicPartition, dict[int, str]
        ] = {}
        for consumer, member_metadata in members.items():
            for partition in member_metadata.partitions:
                if partition in sorted_partition_consumers_by_generation:
                    consumers = sorted_partition_consumers_by_generation[partition]
                    if (
                        member_metadata.generation
                        and member_metadata.generation in consumers
                    ):
                        # same partition is assigned to two consumers during the same
                        # rebalance. log a warning and skip this record
                        log.warning(
                            "Partition %r is assigned to multiple consumers "
                            "following sticky assignment generation %d",
                            partition,
                            member_metadata.generation,
                        )
                    else:
                        consumers[member_metadata.generation] = consumer
                else:
                    sorted_consumers = {member_metadata.generation: consumer}
                    sorted_partition_consumers_by_generation[partition] = (
                        sorted_consumers
                    )

        # previous_assignment holds the prior ConsumerGenerationPair (before current) of
        # each partition current and previous consumers are the last two consumers of
        # each partition in the above sorted map
        for partition, consumers in sorted_partition_consumers_by_generation.items():
            generations = sorted(consumers.keys(), reverse=True)
            self.current_assignment[consumers[generations[0]]].append(partition)
            # now update previous assignment if any
            if len(generations) > 1:
                self.previous_assignment[partition] = ConsumerGenerationPair(
                    consumer=consumers[generations[1]], generation=generations[1]
                )

        self.is_fresh_assignment = len(self.current_assignment) == 0

        for consumer_id, partitions in self.current_assignment.items():
            for partition in partitions:
                self.current_partition_consumer[partition] = consumer_id

    def _are_subscriptions_identical(self) -> bool:
        """
        Returns:
            true, if both potential consumers of partitions and potential partitions
            that consumers can consume are the same
        """
        if not has_identical_list_elements(
            list(self.partition_to_all_potential_consumers.values())
        ):
            return False
        return has_identical_list_elements(
            list(self.consumer_to_all_potential_partitions.values())
        )

    def _populate_sorted_partitions(self) -> None:
        # set of topic partitions with their respective potential consumers
        all_partitions = {
            (tp, tuple(consumers))
            for tp, consumers in self.partition_to_all_potential_consumers.items()
        }
        partitions_sorted_by_num_of_potential_consumers = sorted(
            all_partitions, key=partitions_comparator_key
        )

        self.sorted_partitions = []
        if not self.is_fresh_assignment and self._are_subscriptions_identical():
            # if this is a reassignment and the subscriptions are identical (all
            # consumers can consumer from all topics) then we just need to simply list
            # partitions in a round robin fashion (from consumers with most assigned
            # partitions to those with least)
            assignments = deepcopy(self.current_assignment)
            for partitions in assignments.values():
                partitions[:] = [
                    partition
                    for partition in partitions
                    if partition in self.partition_to_all_potential_consumers
                ]

            sorted_consumers = SortedSet(
                iterable=[
                    (consumer, tuple(partitions))
                    for consumer, partitions in assignments.items()
                ],
                key=subscriptions_comparator_key,
            )
            # at this point, sorted_consumers contains an ascending-sorted list of
            # consumers based on how many valid partitions are currently assigned to
            # them

            while sorted_consumers:
                # take the consumer with the most partitions
                consumer, _ = sorted_consumers.pop_last()
                # currently assigned partitions to this consumer
                remaining_partitions = assignments[consumer]
                # from partitions that had a different consumer before,
                # keep only those that are assigned to this consumer now
                previous_partitions = set(self.previous_assignment.keys()).intersection(
                    set(remaining_partitions)
                )
                if previous_partitions:
                    # if there is a partition of this consumer that was assigned to
                    # another consumer before mark it as good options for reassignment
                    partition = previous_partitions.pop()
                    remaining_partitions.remove(partition)
                    self.sorted_partitions.append(partition)
                    sorted_consumers.add((consumer, tuple(assignments[consumer])))
                elif remaining_partitions:
                    # otherwise, mark any other one of the current partitions as a
                    # reassignment candidate
                    self.sorted_partitions.append(remaining_partitions.pop())
                    sorted_consumers.add((consumer, tuple(assignments[consumer])))

            while partitions_sorted_by_num_of_potential_consumers:
                partition = partitions_sorted_by_num_of_potential_consumers.pop(0)[0]
                if partition not in self.sorted_partitions:
                    self.sorted_partitions.append(partition)
        else:
            while partitions_sorted_by_num_of_potential_consumers:
                self.sorted_partitions.append(
                    partitions_sorted_by_num_of_potential_consumers.pop(0)[0]
                )

    def _populate_partitions_to_reassign(self) -> None:
        self.unassigned_partitions = deepcopy(self.sorted_partitions)

        assignments_to_remove: list[str] = []
        for consumer_id, partitions in self.current_assignment.items():
            if consumer_id not in self.members:
                # if a consumer that existed before (and had some partition assignments)
                # is now removed, remove it from current_assignment
                for partition in partitions:
                    del self.current_partition_consumer[partition]
                assignments_to_remove.append(consumer_id)
            else:
                # otherwise (the consumer still exists)
                partitions_to_remove = []
                for partition in partitions:
                    if partition not in self.partition_to_all_potential_consumers:
                        # if this topic partition of this consumer no longer exists
                        # remove it from current_assignment of the consumer
                        partitions_to_remove.append(partition)
                    elif partition.topic not in self.members[consumer_id].subscription:
                        # if this partition cannot remain assigned to its current
                        # consumer because the consumer is no longer subscribed to its
                        # topic remove it from current_assignment of the consumer
                        partitions_to_remove.append(partition)
                        self.revocation_required = True
                    else:
                        # otherwise, remove the topic partition from those that need to
                        # be assigned only if its current consumer is still subscribed
                        # to its topic (because it is already assigned and we would want
                        # to preserve that assignment as much as possible)
                        self.unassigned_partitions.remove(partition)
                for partition in partitions_to_remove:
                    self.current_assignment[consumer_id].remove(partition)
                    del self.current_partition_consumer[partition]
        for consumer_id in assignments_to_remove:
            del self.current_assignment[consumer_id]

    def _initialize_current_subscriptions(self) -> None:
        self.sorted_current_subscriptions = SortedSet(
            iterable=[
                ConsumerSubscription(consumer=consumer, partitions=tuple(partitions))
                for consumer, partitions in self.current_assignment.items()
            ],
            key=subscriptions_comparator_key,
        )

    def _get_consumer_with_least_subscriptions(self) -> str:
        if current_subscription := self.sorted_current_subscriptions.first():
            return current_subscription[0]
        raise ValueError("sorted_current_subscriptions is empty")

    def _get_consumer_with_most_subscriptions(self) -> str:
        if current_subscription := self.sorted_current_subscriptions.last():
            return current_subscription[0]
        raise ValueError("sorted_current_subscriptions is empty")

    def _remove_consumer_from_current_subscriptions_and_maintain_order(
        self, consumer: str
    ) -> None:
        self.sorted_current_subscriptions.remove(
            ConsumerSubscription(
                consumer=consumer,
                partitions=tuple(self.current_assignment[consumer]),
            )
        )

    def _add_consumer_to_current_subscriptions_and_maintain_order(
        self, consumer: str
    ) -> None:
        self.sorted_current_subscriptions.add(
            ConsumerSubscription(
                consumer=consumer,
                partitions=tuple(self.current_assignment[consumer]),
            )
        )

    def _is_balanced(self) -> bool:
        """Determines if the current assignment is a balanced one"""
        if (
            len(self.current_assignment[self._get_consumer_with_least_subscriptions()])
            >= len(
                self.current_assignment[self._get_consumer_with_most_subscriptions()]
            )
            - 1
        ):
            # if minimum and maximum numbers of partitions assigned to consumers differ
            # by at most one return true
            return True

        # create a mapping from partitions to the consumer assigned to them
        all_assigned_partitions = {}
        for consumer_id, consumer_partitions in self.current_assignment.items():
            for partition in consumer_partitions:
                if partition in all_assigned_partitions:
                    log.error(
                        "Partition %r is assigned to more than one consumer", partition
                    )
                all_assigned_partitions[partition] = consumer_id

        # for each consumer that does not have all the topic partitions it can get
        # make sure none of the topic partitions it could but did not get cannot be
        # moved to it (because that would break the balance)
        for consumer, _ in self.sorted_current_subscriptions:
            consumer_partition_count = len(self.current_assignment[consumer])
            # skip if this consumer already has all the topic partitions it can get
            if consumer_partition_count == len(
                self.consumer_to_all_potential_partitions[consumer]
            ):
                continue

            # otherwise make sure it cannot get any more
            for partition in self.consumer_to_all_potential_partitions[consumer]:
                if partition not in self.current_assignment[consumer]:
                    other_consumer = all_assigned_partitions[partition]
                    other_consumer_partition_count = len(
                        self.current_assignment[other_consumer]
                    )
                    if consumer_partition_count < other_consumer_partition_count:
                        return False
        return True

    def _assign_partition(self, partition: TopicPartition) -> None:
        for consumer, _ in self.sorted_current_subscriptions:
            if partition in self.consumer_to_all_potential_partitions[consumer]:
                self._remove_consumer_from_current_subscriptions_and_maintain_order(
                    consumer
                )
                self.current_assignment[consumer].append(partition)
                self.current_partition_consumer[partition] = consumer
                self._add_consumer_to_current_subscriptions_and_maintain_order(consumer)
                break

    def _can_partition_participate_in_reassignment(
        self, partition: TopicPartition
    ) -> bool:
        return len(self.partition_to_all_potential_consumers[partition]) >= 2

    def _can_consumer_participate_in_reassignment(self, consumer: str) -> bool:
        current_partitions = self.current_assignment[consumer]
        current_assignment_size = len(current_partitions)
        max_assignment_size = len(self.consumer_to_all_potential_partitions[consumer])
        if current_assignment_size > max_assignment_size:
            log.error(
                "The consumer %r is assigned more partitions than the maximum "
                "possible",
                consumer,
            )
        if current_assignment_size < max_assignment_size:
            # if a consumer is not assigned all its potential partitions it is subject
            # to reassignment
            return True
        for partition in current_partitions:
            # if any of the partitions assigned to a consumer is subject to reassignment
            # the consumer itself is subject to reassignment
            if self._can_partition_participate_in_reassignment(partition):
                return True
        return False

    def _perform_reassignments(
        self, reassignable_partitions: list[TopicPartition]
    ) -> bool:
        reassignment_performed = False

        # repeat reassignment until no partition can be moved to improve the balance
        while True:
            modified = False
            # reassign all reassignable partitions until the full list is processed or
            # a balance is achieved (starting from the partition with least potential
            # consumers and if needed)
            for partition in reassignable_partitions:
                if self._is_balanced():
                    break
                # the partition must have at least two potential consumers
                if len(self.partition_to_all_potential_consumers[partition]) <= 1:
                    log.error(
                        "Expected more than one potential consumer for partition %r",
                        partition,
                    )
                # the partition must have a current consumer
                consumer = self.current_partition_consumer.get(partition)
                if consumer is None:
                    log.error(
                        "Expected partition %r to be assigned to a consumer", partition
                    )
                    continue

                if (
                    partition in self.previous_assignment
                    and len(self.current_assignment[consumer])
                    > len(
                        self.current_assignment[
                            self.previous_assignment[partition].consumer
                        ]
                    )
                    + 1
                ):
                    self._reassign_partition_to_consumer(
                        partition,
                        self.previous_assignment[partition].consumer,
                    )
                    reassignment_performed = True
                    modified = True
                    continue

                # check if a better-suited consumer exist for the partition; if so,
                # reassign it
                for other_consumer in self.partition_to_all_potential_consumers[
                    partition
                ]:
                    if (
                        len(self.current_assignment[consumer])
                        > len(self.current_assignment[other_consumer]) + 1
                    ):
                        self._reassign_partition(partition)
                        reassignment_performed = True
                        modified = True
                        break

            if not modified:
                break
        return reassignment_performed

    def _reassign_partition(self, partition: TopicPartition) -> None:
        new_consumer = None
        for another_consumer, _ in self.sorted_current_subscriptions:
            if partition in self.consumer_to_all_potential_partitions[another_consumer]:
                new_consumer = another_consumer
                break
        assert new_consumer is not None
        self._reassign_partition_to_consumer(partition, new_consumer)

    def _reassign_partition_to_consumer(
        self, partition: TopicPartition, new_consumer: str
    ) -> None:
        consumer = self.current_partition_consumer[partition]
        # find the correct partition movement considering the stickiness requirement
        partition_to_be_moved = self.partition_movements.get_partition_to_be_moved(
            partition, consumer, new_consumer
        )
        self._move_partition(partition_to_be_moved, new_consumer)

    def _move_partition(self, partition: TopicPartition, new_consumer: str) -> None:
        old_consumer = self.current_partition_consumer[partition]
        self._remove_consumer_from_current_subscriptions_and_maintain_order(
            old_consumer
        )
        self._remove_consumer_from_current_subscriptions_and_maintain_order(
            new_consumer
        )

        self.partition_movements.move_partition(partition, old_consumer, new_consumer)

        self.current_assignment[old_consumer].remove(partition)
        self.current_assignment[new_consumer].append(partition)
        self.current_partition_consumer[partition] = new_consumer

        self._add_consumer_to_current_subscriptions_and_maintain_order(new_consumer)
        self._add_consumer_to_current_subscriptions_and_maintain_order(old_consumer)

    @staticmethod
    def _get_balance_score(assignment: dict[str, list[TopicPartition]]) -> int:
        """Calculates a balance score of a give assignment
        as the sum of assigned partitions size difference of all consumer pairs.
        A perfectly balanced assignment (with all consumers getting the same number of
        partitions) has a balance score of 0. Lower balance score indicates a more
        balanced assignment.

        Arguments:
          assignment (dict): {consumer: list of assigned topic partitions}

        Returns:
          the balance score of the assignment
        """
        score = 0
        consumer_to_assignment: dict[str, int] = {}
        for consumer_id, partitions in assignment.items():
            consumer_to_assignment[consumer_id] = len(partitions)

        consumers_to_explore = set(consumer_to_assignment.keys())
        for consumer_id in consumer_to_assignment:
            if consumer_id in consumers_to_explore:
                consumers_to_explore.remove(consumer_id)
                for other_consumer_id in consumers_to_explore:
                    score += abs(
                        consumer_to_assignment[consumer_id]
                        - consumer_to_assignment[other_consumer_id]
                    )
        return score


class StickyPartitionAssignor(AbstractPartitionAssignor):
    """
    https://cwiki.apache.org/confluence/display/KAFKA/KIP-54+-+Sticky+Partition+Assignment+Strategy

    The sticky assignor serves two purposes. First, it guarantees an assignment that is
    as balanced as possible, meaning either:
    - the numbers of topic partitions assigned to consumers differ by at most one; or
    - each consumer that has 2+ fewer topic partitions than some other consumer cannot
    get any of those topic partitions transferred to it.

    Second, it preserved as many existing assignment as possible when a reassignment
    occurs. This helps in saving some of the overhead processing when topic partitions
    move from one consumer to another.

    Starting fresh it would work by distributing the partitions over consumers as evenly
    as possible. Even though this may sound similar to how round robin assignor works,
    the second example below shows that it is not. During a reassignment it would
    perform the reassignment in such a way that in the new assignment
    - topic partitions are still distributed as evenly as possible, and
    - topic partitions stay with their previously assigned consumers as much as
    possible.

    The first goal above takes precedence over the second one.

    Example 1.
    Suppose there are three consumers C0, C1, C2,
    four topics t0, t1, t2, t3, and each topic has 2 partitions,
    resulting in partitions t0p0, t0p1, t1p0, t1p1, t2p0, t2p1, t3p0, t3p1.
    Each consumer is subscribed to all three topics.

    The assignment with both sticky and round robin assignors will be:
    - C0: [t0p0, t1p1, t3p0]
    - C1: [t0p1, t2p0, t3p1]
    - C2: [t1p0, t2p1]

    Now, let's assume C1 is removed and a reassignment is about to happen. The round
    robin assignor would produce:
    - C0: [t0p0, t1p0, t2p0, t3p0]
    - C2: [t0p1, t1p1, t2p1, t3p1]

    while the sticky assignor would result in:
    - C0 [t0p0, t1p1, t3p0, t2p0]
    - C2 [t1p0, t2p1, t0p1, t3p1]
    preserving all the previous assignments (unlike the round robin assignor).


    Example 2.
    There are three consumers C0, C1, C2,
    and three topics t0, t1, t2, with 1, 2, and 3 partitions respectively.
    Therefore, the partitions are t0p0, t1p0, t1p1, t2p0, t2p1, t2p2.
    C0 is subscribed to t0;
    C1 is subscribed to t0, t1;
    and C2 is subscribed to t0, t1, t2.

    The round robin assignor would come up with the following assignment:
    - C0 [t0p0]
    - C1 [t1p0]
    - C2 [t1p1, t2p0, t2p1, t2p2]

    which is not as balanced as the assignment suggested by sticky assignor:
    - C0 [t0p0]
    - C1 [t1p0, t1p1]
    - C2 [t2p0, t2p1, t2p2]

    Now, if consumer C0 is removed, these two assignors would produce the following
    assignments. Round Robin (preserves 3 partition assignments):
    - C1 [t0p0, t1p1]
    - C2 [t1p0, t2p0, t2p1, t2p2]

    Sticky (preserves 5 partition assignments):
    - C1 [t1p0, t1p1, t0p0]
    - C2 [t2p0, t2p1, t2p2]
    """

    DEFAULT_GENERATION_ID = -1

    name = "sticky"
    version = 0

    member_assignment: Optional[list[TopicPartition]] = None
    generation: int = DEFAULT_GENERATION_ID

    _latest_partition_movements: Optional[PartitionMovements] = None

    @classmethod
    def assign(
        cls,
        cluster: ClusterMetadata,
        members: Mapping[str, ConsumerProtocolMemberMetadata],
    ) -> dict[str, ConsumerProtocolMemberAssignment]:
        """Performs group assignment given cluster metadata and member subscriptions

        Arguments:
            cluster (ClusterMetadata): cluster metadata
            members (dict of {member_id: MemberMetadata}): decoded metadata for each
            member in the group.

        Returns:
          dict: {member_id: MemberAssignment}
        """
        members_metadata: dict[str, StickyAssignorMemberMetadataV1] = {}
        for consumer, member_metadata in members.items():
            members_metadata[consumer] = cls.parse_member_metadata(member_metadata)

        executor = StickyAssignmentExecutor(cluster, members_metadata)
        executor.perform_initial_assignment()
        executor.balance()

        cls._latest_partition_movements = executor.partition_movements

        assignment: dict[str, ConsumerProtocolMemberAssignment] = {}
        for member_id in members:
            assignment[member_id] = ConsumerProtocolMemberAssignment(
                cls.version,
                sorted(executor.get_final_assignment(member_id)),
                b"",
            )
        return assignment

    @classmethod
    def parse_member_metadata(
        cls, metadata: ConsumerProtocolMemberMetadata
    ) -> StickyAssignorMemberMetadataV1:
        """
        Parses member metadata into a python object.
        This implementation only serializes and deserializes the
        StickyAssignorMemberMetadataV1 user data, since no StickyAssignor written in
        Python was deployed ever in the wild with version V0, meaning that there is no
        need to support backward compatibility with V0.

        Arguments:
          metadata (MemberMetadata): decoded metadata for a member of the group.

        Returns:
          parsed metadata (StickyAssignorMemberMetadataV1)
        """
        user_data = metadata.user_data
        if not user_data:
            return StickyAssignorMemberMetadataV1(
                partitions=[],
                generation=cls.DEFAULT_GENERATION_ID,
                subscription=metadata.subscription,
            )

        try:
            decoded_user_data = StickyAssignorUserDataV1.decode(user_data)
        except Exception as e:  # noqa: BLE001
            # ignore the consumer's previous assignment if it cannot be parsed
            log.error("Could not parse member data: %r", e)
            return StickyAssignorMemberMetadataV1(
                partitions=[],
                generation=cls.DEFAULT_GENERATION_ID,
                subscription=metadata.subscription,
            )

        member_partitions: list[TopicPartition] = []
        for (
            topic,
            partitions,
        ) in decoded_user_data.previous_assignment:
            member_partitions.extend(
                [TopicPartition(topic, partition) for partition in partitions]
            )
        return StickyAssignorMemberMetadataV1(
            partitions=member_partitions,
            generation=decoded_user_data.generation,
            subscription=metadata.subscription,
        )

    @classmethod
    def metadata(cls, topics: Iterable[str]) -> ConsumerProtocolMemberMetadata:
        return cls._metadata(topics, cls.member_assignment, cls.generation)

    @classmethod
    def _metadata(
        cls,
        topics: Iterable[str],
        member_assignment_partitions: Optional[list[TopicPartition]],
        generation: int = -1,
    ) -> ConsumerProtocolMemberMetadata:
        if member_assignment_partitions is None:
            log.debug("No member assignment available")
            user_data = b""
        else:
            log.debug(
                "Member assignment is available, generating the metadata: "
                "generation %d",
                cls.generation,
            )
            partitions_by_topic = defaultdict(list)
            for topic_partition in member_assignment_partitions:
                partitions_by_topic[topic_partition.topic].append(
                    topic_partition.partition
                )
            data = StickyAssignorUserDataV1(partitions_by_topic.items(), generation)
            user_data = data.encode()
        return ConsumerProtocolMemberMetadata(cls.version, list(topics), user_data)

    @classmethod
    def on_assignment(cls, assignment: ConsumerProtocolMemberAssignment) -> None:
        """Callback that runs on each assignment. Updates assignor's state.

        Arguments:
          assignment: MemberAssignment
        """
        log.debug("On assignment: assignment=%r", assignment)
        cls.member_assignment = assignment.partitions()

    @classmethod
    def on_generation_assignment(cls, generation: int) -> None:
        """Callback that runs on each assignment. Updates assignor's generation id.

        Arguments:
          generation: generation id
        """
        log.debug("On generation assignment: generation=%d", generation)
        cls.generation = generation
