#!/usr/bin/python3

DEFAULT_COLLECTOR = 'collector.opensciencegrid.org:9619'
DEFAULT_COLLECTOR_PORT = '9619'

import os
import sys
import optparse

os.environ.setdefault("CONDOR_CONFIG", "/etc/condor-ce/condor_config")

import classad2 as classad
import htcondor2 as htcondor

import htcondorce.info_query


def formatResourceAdsTable(resources, width):
    """Create a table of resource information obtained from `resources`,
    formatted in a text block `width` characters wide. Returns the table
    as a string

    """
    if width < 50:
        width = 50
    cpus_width = 4
    memory_width = 8
    max_wall_time_width = 12

    remaining_width = width - cpus_width - memory_width - max_wall_time_width
    name_width = remaining_width * 2 // 5
    allowed_vos_width = remaining_width - name_width

    format_str = ("%%-%(name_width)d.%(name_width)ds "
                  "%%%(cpus_width)ds "
                  "%%%(memory_width)ds "
                  "%%%(max_wall_time_width)ds "
                  "%%-.%(allowed_vos_width)ds" % locals())
    table = format_str % ("Name",
                          "CPUs",
                          "Memory",
                          "MaxWallTime",
                          "AllowedVOs")
    table += "\n"
    for resource in resources:
        catalog_entry = resource.catalog_entry
        table += "\n" + format_str % (catalog_entry.get('Name','???'),
                                      catalog_entry.get('CPUs', catalog_entry.get('Cpus', '???')),
                                      catalog_entry.get('Memory','???'),
                                      catalog_entry.get('MaxWallTime', ''),
                                      ", ".join(list(catalog_entry.get('AllowedVOs', []))))

    return table


def parseOpts(argv):
    """Parse command-line options and arguments.
    Returns the options and the address of the collector to contact

    """
    usage = "%prog [options] [--pool=<collector addr>[:<port>]]"
    epilog = '''

Here are some example queries.
To list resources that provide at least 1 GB of memory and are accessible
by the 'osg' VO:

    %(prog)s  --memory 1024  --vo osg

which results in:

    Name                   CPUs   Memory  MaxWallTime AllowedVOs

    matyasfermicloud085       1     1877         1440 fermi, osg

To list resources that have a MaxWallTime of at least 1 day (1440 minutes):

    %(prog)s  --walltime 1440

To print a full classad containing resource attributes and some attributes
from the CE containing the resource:

    %(prog)s  --walltime 1440  --verbose

which results in:

        [
            OSG_BatchSystems = "Condor";
            Name = "matyasfermicloud085";
            MaxWallTime = 1440;
            CPUs = 1;
            Memory = 1877;
            ...
        ]

''' % {'prog': os.path.basename(argv[0])}

    parser = optparse.OptionParser(usage=usage, epilog=epilog)
    # Don't reformat the epilog string (https://stackoverflow.com/questions/1857346)
    optparse.OptionParser.format_epilog = lambda self, formatter: self.epilog

    parser.add_option("--constrain", default=None,
                      help="Show only results matching a ClassAd expression")
    parser.add_option("--cpus", default=1, type='int',
                      help="CPUs requested")
    parser.add_option("--memory", default=1, type='int',
                      help="Memory requested (in megabytes)")
    parser.add_option("--name", default=None,
                      help="Name of the resource")
    parser.add_option("--walltime", default=0, type='int',
                      help="Walltime requested (in minutes)")
    parser.add_option("--vo", default=None,
                      help="VO membership")
    parser.add_option("--width", default=os.environ.get('COLUMNS', 80), type='int',
                      help="Output width for table")
    parser.add_option("--pool", default=None, metavar='ADDRESS',
                      help="Collector to get information from")
    parser.add_option("--submitfile", default=False, dest='submitfile', action='store_true',
                      help="Show entries as lines to be added to a submit file")
    parser.add_option("-v", "--verbose", "-l", "--long", dest='verbose', default=False, action='store_true',
                      help="Show matching classads instead of a table")
    parser.add_option("-d", "--debug", default=False, action='store_true',
                      help="Show debug output")

    opts = parser.parse_args(argv[1:])[0]

    if opts.debug:
        htcondor.enable_debug()
    global g_debug
    g_debug = opts.debug

    if opts.constrain:
        try:
            # catch this one early
            classad.ExprTree(opts.constrain)
        except SyntaxError:
            parser.error("Invalid classad expression: %s" % opts.constrain)

    if opts.pool:
        collector = opts.pool
    else:
        collector = htcondor.param.get('CONDOR_CE_COLLECTOR_HOST', DEFAULT_COLLECTOR)
    if ':' not in collector:
        collector += ":" + DEFAULT_COLLECTOR_PORT

    return opts, collector


def main(argv):
    """Main function.

    Get a list of ResourceAd objects filled with data from querying a
    collector; pass the list through a series of filters and either
    display the remaining resource ads in full, or create a table
    from the resource ads and display that.

    """

    opts, collector_addr = parseOpts(argv)

    constraints = {}
    for key in ['constrain', 'cpus', 'memory', 'name', 'vo', 'walltime']:
        try:
            constraints[key] = getattr(opts, key)
        except AttributeError:
            pass

    ce_ads = htcondorce.info_query.fetchCEAds(collector_addr)
    resources = htcondorce.info_query.filterResourceAds(constraints, htcondorce.info_query.getResourceAdsIter(ce_ads))

    if opts.verbose:
        for res in resources:
            print(res)
    elif opts.submitfile:
        resources = [x for x in resources if x.get('Name')]
        if not resources:
            print("No resources match", file=sys.stderr)
            return 1
        else:
            if len(resources) > 1:
                print("%d resources match; only printing first" % len(resources), file=sys.stderr)
            sorted_resources = sorted(resources, key=lambda x: x['Name'])
            res = sorted_resources[0]
            print("# Submit file for resource " + res['Name'])
            print("\n".join(htcondorce.info_query.getSubmitFileAdditions(res)) + "\n")
    else:
        print(formatResourceAdsTable(resources, opts.width))



if __name__ == '__main__':
    sys.exit(main(sys.argv))
