#!/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 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


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 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.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, handler_dir):
        dict.__init__(self)
        self._handlers = []
        if handler_dir:
            self.HANDLER_DIR = handler_dir
        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.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 sorted(result, key=lambda x: x.prio)


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


nodes = get_nodes(get_file('NODE_LIST_FILE'))
resources = get_resources(get_file('RESOURCE_LIST_FILE'))

rcache = ResourceDescriptorCache(os.environ.get('HANDLER_DIR'))

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
    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(nodes, 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))
