#!/usr/bin/env python
# encoding: utf-8
"""
clusterinfo.py

Created by Nitin Madnani on 2006-07-09.
Copyright (c) 2006 UMIACS. All rights reserved.

"""

import sys, os, re, xml.parsers.expat
from optparse import OptionParser
from xml.dom.minidom import parseString

class job:
    def __init__(self, args):
        self.jobid = args[0].split('.')[0]
        self.user = args[1]
        self.queue = args[2]
        self.name = args[3]
        self.status = args[9]
        self.queue = args[10]

def parse_qstat_output(qstat_output):
    # Get relevant portion of the qstat output
    qstat_lines = qstat_output.split('\n')
    qstat_lines = qstat_lines[5:-1]
    jobs = {}
    
    # Iterate over the lines and create a job() object for each line
    for ql in qstat_lines:
        ql_fields = re.split('\s+',ql.strip())
        jobnode = ql_fields[-1].split('/')[0]            
        if jobs.has_key(jobnode):
            jobs[jobnode].append(job(ql_fields[:-1]))
        else:
            jobs[jobnode] = [job(ql_fields[:-1])]
        
    return jobs        

# This function formats the output to be printed on the screen
def format_output(dom, qstatout, verbose=False):
    free = []
    exclusive = []
    partial = []
    offline = []
    if verbose:
        nodenames = []
    
    # Extract all <node> elements from the dom
    nodes = dom.getElementsByTagName('Node')
    
    # Iterate over each <node> element and get the relevant info
    for node in nodes:
        # Get node name
        propertyNode = node.getElementsByTagName('properties')[0]
        (foo, bar, name) = propertyNode.firstChild.nodeValue.split(',')

        if verbose:
            nodenames.append(name)

        # Get node state
        state = node.getElementsByTagName('state')[0].firstChild.nodeValue.lower()
        
        if state == 'job-exclusive':
            exclusive.append(name)                                
        elif state == 'offline':
            offline.append(name)
            continue
        
        # Get jobs running, if any, on this node
        jobsNode = node.getElementsByTagName('jobs')
        if jobsNode:
            jobs = jobsNode[0].firstChild.nodeValue.split(',')
            if 1 <= len(jobs) < 4:
                partial.append((name, len(jobs)))
        else:
            free.append(name)

    # Calculate the total number of jobs running
    runningjobs = sum([x[1] for x in partial]) + 4*len(exclusive)
    
    # Write out information to stdout
    outputstr = '\nTotal Nodes: %d, Online: %d, Offline: %d\n\n' % (len(nodes), len(nodes) - len(offline), len(offline))
    outputstr += 'Total jobs running: %d\n' % runningjobs
    outputstr += '\nSummary:\n'
    outputstr += '  Free [%d] : %s\n' % (len(free), ', '.join(free))
    outputstr += '  Job-Exclusive [%d] : %s\n' % (len(exclusive), ', '.join(exclusive))
    outputstr += '  Partial [%d] : %s\n\n' % (len(partial), ', '.join([str(x[0]) + '(' + str(x[1]) + ')' for x in partial]))

    # If the user requested verbose output, then process the qstat_output as well
    if verbose:
        outputstr += 'Breakdown by nodes:\n'
        jobhash = parse_qstat_output(qstatout)
        
        # Go over each node and find which jobs, if any are running on it
        for nodename in nodenames:
            outputstr += '  %s: \n' % nodename
            if nodename in offline:
                outputstr += '      Offline.\n'
                continue
            elif nodename in free:
                outputstr += '      No jobs running.\n'
            elif jobhash.has_key(nodename):
                nodejobs = jobhash[nodename]
                userhash = {}
                for nodejob in nodejobs:
                    if userhash.has_key(nodejob.user):
                        userhash[nodejob.user].append(nodejob.jobid)
                    else:
                        userhash[nodejob.user] = [nodejob.jobid]
                    
                for user in userhash.keys():
                    outputstr += '      %s[%d] : %s\n' % (user, len(userhash[user]), ' '.join(userhash[user]))
        
        
        # Print out jobs on hold or queued, if any
        heldstr = ''
        queuestr = ''
        held_userhash = {}
        queue_userhash = {}
        if jobhash.has_key('--'):
            heldstr += '\nJobs on hold: \n'
            queuestr += '\nJobs queued: \n'
            for thisjob in jobhash['--']:
                if thisjob.status == 'Q':
                    if queue_userhash.has_key(thisjob.user):
                        queue_userhash[thisjob.user].append(thisjob.jobid)
                    else:
                        queue_userhash[thisjob.user] = [thisjob.jobid]
                elif thisjob.status == 'H':
                    if held_userhash.has_key(thisjob.user):
                        held_userhash[thisjob.user].append(thisjob.jobid)
                    else:
                        held_userhash[thisjob.user] = [thisjob.jobid]                    
            for queueuser in queue_userhash:
                queuestr += '  %s[%d] : %s\n' % (queueuser, len(queue_userhash[queueuser]), ' '.join(queue_userhash[queueuser]))
            for helduser in held_userhash:
                heldstr += '  %s[%d] : %s\n' % (helduser, len(held_userhash[helduser]), ' '.join(held_userhash[helduser]))
                
        # Append the queuestr and heldstr to outputstr
        if held_userhash:
            outputstr += heldstr
        if queue_userhash:
            outputstr += queuestr                         
    return outputstr

#def main(argv=None):

parser = OptionParser()
parser.add_option("-v", "--verbose", action="store_true", dest="verbose",default=False, help="show more detailed information")
(options, args) = parser.parse_args()

# Get the user name
username = os.getenv('USER')

# Get server name
submitnode = os.getenv('SUBMITNODE')
if not submitnode:
    print >> sys.stderr, "Error: Please define the SUBMITNODE environment variable."

# Get the cluster usage information from the submit node
try:
    pbs_cmd = 'ssh %s@%s pbsnodes -x' % (username, submitnode)
    pipe = os.popen(pbs_cmd, 'r')
    cluster_xml = pipe.read()
    pbs_status = pipe.close()
    status_dom = parseString(cluster_xml)
    if pbs_status:
        print >> sys.stderr, "Error: Could not execute ``pbsnodes'' on cluster submit node."
        sys.exit(1)
except xml.parsers.expat.ExpatError:
    print >> sys.stderr, "Error: Malformed XML from ``pbsnodes''."
    sys.exit(1)

# If the user requested verbose information, then we need to do a qstat as well
if options.verbose:
    qstat_cmd = 'ssh %s@csubmit01.umiacs.umd.edu qstat -n -1' % username
    pipe = os.popen(qstat_cmd, 'r')
    qstat_output = pipe.read()
    qstat_status = pipe.close() 
    if qstat_status:
        print >> sys.stderr, "Error: Could not execute ``qstat'' on cluster submit node."
        sys.exit(1)
                
# print the formatted information to the screen
if options.verbose:
    print format_output(status_dom, qstat_output, options.verbose)
else:
    print format_output(status_dom, None, options.verbose)    
                
#if _name__ == "__main__":
#    main()
