#!/usr/bin/python
import classad
import htcondor
import time

# Ignoreable schedds - any failures when quering these schedds will be ignored
ignoreSchedFailures = ['']

vosa = []
jobsW = {}
jobsC = {}

data = ''

coll = htcondor.Collector()
results = coll.query(htcondor.AdTypes.Schedd, "true", ["Name"])
for result in results:
   host = result["Name"]
   scheddAd = coll.locate(htcondor.DaemonTypes.Schedd, host)
   schedd = htcondor.Schedd(scheddAd)
   try:
      ajobs = schedd.query('',["x509UserProxyVOName", "RequestCpus", "RequestMemory", "JobStatus", "JobCurrentStartDate", "RemoteSysCpu", "RemoteUserCpu", "isPreemptable"])
   except:
      print '# Unable to query schedd:',host,'so will try once more'
      if host not in ignoreSchedFailures:
         time.sleep(5)
         try:
            ajobs = schedd.query('',["x509UserProxyVOName", "RequestCpus", "RequestMemory", "JobStatus", "JobCurrentStartDate", "RemoteSysCpu", "RemoteUserCpu", "isPreemptable"])
         except:
            print '# Unable to query schedd:',host,'again'
            exit(1)
      else:
         print '# - ignoring failure'
         continue

   vos = []
   cpus = []

   jobsR = {}
   jobsI = {}
   jobsH = {}

   jobsRp = {}
   jobsIp = {}
   jobsHp = {}

   for ajob in ajobs:

      preemptable = 0
      if "isPreemptable" in ajob:
         if ajob["isPreemptable"]:
            preemptable = 1

      if "x509UserProxyVOName" not in ajob:
         ajob["x509UserProxyVOName"] = "undefined"

      if ajob["x509UserProxyVOName"] not in vos:
         vos.append(ajob["x509UserProxyVOName"])

      if ajob["x509UserProxyVOName"] not in vosa:
         vosa.append(ajob["x509UserProxyVOName"])

      if ajob["RequestCpus"] not in cpus:
         cpus.append(ajob["RequestCpus"])

      if ajob["JobStatus"] == 2:
         if not jobsR.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
            jobsR[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
         jobsR[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsR[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

         if preemptable == 1:
            if not jobsRp.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
               jobsRp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
            jobsRp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsRp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

         if "JobCurrentStartDate" in ajob and "RemoteSysCpu" in ajob and "RemoteUserCpu" in ajob:
            if not jobsW.has_key(ajob["x509UserProxyVOName"]):
               jobsW[ajob["x509UserProxyVOName"]] = 0
            jobsW[ajob["x509UserProxyVOName"]] += ajob["RequestCpus"]*(time.time() - ajob["JobCurrentStartDate"])

            if not jobsC.has_key(ajob["x509UserProxyVOName"]):
               jobsC[ajob["x509UserProxyVOName"]] = 0
            jobsC[ajob["x509UserProxyVOName"]] += ajob["RemoteSysCpu"] + ajob["RemoteUserCpu"]

      if ajob["JobStatus"] == 1:
         if not jobsI.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
            jobsI[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
         jobsI[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsI[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

         if preemptable == 1:
            if not jobsIp.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
               jobsIp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
            jobsIp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsIp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

      if ajob["JobStatus"] == 5:
         if not jobsH.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
            jobsH[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
         jobsH[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsH[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

         if preemptable == 1:
            if not jobsHp.has_key((ajob["x509UserProxyVOName"], ajob["RequestCpus"])):
               jobsHp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = 0
            jobsHp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] = jobsHp[(ajob["x509UserProxyVOName"], ajob["RequestCpus"])] + 1

   hostS = host.split('.')[0]
   for vo in vos:
      coresR = 0
      coresI = 0
      coresH = 0
      for cpu in cpus:
         if jobsR.has_key((vo, cpu)):
            data += 'jobs,vo='+vo+',status=running,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsR[(vo,cpu)])+'\n'
            coresR += jobsR[(vo,cpu)] * cpu

         if jobsI.has_key((vo, cpu)):
            data += 'jobs,vo='+vo+',status=idle,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsI[(vo,cpu)])+'\n'
            coresI += jobsI[(vo,cpu)] * cpu

         if jobsH.has_key((vo, cpu)):
            data += 'jobs,vo='+vo+',status=held,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsH[(vo,cpu)])+'\n'
            coresH += jobsH[(vo,cpu)] * cpu

         if jobsRp.has_key((vo, cpu)):
            data += 'pjobs,vo='+vo+',status=running,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsRp[(vo,cpu)])+'\n'

         if jobsIp.has_key((vo, cpu)):
            data += 'pjobs,vo='+vo+',status=idle,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsIp[(vo,cpu)])+'\n'

         if jobsHp.has_key((vo, cpu)):
            data += 'pjobs,vo='+vo+',status=held,schedd='+hostS+',cpus='+str(cpu)+' value='+str(jobsHp[(vo,cpu)])+'\n'

      data += 'cores,vo='+vo+',status=running,schedd='+hostS+' value='+str(coresR)+'\n'
      data += 'cores,vo='+vo+',status=idle,schedd='+hostS+' value='+str(coresI)+'\n'
      data += 'cores,vo='+vo+',status=held,schedd='+hostS+' value='+str(coresH)+'\n'

for vo in vosa:
   if vo in jobsC:
      eff = -1.0
      if jobsW[vo] > 0:
         eff = jobsC[vo]/jobsW[vo]*100.0
      data += 'cpu-efficiency,vo='+vo+' value='+str(eff)+'\n'

print data
