# -*- cpy-indent-level: 4; indent-tabs-mode: nil -*-
# ex: set expandtab softtabstop=4 shiftwidth=4:
#
# Copyright (C) 2009,2010,2011,2012,2013,2014,2015,2016  Contributor
#
# 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.
"""Provides various utilities around services."""

from collections import defaultdict
from random import choice

from sqlalchemy.orm.session import object_session

from aquilon.exceptions_ import ArgumentError, InternalError
from aquilon.aqdb.model import Host, Cluster, ServiceMap, MetaCluster
from aquilon.worker.templates import PlenaryServiceInstanceServer


class ChooserCache(object):
    def __init__(self):
        self.service_maps = defaultdict(dict)
        self.client_counts = dict()

    def cache_service_maps(self, services, key):
        # key is (personality, location, network)
        missing_services = set(srv for srv in services
                               if key not in self.service_maps[srv])
        if not missing_services:
            return

        maps = ServiceMap.get_mapped_instance_cache(missing_services, key[0],
                                                    key[1], key[2])
        for dbsrv in maps:
            self.service_maps[dbsrv][key] = maps[dbsrv]

    def client_count(self, dbsi):
        try:
            return self.client_counts[dbsi]
        except KeyError:
            self.client_counts[dbsi] = dbsi.client_count
            return self.client_counts[dbsi]

    def allocate(self, dbsi, footprint):
        if dbsi not in self.client_counts:
            self.client_counts[dbsi] = dbsi.client_count
        self.client_counts[dbsi] += footprint

    def unallocate(self, dbsi, footprint):
        if dbsi not in self.client_counts:
            self.client_counts[dbsi] = dbsi.client_count
        self.client_counts[dbsi] -= footprint


class Chooser(object):
    """Helper for choosing services for an object."""

    def __new__(cls, dbobj, *args, **kwargs):
        if isinstance(dbobj, Host):
            chooser = super(Chooser, HostChooser).__new__(HostChooser)
        elif isinstance(dbobj, Cluster):
            chooser = super(Chooser, ClusterChooser).__new__(ClusterChooser)
        else:
            # Just assume the consumer invoked the right subclass...
            chooser = super(Chooser, cls).__new__(cls)

        # Lock the owner in the DB to avoid problems with parallel runs
        dbobj.lock_row()

        return chooser

    def __init__(self, dbobj, plenaries, logger, required_only=False, cache=None):
        """Initialize the chooser.

        To clear out bindings that are not required, pass in
        required_only=True.

        Several staging areas and caches are set up within this object.
        The general flow is that potential service instance choices
        are kept in staging_services (dictionary of service to list of
        service instances) and finalized into chosen_services (dictionary
        of service to single service instance).

        The original state of the object is held in the cache
        original_service_instances (dictionary of service to single service
        instance).

        The instances_bound and instances_unbound lists are populated
        after chosen_services with the differences between chosen_services
        and original_service_instances.

        Subclasses should call this before starting their own
        initialization.

        """
        logger.debug("Creating service chooser for {0:l}.".format(dbobj))

        self.dbobj = dbobj
        self.session = object_session(dbobj)
        self.required_only = required_only
        self.logger = logger
        self.network = None
        self.location = None
        self.map_key = None

        if cache is None:
            self.cache = ChooserCache()
        else:
            self.cache = cache

        # Stores interim service instance lists
        self.staging_services = {}

        # Report as many errors as possible in one shot
        self.errors = []

        # Cache the servers backing service instances, used to determine
        # affinity
        self.servers = set()

        # Set of service instances with a new client
        self.instances_bound = set()

        # Set of service instances losing a client
        self.instances_unbound = set()

        # Track the chosen services
        self.chosen_services = {}

        # Keep stashed plenaries for rollback purposes
        self.plenaries = plenaries

        # Cache the original service bindings
        self.original_service_instances = {}

        # Services which should be aligned to the parent object (cluster,
        # metacluster etc.)
        self.aligned_services = {}

        # Parent object whose services we should align to
        self.aligned_parents = []

        # Cache of any already bound services (keys) and the instance
        # that was bound (values).
        for si in dbobj.services_used:
            self.original_service_instances[si.service] = si
            logger.debug("{0} original binding: {1}".format(dbobj, si))

        self.required_services = dbobj.required_services

    def apply_changes(self):

        for instance in self.instances_unbound:
            self.logger.client_info("{0} removing binding for {1:l}"
                                    .format(self.dbobj, instance))
            self.dbobj.services_used.remove(instance)
            self.cache.unallocate(instance, self.get_footprint(instance))
        for instance in self.instances_bound:
            self.logger.client_info("{0} adding binding for {1:l}"
                                    .format(self.dbobj, instance))
            self.dbobj.services_used.append(instance)
            self.cache.allocate(instance, self.get_footprint(instance))

    def verify_init(self):
        """This is more of a verify-and-finalize method..."""
        if not self.required_only:
            for (service, instance) in self.original_service_instances.items():
                self.staging_services[service] = [instance]

        for parent in self.aligned_parents:
            for si in parent.services_used:
                self.aligned_services[si.service] = (si, parent)

            for service in parent.required_services:
                if service not in self.aligned_services:
                    # Don't just error here because the error() call
                    # has not yet been set up.  Will error out later.
                    self.aligned_services[service] = (None, parent)

        self.map_key = (self.dbobj.personality_stage, self.location,
                        self.network)

    def error(self, msg, *args, **kwargs):
        """Errors are consolidated so that many can be reported at once."""
        formatted = msg % args
        self.errors.append(formatted)
        self.logger.info(msg, *args, **kwargs)

    def set_required(self):
        """Main entry point when setting the required services for a host."""
        self.verify_init()
        self.prestash_primary()
        self.logger.debug("Setting required services")
        self.cache.cache_service_maps(self.required_services, self.map_key)
        for dbservice in self.required_services:
            self.find_service_instances(dbservice)
        self.check_errors()
        for dbservice in self.required_services:
            self.choose_aligned(dbservice)
            self.choose_available_capacity(dbservice)
            self.choose_past_use(dbservice)
        self.check_errors()
        # If this code needs to be made more efficient, this could
        # be refactored.  We don't always need count_servers()...
        # In theory don't always need the loop above, either.
        self.count_servers()
        for dbservice in self.required_services:
            self.reduce_service_instances(dbservice)
        self.finalize_service_instances()
        self.analyze_changes()
        self.stash_services()
        self.apply_changes()
        self.check_errors()

    def set_single(self, service, instance=None, force=False):
        """Use this to update a single service.

        If planning to use this method, construct the Chooser with
        required_only=False.  If required_only is True, all other
        bindings will be cleared.

        """
        self.verify_init()
        # When modifying just one service, we allow the object to be left in an
        # incomplete state
        self.prestash_primary(allow_incomplete=True)
        if instance:
            self.logger.debug("Setting service %s instance %s",
                              service.name, instance.name)
            self.staging_services[service] = [instance]
        else:
            self.logger.debug("Setting service %s with auto-bind",
                              service.name)
            self.staging_services[service] = None
            self.cache.cache_service_maps([service], self.map_key)
            self.find_service_instances(service)
        self.check_errors()
        self.choose_aligned(service)
        self.choose_available_capacity(service)
        self.check_errors()
        self.choose_past_use(service)
        # If this code needs to be made more efficient, this could
        # be refactored.  We don't always need count_servers()...
        self.count_servers()
        self.reduce_service_instances(service)
        self.finalize_service_instances()
        self.analyze_changes()
        if not force and self.instances_bound and self.instances_unbound:
            si = list(self.instances_unbound)[0]
            self.error("{0} is already bound to {1:l}, use unbind "
                       "to clear first or rebind to force."
                       .format(self.dbobj, si))
            self.check_errors()
        self.stash_services()
        self.apply_changes()
        self.check_errors()

    def find_service_instances(self, dbservice):
        """This finds the "closest" service instances, based on the known maps.

        It expects that cache_service_maps has been run.

        """
        instances = self.cache.service_maps[dbservice].get(self.map_key, [])
        if instances:
            for instance in instances:
                self.logger.debug("Found {0:l} in the maps.".format(instance))
            self.staging_services[dbservice] = instances
            return
        self.error("Could not find a relevant service map for {0:l} "
                   "on {1:l}".format(dbservice, self.dbobj))

    def check_errors(self):
        if self.errors:
            raise ArgumentError("\n".join(self.errors))

    def choose_aligned(self, dbservice):
        if dbservice not in self.aligned_services:
            return

        si, parent = self.aligned_services[dbservice]
        if not si:
            self.error("No instance set for {0:l} aligned service {1.name}.  "
                       "Please run `make cluster --cluster {0.name}` "
                       "to resolve.".format(parent, dbservice))
            return

        # This check is necessary to prevent bind_client from overriding
        # the cluster's binding.  The error message will be misleading...
        if si not in self.staging_services[dbservice]:
            self.error("{0} is set to use {1:l}, but that instance is not in a "
                       "service map for {2:l}.".format(parent, si, self.dbobj))
            return
        self.logger.debug("Chose {0:l} because it is bound to {1:l}."
                          .format(si, parent))
        self.staging_services[dbservice] = [si]

    def get_footprint(self, instance):  # pylint: disable=W0613
        return 1

    def instance_full(self, instance, max_clients, current_clients):
        """Check if the instance is effectively full.

        This check is complicated because clusters have a larger impact
        than a single host does.

        """
        if max_clients is None:
            return False
        if instance == self.original_service_instances.get(instance.service):
            if current_clients > max_clients:
                return True
            return False
        return current_clients + self.get_footprint(instance) > max_clients

    def choose_available_capacity(self, dbservice):
        """Verify that the available instances have spare capacity.

        Error out if none should be used.

        """
        maxed_out_instances = set()
        for instance in self.staging_services[dbservice][:]:
            max_clients = instance.enforced_max_clients
            current_clients = self.cache.client_count(instance)
            if self.instance_full(instance, max_clients, current_clients):
                self.staging_services[dbservice].remove(instance)
                maxed_out_instances.add(instance)
                self.logger.debug("Rejected service %s instance %s with "
                                  "max_client value of %s since client_count "
                                  "is %s.",
                                  instance.service.name, instance.name,
                                  max_clients, current_clients)
        if len(self.staging_services[dbservice]) < 1:
            self.error("The available instances %s for service %s are "
                       "at full capacity.",
                       [instance.name for instance in maxed_out_instances],
                       dbservice.name)
        return

    def choose_past_use(self, dbservice):
        """If more than one service instance was found in the maps,
        this method checks to see if we can reduce the list to a single
        choice by checking to see if any of the instances was already in use.

        """
        if len(self.staging_services[dbservice]) > 1 and \
           dbservice in self.original_service_instances and \
           self.original_service_instances[dbservice] in \
           self.staging_services[dbservice]:
            instance = self.original_service_instances[dbservice]
            self.logger.debug("Chose {0:l} because of past use."
                              .format(instance))
            self.staging_services[dbservice] = [instance]
        return

    def append_server_object(self, container, srv):
        """Append server objects to a container.

        This is a helper used for implementing server affinity.

        """
        if srv.host:
            container.add(srv.host)
        if srv.cluster:
            container.add(srv.cluster)

        # If the service is not bound to a host or cluster, then we have
        # no choice than to use the alias for providing server affinity
        if srv.alias and not srv.host and not srv.cluster:
            container.add(srv.alias)

    def count_servers(self, dbservice=None):
        """Get a count of the number of times a server backs
        service instances in use by this host.

        This method is called both to initialize the count and to update
        it as service instances are locked in.

        """
        if dbservice:
            instance_lists = [self.staging_services[dbservice]]
        else:
            instance_lists = self.staging_services.values()

        for instances in instance_lists:
            if len(instances) > 1:
                # Ignore any services where an instance has not been chosen.
                continue
            for srv in instances[0].servers:
                self.append_server_object(self.servers, srv)

    def reduce_service_instances(self, dbservice):
        if len(self.staging_services[dbservice]) == 1:
            self.count_servers(dbservice)
            return
        self.choose_affinity(dbservice)
        if len(self.staging_services[dbservice]) == 1:
            self.count_servers(dbservice)
            return
        self.choose_least_loaded(dbservice)
        if len(self.staging_services[dbservice]) == 1:
            self.count_servers(dbservice)
            return
        self.choose_random(dbservice)
        self.count_servers(dbservice)
        return

    def choose_affinity(self, dbservice):
        """Attempt to choose a service based on server affinity,
        also known as stickiness.

        This could be extremely complicated when trying to deal with
        instances backed by multiple servers.  Starting simple.
        Count the number of servers backing this instance that
        back other instances used by client.  Any instance that does
        not have the largest count gets tossed.

        """
        max_servers = 0
        max_instances = []
        for instance in self.staging_services[dbservice]:
            self.logger.debug("Checking service %s instance %s servers %s",
                              instance.service.name, instance.name,
                              [srv.fqdn for srv in instance.servers])
            instance_servers = set()
            for srv in instance.servers:
                self.append_server_object(instance_servers, srv)

            common_servers = self.servers & instance_servers
            if len(common_servers) > max_servers:
                max_servers = len(common_servers)
                max_instances = [instance]
            elif len(common_servers) == max_servers:
                max_instances.append(instance)
        if max_instances and \
           len(max_instances) < len(self.staging_services[dbservice]):
            for instance in self.staging_services[dbservice]:
                if instance not in max_instances:
                    self.logger.debug("Discounted {0:l} due to server affinity "
                                      "(stickiness).".format(instance))
            self.staging_services[dbservice] = max_instances

    def choose_least_loaded(self, dbservice):
        """Choose a service instance based on load."""
        least_clients = None
        least_loaded = []
        for instance in self.staging_services[dbservice]:
            client_count = self.cache.client_count(instance)
            if not least_loaded or client_count < least_clients:
                least_clients = client_count
                least_loaded = [instance]
            elif client_count == least_clients:
                least_loaded.append(instance)
        if len(least_loaded) < len(self.staging_services[dbservice]):
            for instance in self.staging_services[dbservice]:
                if instance not in least_loaded:
                    self.logger.debug("Discounted {0:l} due to load."
                                      .format(instance))
            self.staging_services[dbservice] = least_loaded

    def choose_random(self, dbservice):
        """Pick a service instance randomly."""
        self.staging_services[dbservice] = [
            choice(self.staging_services[dbservice])]
        self.logger.debug("Randomly chose service %s instance %s "
                          "from remaining choices.",
                          dbservice.name,
                          self.staging_services[dbservice][0].name)

    def finalize_service_instances(self):
        """Fill out the list of chosen services."""
        for (service, instances) in self.staging_services.items():
            if len(instances) < 1:  # pragma: no cover
                self.error("Internal Error: Attempt to finalize on "
                           "service %s without any candidates." %
                           service.name)
                continue
            if len(instances) > 1:  # pragma: no cover
                self.error("Internal Error: Attempt to finalize on "
                           "service %s with too many candidates %s." %
                           (service.name,
                            ["service %s instance %s" %
                             (instance.service.name, instance.name)
                             for instance in instances]))
            self.chosen_services[service] = instances[0]

    def analyze_changes(self):
        """Determine what changed."""
        for (service, instance) in self.chosen_services.items():
            if service not in self.original_service_instances or \
               self.original_service_instances[service] != instance:
                self.instances_bound.add(instance)
        for (service, instance) in self.original_service_instances.items():
            if service not in self.chosen_services or \
               self.chosen_services[service] != instance:
                self.instances_unbound.add(instance)

    def stash_services(self):
        changed_servers = set()
        for instance in self.instances_bound.union(self.instances_unbound):
            if not instance.service.need_client_list:
                continue

            for srv in instance.servers:
                if srv.host:
                    changed_servers.add(srv.host)
                if srv.cluster:
                    changed_servers.update(srv.cluster.all_objects())

            self.plenaries.add(instance, cls=PlenaryServiceInstanceServer)

        for dbobj in changed_servers:
            # Skip servers that do not have a profile
            if not dbobj.archetype.is_compileable:
                continue

            # Skip servers that are in a different domain/sandbox
            if dbobj.branch != self.dbobj.branch or \
               dbobj.sandbox_author_id != self.dbobj.sandbox_author_id:
                continue

            # Ignore IncompleteError for servers of service instances. It is
            # debatable if this is the right thing to do, but it preserves the
            # status quo, and can be revisited later.
            self.plenaries.add(dbobj)

    def get_key(self):
        return self.plenaries.get_key()

    def prestash_primary(self, allow_incomplete=False):
        for dbobj in self.dbobj.all_objects():
            self.plenaries.add(dbobj, allow_incomplete=allow_incomplete)

            # This may be too much action at a distance... however, if
            # we are potentially re-writing a host plenary, it seems like
            # a good idea to also verify and refresh known dependencies.
            if dbobj.resholder:
                for dbres in dbobj.resholder.resources:
                    self.plenaries.add(dbres)
            if hasattr(dbobj, 'hardware_entity'):
                self.plenaries.add(dbobj.hardware_entity)


class HostChooser(Chooser):
    """Choose services for a host."""

    def __init__(self, dbhost, *args, **kwargs):
        """Provide initialization specific for host bindings."""
        if not isinstance(dbhost, Host):
            raise InternalError("HostChooser can only choose services for "
                                "hosts, got %r (%s)" % (dbhost, type(dbhost)))
        super(HostChooser, self).__init__(dbhost, *args, **kwargs)

        self.location = dbhost.hardware_entity.location
        if dbhost.cluster:
            self.aligned_parents.append(dbhost.cluster)
            if dbhost.cluster.metacluster:
                self.aligned_parents.append(dbhost.cluster.metacluster)

        # If the primary name is a ReservedName, then it does not have a network
        # attribute
        if hasattr(dbhost.hardware_entity.primary_name, 'network'):
            self.network = dbhost.hardware_entity.primary_name.network


class ClusterChooser(Chooser):
    """Choose services for a cluster."""

    def __init__(self, dbcluster, *args, **kwargs):
        """Provide initialization specific for cluster bindings."""
        if not isinstance(dbcluster, Cluster):
            raise InternalError("ClusterChooser can only choose services for "
                                "clusters, got %r (%s)" % (dbcluster, type(dbcluster)))
        super(ClusterChooser, self).__init__(dbcluster, *args, **kwargs)

        self.location = dbcluster.location_constraint
        if dbcluster.metacluster:
            self.aligned_parents.append(dbcluster.metacluster)

    def get_footprint(self, instance):
        """If this cluster is bound to a service, how many hosts bind?"""
        if instance.service in self.required_services:
            if isinstance(self.dbobj, MetaCluster):
                return 0
            # max_hosts can be None, but we need to return a number
            return self.dbobj.max_hosts or 0
        return 0
