# -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# Licensed 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.
from __future__ import absolute_import

import argparse
import logging
import sys

import six
from six.moves import filter
from six.moves import range

from .cluster_balancer import ClusterBalancer
from .error import BrokerDecommissionError
from .error import EmptyReplicationGroupError
from .error import InvalidBrokerIdError
from .error import InvalidPartitionError
from .error import InvalidReplicationFactorError
from .error import NotEligibleGroupError
from .error import RebalanceError
from .util import compute_optimum
from .util import separate_groups


class PartitionCountBalancer(ClusterBalancer):
    """An implementation of cluster rebalancing that tries to achieve balance
    by considering the number of partitions and leaders on each broker.

    :param cluster_topology: The ClusterTopology object that should be acted
        on.
    :param args: The program arguments.
    """

    def __init__(self, cluster_topology, args):
        super(PartitionCountBalancer, self).__init__(cluster_topology, args)
        self.log = logging.getLogger(self.__class__.__name__)

    def _set_arg_default(self, arg, value):
        if not hasattr(self.args, arg):
            setattr(self.args, arg, value)

    def parse_args(self, balancer_args):
        self._set_arg_default('replication_groups', False)
        self._set_arg_default('brokers', False)
        self._set_arg_default('leaders', False)
        self._set_arg_default('max_partition_movements', None)
        self._set_arg_default('max_movement_size', None)
        self._set_arg_default('max_leader_changes', None)
        parser = argparse.ArgumentParser(
            prog=self.__class__.__name__,
            description='Balance the cluster based on the number of partitions'
            ' per broker and replication-group.',
        )
        parser.parse_args(balancer_args, self.args)

    def decommission_brokers(self, broker_ids):
        """Decommission a list of brokers trying to keep the replication group
        the brokers belong to balanced.

        :param broker_ids: list of string representing valid broker ids in the cluster
        :raises: InvalidBrokerIdError when the id is invalid.
        """
        groups = set()
        for b_id in broker_ids:
            try:
                broker = self.cluster_topology.brokers[b_id]
            except KeyError:
                self.log.error("Invalid broker id %s.", b_id)
                # Raise an error for now. As alternative we may ignore the
                # invalid id and continue with the others.
                raise InvalidBrokerIdError(
                    "Broker id {} does not exist in cluster".format(b_id),
                )
            broker.mark_decommissioned()
            groups.add(broker.replication_group)

        for group in groups:
            self._decommission_brokers_in_group(group)

    def _decommission_brokers_in_group(self, group):
        """Decommission the marked brokers of a group."""
        try:
            group.rebalance_brokers()
        except EmptyReplicationGroupError:
            self.log.warning("No active brokers left in replication group %s", group)
        for broker in group.brokers:
            if broker.decommissioned and not broker.empty():
                # In this case we need to reassign the remaining partitions
                # to other replication groups
                self.log.info(
                    "Broker %s can't be decommissioned within the same "
                    "replication group %s. Moving partitions to other "
                    "replication groups.",
                    broker,
                    broker.replication_group,
                )
                self._force_broker_decommission(broker)
                # Broker should be empty now
                if not broker.empty():
                    # Decommission may be impossible if there are not enough
                    # brokers to redistributed the replicas.
                    self.log.error(
                        "Could not decommission broker %s. "
                        "Partitions %s cannot be reassigned.",
                        broker,
                        broker.partitions,
                    )
                    raise BrokerDecommissionError("Broker decommission failed.")

    def _force_broker_decommission(self, broker):
        available_groups = [
            rg for rg in six.itervalues(self.cluster_topology.rgs)
            if rg is not broker.replication_group
        ]

        for partition in broker.partitions.copy():  # partitions set changes during loop
            groups = sorted(
                available_groups,
                key=lambda x: x.count_replica(partition),
            )
            for group in groups:
                self.log.debug(
                    "Try to move partition: %s from broker %s to "
                    "replication group %s",
                    partition,
                    broker,
                    broker.replication_group,
                )
                try:
                    group.acquire_partition(partition, broker)
                    break
                except NotEligibleGroupError:
                    pass

    def rebalance(self):
        if self.args.max_movement_size:
            self.log.error(
                '--max-movement-size can not be specified for {balancer}.'
                ' Exiting.'.format(
                    balancer=self.__class__.__name__,
                ),
            )
            sys.exit(1)

        if self.args.replication_groups:
            self.log.info(
                'Re-balancing replica-count over replication groups: %s',
                ', '.join(str(rg) for rg in self.cluster_topology.rgs.keys()),
            )
            self.rebalance_replication_groups()

        if self.args.brokers:
            self.log.info(
                'Re-balancing partition-count across brokers: %s',
                ', '.join(str(e) for e in self.cluster_topology.brokers.keys()),
            )
            self.rebalance_brokers()

        if self.args.leaders:
            self.log.info(
                'Re-balancing leader-count across brokers: %s',
                ', '.join(str(e) for e in self.cluster_topology.brokers.keys()),
            )
            self.rebalance_leaders()

    def rebalance_replication_groups(self):
        """Rebalance partitions over replication groups.

        First step involves rebalancing replica-count for each partition across
        replication-groups.
        Second step involves rebalancing partition-count across replication-groups
        of the cluster.