#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Hostlist utility
#
# Version 1.1
#
# Copyright (C) 2008 Kent Engström <kent@nsc.liu.se> and
#                    Thomas Bellman <bellman@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.

import sys
import os
import optparse
import operator

from hostlist import expand_hostlist, collect_hostlist, numerically_sorted, BadHostlist

def func_union(args):
    return reduce(operator.or_, args)

def func_intersection(args):
    return reduce(operator.and_, args)

def func_difference(args):
    return reduce(operator.sub, args)

def func_xor(args):
    return reduce(operator.xor, args)

op = optparse.OptionParser(usage="usage: %prog [options] {hostlist arguments}")
op.add_option("-u", "--union",
              action="store_const", dest="func", const=func_union,
              default=func_union,
              help="compute the union of the hostlist arguments (default)")
op.add_option("-i", "--intersection",
              action="store_const", dest="func", const=func_intersection,
              help="compute the intersection of the hostlist arguments")
op.add_option("-d", "--difference",
              action="store_const", dest="func", const=func_difference,
              help="compute the difference between the first hostlist argument and the rest")
op.add_option("-x", "--symmetric-difference",
              action="store_const", dest="func", const=func_xor,
              help="compute the symmetric difference between the first hostlist argument and the rest")
op.add_option("-w", "--expand",
              action="store_true",
              help="output the results as an expanded list")
op.add_option("-c", "--collapse",
              action="store_false", dest="expand",
              help="output the results as a hostlist expression (default)")
op.add_option("-n", "--count",
              action="store_true",
              help="output the number of hosts instead of a hostlist")
(opts, args) = op.parse_args()

func = opts.func

func_args  = []

try:
    for a in args:
        if a == "-":
            for a in sys.stdin.read().split():
                func_args.append(set(expand_hostlist(a)))
        else:
            func_args.append(set(expand_hostlist(a)))
except BadHostlist, e:
    sys.stderr.write("Bad hostlist ``%s'' encountered: %s\n"
                     % ((a,) + e.args))
    sys.exit(os.EX_DATAERR)

if not func_args:
    op.print_help()
    sys.exit(os.EX_USAGE)

res = func(func_args)

if opts.count:
    print len(res)
elif opts.expand:
    for host in numerically_sorted(res):
        print host
else:
    try:
        print collect_hostlist(res)
    except BadHostlist, e:
        sys.stderr.write("Bad hostname encountered: %s\n" % e.args)
        sys.exit(os.EX_DATAERR)
