#!/usr/bin/python

import os
import sys
import json
import tempfile
import time
import optparse
import collections
import multiprocessing

import rrdtool

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

# htcondor and classad modules are imported within the functions that need them 
# to avoid loading HTCondor code in the parent then forking. Otherwise, issues
# arise in the child processes.

from htcondorce import rrd, web_utils

def parse_opts():
    parser = optparse.OptionParser()
    parser.add_option("-s", "--spool", help="Spool directory", dest="spool")

    return parser.parse_args()[0]

def secure_json_write(fname, contents):
    temp_fd, temp_fname = tempfile.mkstemp(dir=os.path.dirname(fname))
    temp_file = os.fdopen(temp_fd, 'w')
    try:
        try:
            json.dump(contents, temp_file)
            os.rename(temp_fname, fname)
        except:
            print "Error writing json file"
    finally:
        temp_file.close()

def process_one_schedd(ad):
    htcondor = __import__("htcondor")
    classad = __import__("classad")
    ad = classad.ClassAd(ad)
    schedd = htcondor.Schedd(ad)
    total_results = {"Running": 0, "Idle": 0, "Held": 0, "Jobs": 0}
    vo_results = {}
    job_count = {}
    print "Processing CE %s." % ad.get("Name", "Unknown")
    try:
        for job in schedd.xquery("true", ["JobStatus", 'x509UserProxyVOName',
                                          'x509UserProxyFirstFQAN', 'x509userproxysubject']):
            jobad = {}
            jobad['dn'] = job.get("x509userproxysubject", 'Unknown')
            jobad['vo'] = job.get('x509UserProxyVOName', 'Unknown')
            jobad['voms'] = job.get('x509UserProxyFirstFQAN', '').replace("/Capability=NULL", "").replace("/Role=NULL", "")
            for attr, val in jobad.iteritems():
                if not val:
                    jobad[attr] = 'Unknown'
            job_key = tuple(jobad.values())
            if job_key not in job_count:
                job_count[job_key] = {"Running": 0, "Idle": 0, "Held": 0, "Jobs": 0}
            if jobad['vo'] not in vo_results:
                vo_results[jobad['vo']] = {"Running": 0, "Idle": 0, "Held": 0, "Jobs": 0}
            results = vo_results[jobad['vo']]
            job_results = job_count[job_key]
            results["Jobs"] += 1
            total_results['Jobs'] += 1
            job_results['Jobs'] += 1
            if job.get("JobStatus") == 1:
                results['Idle'] += 1
                job_results['Idle'] += 1
                total_results['Idle'] += 1
            elif job.get("JobStatus") == 2:
                results['Running'] += 1
                job_results['Running'] += 1
                total_results['Running'] += 1
            elif job.get("JobStatus") == 5:
                results['Held'] += 1
                job_results['Held'] += 1
                total_results['Held'] += 1
    except RuntimeError, re:
        print "Failure to query CE %s: %s" % (ad.get("Name", "Unknown"), str(re))
    return [total_results, vo_results, job_count, ad]


def get_ads(spooldir):
    if not spooldir:
        spooldir = web_utils.get_spooldir()
    environ = {'htcondorce.spool': spooldir}
    ads = [str(ad) for ad in web_utils.get_schedd_ads({})]
    return [ads, environ]


def main():
    opts = parse_opts()

    htcondor = __import__("htcondor")
    # Check if the CEVIEW daemon is in the daemon list, exit if it is not
    if htcondor.param["DAEMON_LIST"].find("CEVIEW") == -1:
        print "CEVIEW is not in DAEMON_LIST, exiting..."
        sys.exit(0)

    pool = multiprocessing.Pool(10)
    ads, environ = pool.map(get_ads, [opts.spool])[0]

    total_results = collections.defaultdict(int, {"Running": 0, "Idle": 0, "Held": 0})
    vo_results = {}
    job_results = {}

    map_results = pool.map(process_one_schedd, ads, 1)

    for schedd_total_results, schedd_vo_results, schedd_job_results, ad in map_results:

        for key, val in schedd_total_results.items():
            total_results[key] += val

        for vo, schedd_results in schedd_vo_results.items():
            results = vo_results.setdefault(vo, collections.defaultdict(int))
            for key, val in schedd_results.items():
                results[key] += val

        for key, val in schedd_job_results.items():
            results = job_results.setdefault(key, collections.defaultdict(int))
            for key2, val2 in val.items():
                results[key2] += val2

        total_fname = rrd.check_rrd(environ, ad['Name'], "jobs")
        if os.isatty(2):
            print >> sys.stderr, "Schedd", ad['Name'], "totals:", schedd_total_results
        rrdtool.update(total_fname, "N:%d:%d:%d" % (schedd_total_results["Running"], schedd_total_results["Idle"],
                                                    schedd_total_results["Held"]))

        for vo, info in schedd_vo_results.items():
            vo_fname = rrd.check_rrd(environ, ad['Name'], "vos", vo)
            if os.isatty(2):
                print >> sys.stderr, "Schedd", ad['Name'], "VO", vo, "totals:", info
            rrdtool.update(vo_fname, "N:%d:%d:%d:%d" % (info['Running'], info['Idle'], info['Held'], info['Jobs']))

    job_results_final = []
    for key, val in job_results.items():
        val['DN'] = key[0]
        val['VO'] = key[1]
        val['VOMS'] = key[2]
        job_results_final.append(val)
    pilots_fname = rrd.path_with_spool(environ, "pilots")
    secure_json_write(pilots_fname, job_results_final)

    total_results['UpdateTime'] = time.time()
    totals_fname = rrd.path_with_spool(environ, "totals")
    secure_json_write(totals_fname, total_results)

    total_fname = rrd.check_rrd(environ, None, "jobs")
    if os.isatty(2):
        print >> sys.stderr, "Totals:", total_results
    rrdtool.update(total_fname, "N:%d:%d:%d" % (total_results["Running"], total_results["Idle"], total_results["Held"]))


    vos_fname = rrd.path_with_spool(environ, "vos.json")
    secure_json_write(vos_fname, vo_results)
    for vo, info in vo_results.items():
        vo_fname = rrd.check_rrd(environ, None, "vos", vo)
        if os.isatty(2):
            print >> sys.stderr, "VO", vo, "totals:", info
        rrdtool.update(vo_fname, "N:%d:%d:%d:%d" % (info['Running'], info['Idle'], info['Held'], info['Jobs']))


if __name__ == '__main__':
    main()

