#!/usr/bin/env python
#
# Hostlist library
#
# Copyright (C) 2008-2018
#                    Kent Engström <kent@nsc.liu.se>,
#                    Thomas Bellman <bellman@nsc.liu.se>,
#                    Pär Lindfors <paran@nsc.liu.se> and
#                    Torbjörn Lönnemark <ketl@nsc.liu.se>,
#                    National Supercomputer Centre
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

"""Handle hostlist expressions.

This module provides operations to expand and collect hostlist
expressions.

The hostlist expression syntax is the same as in several programs
developed at LLNL (https://computing.llnl.gov/linux/). However in
corner cases the behaviour of this module have not been compared for
compatibility with pdsh/dshbak/SLURM et al.
"""

__version__ = "2.2.1"

import re
import itertools

# Replace range with xrange on Python 2, do nothing on Python 3 (where xrange
# does not exist, and range returns an iterator)
try:
    range = xrange
except:
    pass

# Exception used for error reporting to the caller
class BadHostlist(Exception): pass

# Configuration to guard against ridiculously long expanded lists
MAX_SIZE = 100000

# Hostlist expansion

def expand_hostlist(hostlist, allow_duplicates=False, sort=False):
    """Expand a hostlist expression string to a Python list.

    Example: expand_hostlist("n[9-11],d[01-02]") ==> 
             ['n9', 'n10', 'n11', 'd01', 'd02']

    Unless allow_duplicates is true, duplicates will be purged
    from the results. If sort is true, the output will be sorted.
    """

    results = []
    bracket_level = 0
    part = ""
    
    for c in hostlist + ",":
        if c == "," and bracket_level == 0:
            # Comma at top level, split!
            if part: results.extend(expand_part(part))
            part = ""
            bad_part = False
        else:
            part += c

        if c == "[": bracket_level += 1
        elif c == "]": bracket_level -= 1

        if bracket_level > 1:
            raise BadHostlist("nested brackets")
        elif bracket_level < 0:
            raise BadHostlist("unbalanced brackets")

    if bracket_level > 0:
        raise BadHostlist("unbalanced brackets")

    if not allow_duplicates:
        results = remove_duplicates(results)
    if sort:
        results = numerically_sorted(results)
    return results

def expand_part(s):
    """Expand a part (e.g. "x[1-2]y[1-3][1-3]") (no outer level commas)."""

    # Base case: the empty part expand to the singleton list of ""
    if s == "":
        return [""]

    # Split into:
    # 1) prefix string (may be empty)
    # 2) rangelist in brackets (may be missing)
    # 3) the rest

    m = re.match(r'([^,\[]*)(\[[^\]]*\])?(.*)', s)
    (prefix, rangelist, rest) = m.group(1,2,3)

    # Expand the rest first (here is where we recurse!)
    rest_expanded = expand_part(rest)

    # Expand our own part
    if not rangelist:
        # If there is no rangelist, our own contribution is the prefix only
        us_expanded = [prefix]
    else:
        # Otherwise expand the rangelist (adding the prefix before)
        us_expanded = expand_rangelist(prefix, rangelist[1:-1])

    # Combine our list with the list from the expansion of the rest
    # (but guard against too large results first)
    if len(us_expanded) * len(rest_expanded) > MAX_SIZE:
        raise BadHostlist("results too large")

    return [us_part + rest_part
            for us_part in us_expanded
            for rest_part in rest_expanded]

def expand_rangelist(prefix, rangelist):
    """ Expand a rangelist (e.g. "1-10,14"), putting a prefix before."""
    
    # Split at commas and expand each range separately
    results = []
    for range_ in rangelist.split(","):
        results.extend(expand_range(prefix, range_))
    return results

def expand_range(prefix, range_):
    """ Expand a range (e.g. 1-10 or 14), putting a prefix before."""

    # Check for a single number first
    m = re.match(r'^[0-9]+$', range_)
    if m:
        return ["%s%s" % (prefix, range_)]

    # Otherwise split low-high
    m = re.match(r'^([0-9]+)-([0-9]+)$', range_)
    if not m:
        raise BadHostlist("bad range")

    (s_low, s_high) = m.group(1,2)
    low = int(s_low)
    high = int(s_high)
    width = len(s_low)

    if high < low:
        raise BadHostlist("start > stop")
    elif high - low > MAX_SIZE:
        raise BadHostlist("range too large")

    results = []
    for i in range(low, high+1):
        results.append("%s%0*d" % (prefix, width, i))
    return results

def remove_duplicates(l):
    """Remove duplicates from a list (but keep the order)."""
    seen = set()
    results = []
    for e in l:
        if e not in seen:
            results.append(e)
            seen.add(e)
    return results

# Hostlist collection

def collect_hostlist(hosts, silently_discard_bad = False):
    """Collect a hostlist string from a Python list of hosts.

    We start grouping from the rightmost numerical part.
    Duplicates are removed.

    A bad hostname raises an exception (unless silently_discard_bad
    is true causing the bad hostname to be silently discarded instead).
    """

    # Split hostlist into a list of (host, "") for the iterative part.
    # (Also check for bad node names now)
    # The idea is to move already collected numerical parts from the
    # left side (seen by each loop) to the right side (just copied).

    left_right = []
    for host in hosts:
        # We remove leading and trailing whitespace first, and skip empty lines
        host = host.strip()
        if host == "": continue

        # We cannot accept a host containing any of the three special
        # characters in the hostlist syntax (comma and flat brackets)
        if re.search(r'[][,]', host):
            if silently_discard_bad:
                continue
            else:
                raise BadHostlist("forbidden character")

        left_right.append((host, ""))

    # Call the iterative function until it says it's done
    looping = True
    while looping:
        left_right, looping = collect_hostlist_1(left_right)
    return ",".join([left + right for left, right in left_right])

def collect_hostlist_1(left_right):
    """Collect a hostlist string from a list of hosts (left+right).

    The input is a list of tuples (left, right). The left part
    is analyzed, while the right part is just passed along
    (it can contain already collected range expressions).
    """

    # Scan the list of hosts (left+right) and build two things:
    # *) a set of all hosts seen (used later)
    # *) a list where each host entry is preprocessed for correct sorting

    sortlist = []
    remaining = set()
    for left, right in left_right:
        host = left + right
        remaining.add(host)

        # Match the left part into parts
        m = re.match(r'^(.*?)([0-9]+)?([^0-9]*)$', left)
        (prefix, num_str, suffix) = m.group(1,2,3)

        # Add the right part unprocessed to the suffix.
        # This ensures than an already computed range expression
        # in the right part is not analyzed again.
        suffix = suffix + right 

        if num_str is None:
            # A left part with no numeric part at all gets special treatment!
            # The regexp matches with the whole string as the suffix,
            # with nothing in the prefix or numeric parts.
            # We do not want that, so we move it to the prefix and put
            # None as a special marker where the suffix should be.
            assert prefix == ""
            sortlist.append(((host, None), None, None, host))
        else:
            # A left part with at least an numeric part
            # (we care about the rightmost numeric part)
            num_int = int(num_str)
            num_width = len(num_str) # This width includes leading zeroes
            sortlist.append(((prefix, suffix), num_int, num_width, host))

    # Sort lexicographically, first on prefix, then on suffix, then on
    # num_int (numerically), then...
    # This determines the order of the final result.

    def sort_key(entry):
        """
        Sort sortlist entries without causing TypeError in Python 3.

        Prefix is always present and a string.

        If suffix is None, use the empty string. The empty string will sort
        before all other strings (as None did in Python 2).

        If num_int is None, use -1 instead. -1 sorts before all *possible*
        numbers (as None did in Python 2). Negative numbers cannot be present
        in a hostlist.

        If num_width is None, use -1 instead. As before, -1 sorts before all
        possible widths. An actual num_width of 0 should not be possible, so we
        could technically use 0, but using -1 doesn't hurt, so we might as well
        use that instead.
        """
        ((prefix, suffix), num_int, num_width, host) = entry
        return (
            (prefix, '' if suffix is None else suffix),
            -1 if num_int is None else num_int,
            -1 if num_width is None else num_width, host
        )


    sortlist.sort(key=sort_key)

    # We are ready to collect the result parts as a list of new (left,
    # right) tuples.

    results = []
    needs_another_loop = False 

    # Now group entries with the same prefix+suffix combination (the
    # key is the first element in the sortlist) to loop over them and
    # then to loop over the list of hosts sharing the same
    # prefix+suffix combination.

    for ((prefix, suffix), group) in itertools.groupby(sortlist,
                                                       key=lambda x:x[0]):

        if suffix is None:
            # Special case: a host with no numeric part
            results.append(("", prefix)) # Move everything to the right part
            remaining.remove(prefix)
        else:
            # The general case. We prepare to collect a list of
            # ranges expressed as (low, high, width) for later
            # formatting.
            range_list = []
    
            for ((prefix2, suffix2), num_int, num_width, host) in group:
                if host not in remaining:
                    # Below, we will loop internally to enumate a whole range
                    # at a time. We then remove the covered hosts from the set.
                    # Therefore, skip the host here if it is gone from the set.
                    continue
                assert num_int is not None

                # Scan for a range starting at the current host
                low = num_int
                while True:
                    host = "%s%0*d%s" % (prefix, num_width, num_int, suffix)
                    if host in remaining:
                        remaining.remove(host)
                        num_int += 1
                    else:
                        break
                high = num_int - 1
                assert high >= low
                range_list.append((low, high, num_width))

            # We have a list of ranges to format. We make sure
            # we move our handled numerical part to the right to
            # stop it from being processed again.
            needs_another_loop = True
            if len(range_list) == 1 and range_list[0][0] == range_list[0][1]:
                # Special case to make sure that n1 is not shown as n[1] etc
                results.append((prefix,
                                "%0*d%s" % 
                               (range_list[0][2], range_list[0][0], suffix)))
            else:
                # General case where high > low
                results.append((prefix, "[" + \
                                   ",".join([format_range(l, h, w)
                                             for l, h, w in range_list]) + \
                                   "]" + suffix))

    # At this point, the set of remaining hosts should be empty and we
    # are ready to return the result, together with the flag that says
    # if we need to loop again (we do if we have added something to a
    # left part).
    assert not remaining
    return results, needs_another_loop

def format_range(low, high, width):
    """Format a range from low to high inclusively, with a certain width."""

    if low == high:
        return "%0*d" % (width, low)
    else:
        return "%0*d-%0*d" % (width, low, width, high)

# Sort a list of hosts numerically

def numerically_sorted(l):
    """Sort a list of hosts numerically.

    E.g. sorted order should be n1, n2, n10; not n1, n10, n2.
    """

    return sorted(l, key=numeric_sort_key)

numeric_sort_key_regexp = re.compile("([0-9]+)|([^0-9]+)")
def numeric_sort_key(x):
    """Compose a sorting key to compare strings "numerically":

    We split numerical (integer) and non-numerical parts into a list,
    making sure that the numerical parts are converted to Python ints,
    and then sort on the lists. Thus, if we sort x10y and x9z8, we will
    compare ["x", 10, "y"] with ["x", 9, "x", "8"] and return x9z8
    before x10y".

    Python 3 complication: We cannot compare int and str, so while we can
    compare x10y and x9z8, we cannot compare x10y and 9z8. Kludge: insert
    a blank string first if the list would otherwise start with an integer.
    This will give the same ordering as before, as integers seem to compare
    smaller than strings in Python 2.
    """

    keylist = [int(i_ni[0]) if i_ni[0] else i_ni[1]
               for i_ni in numeric_sort_key_regexp.findall(x)]
    if keylist and isinstance(keylist[0], int):
        keylist.insert(0, "")
    return keylist

# Parse SLURM_TASKS_PER_NODE into a list of task numbers
#
# Description from the SLURM sbatch man page:
#              Number of tasks to be initiated on each node. Values
#              are comma separated and in the same order as
#              SLURM_NODELIST.  If two or more consecutive nodes are
#              to have the same task count, that count is followed by
#              "(x#)" where "#" is the repetition count. For example,
#              "SLURM_TASKS_PER_NODE=2(x3),1" indicates that the first
#              three nodes will each execute three tasks and the
#              fourth node will execute one task.

def parse_slurm_tasks_per_node(s):
    res = []
    for part in s.split(","):
        m = re.match(r'^([0-9]+)(\(x([0-9]+)\))?$', part)
        if m:
            tasks = int(m.group(1))
            repetitions = m.group(3)
            if repetitions is None:
                repetitions = 1
            else:
                repetitions = int(repetitions)
            if repetitions > MAX_SIZE:
                raise BadHostlist("task list repetitions too large")
            for i in range(repetitions):
                res.append(tasks)
        else:
            raise BadHostlist("bad task list syntax")
    return res

#
# Keep this part to tell users where the command line interface went
#

if __name__ == '__main__':
    import os, sys
    sys.stderr.write("The command line utility has been moved to a separate 'hostlist' program.\n")
    sys.exit(os.EX_USAGE)
