#!/usr/bin/python
# vim: ts=4 sts=4 sw=4 et si
#
# Copyright (c) 2016-2017, Parallels International GmbH
#
# Our contact details: Parallels International GmbH, Vordergasse 59, 8200
# Schaffhausen, Switzerland.
#

from __future__ import print_function
import os
import sys
import subprocess
import functools


def error(msg):
    print(msg, file=sys.stderr)


def get_cpupools_output():
    '''Get the mapping of cluster nodes to CPU pools
       @return: dictionary {'node_id' : 'pool_name'}'''
    try:
        out = subprocess.check_output((os.environ.get('CPUPOOLS_BIN',
                                      '/usr/sbin/cpupools'), 'list-nodes'))
        error("'cpupools list-nodes' returned %r" % out)
    except subprocess.CalledProcessError:
        return {}
    d = {}
    for line in out.splitlines():
        # '<node_id> <pool_name>'
        x = line.strip().split(None, 1)
        if len(x) != 2:
            error("unexpected output from 'cpupools list-nodes': %r" % line)
            return {}
        d.update([x])
    return d


def call_once(f):
    @functools.wraps(f)
    def wrapper(*args, **kw):
        if wrapper.called:
            return wrapper.result
        wrapper.result = f(*args, **kw)
        wrapper.called = True
        return wrapper.result

    wrapper.called = False
    return wrapper


@call_once
def filter_by_cpupool(nodes, src_node):
    '''Filter nodes that are CPU-compatible with a given node
       @nodes: list of Nodes
       @src_node: cluster id of a node, used as a filter
       @return: list of Nodes that are CPU-compatible with a given node'''
    p = get_cpupools_output()
    src_pool = p.get(src_node)
    if src_pool is None:
        # resources are from node that is not included in any pool:
        # schedule resources without considering pools
        return nodes
    result = []
    for n in nodes:
        pool = p.get(n.cluster_id)
        # relocate resources inside the same pool and to the nodes that
        # are not included in any pool
        if pool is None or pool == src_pool:
            result.append(n)
    return result


def filter_by_role(nodes, role):
    '''Filter nodes that support given resource type (role)
       @nodes: list of Nodes
       @role: a string that represents a resource type (role)
       @return: list of Nodes that support given resource type (role)'''
    result = []
    for n in nodes:
        if role in n.roles or ('*' in n.roles and ('~' + role) not in n.roles):
            result.append(n)
    return result


class Node:
    '''A node in cluster'''
    def __init__(self, description):
        self.ip_addr, self.cluster_id = description[0:2]
        # list of resource types, supported by the node
        self.roles = []
        if len(description) > 2:
            self.roles = description[2].strip().split(',')
        # number of resources, relocated to the node, by type
        self._res_count = {}

    def get_count(self, res_type):
        # this is needed in order to choose a node with the smallest
        # amount of already relocated resources of a given type
        return self._res_count.get(res_type, 0)

    def inc(self, res_type):
        try:
            self._res_count[res_type] += 1
        except KeyError:
            self._res_count[res_type] = 1


class Resource:
    '''Cluster resource '''
    name = ''
    # relocation priority
    prio = ''
    # resource data path inside the cluster FS
    path = ''


class ResourceDescriptor:
    '''Description of a resource type'''
    type = ''
    # path to the get-property script
    script = ''
    # whether resource is sensitive to CPU features
    cpupools = False
    # whether a broken resource should be relocated to master node
    relocate_broken = False
    # whether resource type has subtypes (e.g. VM:QEMU, VM:VZ6)
    has_flavors = False


def get_property(script, name, env=None):
    '''Execute get-property script
       @script: path to 'get-property' script
       @name: property name
       @return: script stdout (string) on success'''
    e = os.environ.copy()
    if env:
        e.update(env)
    e.update(PROPERTY_NAME=name)
    try:
        return subprocess.check_output(script, env=e).strip()
    except subprocess.CalledProcessError as ex:
        error("'%s %s' failed with code %d" % (script, env, ex.returncode))
        return None


def parse_flag(value):
    return value.lower() == 'yes'


def get_flag(script, name):
    return parse_flag(get_property(script, name))


def get_all_props(script, descriptor):
    lines = get_property(script, 'all')
    if not lines:
        return False
    d = {}
    for l in lines.splitlines():
        x = l.strip().split('=')
        if len(x) != 2:
            error("invalid output for 'all' resource property: %s" % l)
            return False
        d.update([x])
    try:
        descriptor.cpupools = parse_flag(d['cpupools'])
        descriptor.relocate_broken = parse_flag(d['relocate-broken'])
        descriptor.type = d['type']
        descriptor.has_flavors = parse_flag(d['has-flavors'])
    except KeyError as e:
        error('%s key is not found in script output: %r' % (e, lines))
        return False
    return True


class ResourceDescriptorCache(dict):
    '''Caching facility for resource descriptions'''
    HANDLER_DIR = '/usr/share/shaman'

    def __init__(self):
        dict.__init__(self)
        self.handlers = []
        for p, d, f in os.walk(self.HANDLER_DIR):
            self._handlers = d
            break

    def __getitem__(self, key):
        # find all possible handlers for a given resource:
        # usually, they have the form 'xxx-' e.g. 'vm-', 'ct-'
        h = [x for x in self._handlers if key.startswith(x)]
        if len(h) != 1:
            error("couldn't find handler for %r: candidates are %r"
                  % (key, h))
            return ResourceDescriptor()
        return dict.__getitem__(self, h[0])

    def __missing__(self, key):
        d = ResourceDescriptor()
        script = os.path.join(self.HANDLER_DIR, key, 'get-property')
        if not os.path.exists(script):
            return d
        d.script = script
        # first check whether resource's get-property script supports
        # retrieving all properties at once -- otherwise, fallback to
        # calling the script several times
        if not get_all_props(script, d):
            d.cpupools = get_flag(script, 'cpupools')
            d.relocate_broken = get_flag(script, 'relocate-broken')
            d.type = get_property(script, 'type')
            d.has_flavors = get_flag(script, 'has-flavors')
        # store the newly obtained descriptor in the cache
        dict.__setitem__(self, key, d)
        return d


def get_env(name):
    '''Get the value of a given environment variable'''
    val = os.environ.get(name)
    if not val:
        error('%r environment variable undefined' % name)
        sys.exit(1)
    return val


def get_file(name):
    '''Get file name from a given environment variable'''
    f = get_env(name)
    if not os.path.exists(f):
        error('file %r does not exist' % f)
        sys.exit(1)
    return f


def get_nodes(path):
    '''Obtain list of Nodes from a given file'''
    result = []
    with open(path, 'r') as f:
        for line in f:
            # "<node_ip> <node_cluster_id> [role1,role2,...,roleN]"
            node = line.strip().split()
            if len(node) < 2 or len(node) > 3:
                error('invalid node description: %r' % line)
                sys.exit(1)
            result.append(Node(node))
    return result


def get_resources(path):
    '''Obtain list of Resources from a given file'''
    result = []
    with open(path, 'r') as f:
        count = 0
        # "<name>\n<prio>\n<path possibly with spaces>\n"
        for line in f:
            s = line.strip()
            if not s:
                continue
            if count % 3 == 0:
                res = Resource()
                res.name = s
            elif count % 3 == 1:
                res.prio = s
            else:
                res.path = s
                result.append(res)
            count += 1
        if count % 3 != 0:
            error('invalid contents of resource description file')
            sys.exit(1)
    return result


def print_broken(name, relocate):
    print('broken {} {}'.format('yes' if relocate else 'no',
          name), file=out_file)


src_node_id = get_env('SRC_NODE_ID')
nodes = get_nodes(get_file('NODE_LIST_FILE'))
resources = get_resources(get_file('RESOURCE_LIST_FILE'))
try:
    out_file = os.fdopen(int(get_env('OUTPUT_FD')), 'w')
except ValueError as e:
    error('invalid value for the output file descriptor: %r' % e)
    sys.exit(1)

rcache = ResourceDescriptorCache()

for res in resources:
    d = rcache[res.name]
    if not d.type:
        error('empty resource type for %r' % res.name)
        print_broken(res.name, d.relocate_broken)
        continue
    tmp = filter_by_cpupool(nodes, src_node_id) if d.cpupools else nodes
    if not tmp:
        error('node list is empty after filtering by cpupools')
        print_broken(res.name, d.relocate_broken)
        continue
    role = d.type
    if d.has_flavors:
        role = get_property(d.script, 'flavor', {'RESOURCE_PATH': res.path})
        if not role:
            error('cannot get the flavor of resource %r' % res.name)
            print_broken(res.name, d.relocate_broken)
            continue
    tmp = filter_by_role(tmp, role)
    if not tmp:
        error('node list is empty after filtering by node roles')
        print_broken(res.name, d.relocate_broken)
        continue
    # choose a node with the minimal number of already relocated
    # resources of a given resource type
    node = tmp[0]
    for n in tmp[1:len(tmp)]:
        if n.get_count(d.type) < node.get_count(d.type):
            node = n
    node.inc(d.type)
    print('%s %s %s' % (node.ip_addr, node.cluster_id, res.name),
          file=out_file)
