#!/usr/bin/python
# Copyright (c) Members of the EGEE Collaboration. 2004. 
# See http://www.eu-egee.org/partners/ for details on the copyright
# holders.  
#
# Licensed under the Apache License, Version 2.0 (the "License"); 
# you may not use this file except in compliance with the License. 
# You may obtain a copy of the License at 
#
#     http://www.apache.org/licenses/LICENSE-2.0 
#
# Unless required by applicable law or agreed to in writing, software 
# distributed under the License is distributed on an "AS IS" BASIS, 
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
# See the License for the specific language governing permissions and 
# limitations under the License.

import sys
import re
import socket
import time
import shlex
import subprocess
from threading import Thread




class ErrorHandler(Thread):

    def __init__(self, err_stream):
        Thread.__init__(self)
        self.stream = err_stream
        self.message = ""
    
    def run(self):
        line = self.stream.readline()
        while line:
            self.message = self.message + line
            line = self.stream.readline()


class LRMSVersionHandler(Thread):

    def __init__(self):
        Thread.__init__(self)
        self.version = None
        self.pRegex = re.compile('^\s*pbs_version\s*=([^$]+)$')
    
    def setStream(self, stream):
        self.stream = stream
      
    def run(self):
        line = self.stream.readline()
        while line:
            parsed = self.pRegex.match(line)
            if parsed:
                self.version = parsed.group(1).strip()
            line = self.stream.readline()


class CPUInfoHandler(Thread):

    def __init__(self):
        Thread.__init__(self)
        self.totalCPU = 0
        self.freeCPU = 0
        self.pRegex = re.compile('^\s*([^=\s]+)\s*=([^$]+)$')
    
    def setStream(self, stream):
        self.stream = stream
      
    def run(self):
        currState = None
        line = self.stream.readline()
        while line:
            parsed = self.pRegex.match(line)
            if parsed:
                if parsed.group(1) == 'state':
                
                    currState = parsed.group(2).strip()
                
                elif parsed.group(1) == 'np':
                
                    procNum = int(parsed.group(2).strip())
                    if currState <> 'down' and currState <> 'offline':
                        self.totalCPU += procNum
                    if currState == 'free':
                        self.freeCPU += procNum
                
                elif parsed.group(1) == 'jobs':
                    
                    jobs = parsed.group(2).strip().split(', ')
                    if currState == 'free':
                        self.freeCPU -= len(jobs)
                
            line = self.stream.readline()


class QueueInfoHandler(Thread):

    def __init__(self):
        Thread.__init__(self)
        self.pRegex = re.compile('^\s*([^=\s]+)\s*=([^$]+)$')
        self.cputRegex = re.compile('(\d+):(\d+):(\d+)')
        self.maxCPUtime = -1
        self.defaultCPUtime = -1
        self.maxPCPUtime = -1
        self.defaultPCPUtime = -1
        self.maxTotJobs = -1
        self.maxRunJobs = -1
        self.policyPriority = None
        self.maxWallTime = -1
        self.defaultWallTime = -1
        self.obtWallTime = -1
        self.maxProcCount = -1
        self.defaultProcCount = -1
        self.enabled = False
        self.started = False
        self.state = 'Closed'
        
    def setStream(self, stream):
        self.stream = stream
        
    def conv(self, strtime):
        parsed = self.cputRegex.match(strtime.strip())
        if parsed:
            return int(parsed.group(1)) * 3600 + int(parsed.group(2)) * 60 + int(parsed.group(3))
        if strtime.strip() <> '-':
            return int(strtime.strip())
        return 0

    def run(self):
        line = self.stream.readline()
        while line:
            parsed = self.pRegex.match(line)
            if parsed:
                if parsed.group(1) == 'resources_max.cput':
                
                    self.maxCPUtime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'resources_default.cput':
                
                    self.defaultCPUtime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'resources_max.pcput':
                
                    self.maxPCPUtime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'resources_default.pcput':
                
                    self.defaultPCPUtime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'max_queuable':
                
                    self.maxTotJobs = int(parsed.group(2).strip())
                    
                elif parsed.group(1) == 'Priority':
                    
                    self.policyPriority = parsed.group(2).strip()
                    
                elif parsed.group(1) == 'max_running':
                
                    self.maxRunJobs = int(parsed.group(2).strip())
                    
                elif parsed.group(1) == 'resources_max.walltime':
                
                    self.maxWallTime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'resources_default.walltime':
                
                    self.defaultWallTime = self.conv(parsed.group(2))
                    
                elif parsed.group(1) == 'resources_max.procct':
                
                    self.maxProcCount = int(parsed.group(2).strip())
                    
                elif parsed.group(1) == 'resources_default.procct':
                
                    self.defaultProcCount = int(parsed.group(2).strip())
                    
                elif parsed.group(1) == 'enabled' and parsed.group(2).strip() == 'True':
                    
                    self.enabled = True
                    
                elif parsed.group(1) == 'started' and parsed.group(2).strip() == 'True':
                
                    self.started = True



            line = self.stream.readline()

        if self.enabled:
            if self.started:
                self.state = 'Production'
            else:
                self.state = 'Queueing'
        else:
            if self.started:
                self.state = 'Draining'
            
        if self.defaultCPUtime <> -1 and self.defaultPCPUtime <> -1:
            self.defaultCPUtime = min(self.defaultCPUtime, self.defaultPCPUtime)
        if self.defaultPCPUtime <> -1:
            self.defaultCPUtime = self.defaultPCPUtime
            
        if self.maxCPUtime <> -1 and self.maxPCPUtime <> -1:
            self.maxCPUtime = min(self.maxCPUtime, self.maxPCPUtime)
        if self.maxPCPUtime <> -1:
            self.maxCPUtime = self.maxPCPUtime

        if self.maxCPUtime == -1 and self.defaultCPUtime <> -1:
            self.maxCPUtime = self.defaultCPUtime
        
        if self.maxWallTime <> -1:
            self.obtWallTime = self.maxWallTime
                
        if self.maxWallTime == -1 and self.defaultWallTime <> -1:
            self.maxWallTime = self.defaultWallTime

        if self.maxProcCount == -1 and self.defaultProcCount <> -1:
            self.maxProcCount = self.defaultProcCount


def runCommand(cmdline, stdout_thread):

    try:
        cmd = shlex.split(cmdline)
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
        stdout_thread.setStream(process.stdout)
        stderr_thread = ErrorHandler(process.stderr)
    
        stdout_thread.start()
        stderr_thread.start()
    
        ret_code = process.wait()
    
        stdout_thread.join()
        stderr_thread.join()
    
        return (ret_code, stderr_thread.message)   

    except:
        etype, evalue, etraceback = sys.exc_info()
        return (-125, "%s: (%s)" % (etype, evalue))



def readDNsAndAttr(filename, dnRE, queueRE):
    dnsAndQueues = dict()
    ldifFile = None
    try:
    
        ldifFile = open(filename)
        
        for line in ldifFile.readlines():
            
            tmpm = dnRE.match(line)
            if tmpm <> None:
                currDN = line.strip()
                continue
                
            tmpm = queueRE.match(line)
            if tmpm <> None:
                dnsAndQueues[currDN] = tmpm.group(1).strip()
        
    finally:
        if ldifFile:
            ldifFile.close()

    return dnsAndQueues




def main():
    pRegex = re.compile('^\s*([^=\s]+)\s*=([^$]+)$')
    
    config = {}
    
    hostname = socket.getfqdn()
    now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

    conffile = None
    foundErr = False
    try:
        if len(sys.argv) <> 2:
            raise Exception("Usage: info-dynamic-pbs <config-file>")

        conffile = open(sys.argv[1])
        for line in conffile:
            parsed = pRegex.match(line)
            if parsed:
                config[parsed.group(1)] = parsed.group(2).strip(' \n\t"')
            else:
                tmps = line.strip()
                if len(tmps) > 0 and not tmps.startswith('#'):
                    raise Exception("Error parsing configuration file " + sys.argv[1])

        if not "outputformat" in config:
            if "GlueFormat" in config:
                config["outputformat"] = config["GlueFormat"]
            else:
                config["outputformat"] = "both"
        if config["outputformat"] not in ["glue1", "glue2", "both"]:
            raise Exception("FATAL: Unknown output format specified in config file:%s" % config["outputformat"])

        if not "glue1-static-file-CE" in config:
            config["glue1-static-file-CE"] = "/var/lib/bdii/gip/ldif/static-file-CE.ldif"
        
        if not "glue2-static-file-computing-manager" in config:
            config["glue2-static-file-computing-manager"] = "/var/lib/bdii/gip/ldif/ComputingManager.ldif"
        
        if not "glue2-static-file-computing-share" in config:
            config["glue2-static-file-computing-share"] = "/var/lib/bdii/gip/ldif/ComputingShare.ldif"
            
        if not "pbs-host" in config:
            config["pbs-host"] = hostname

    except Exception, ex:
        sys.stderr.write(str(ex) + '\n')
        foundErr = True

    if conffile:
        conffile.close()
    if foundErr:
        sys.exit(1)


    glue1QueueTable = None
    glue2QueueTable = None
    managerTable = None
    
    if config['outputformat'] <> "glue2":
        
        glue1DNRegex = re.compile("dn:\s*GlueCEUniqueID\s*=\s*[^$]+")
        glue1QueueRegex = re.compile("GlueCEName\s*:\s*([^$]+)")

        try:

            glue1QueueTable = readDNsAndAttr(config["glue1-static-file-CE"], glue1DNRegex, glue1QueueRegex)
        
        except Exception, ex:
            sys.stderr.write(str(ex) + '\n')
            sys.exit(1)




    if config['outputformat'] <> "glue1":
    
        glue2DNRegex = re.compile("dn:\s*GLUE2ShareID\s*=\s*[^$]+")
        glue2QueueRegex = re.compile("GLUE2ComputingShareMappingQueue\s*:\s*([^$]+)")

        managerRegex = re.compile("dn:\s*GLUE2ManagerId\s*=\s*[^$]+")
        manAttrRegex = re.compile("GLUE2ManagerID\s*:\s*([^$]+)")
        
        try:

            glue2QueueTable = readDNsAndAttr(config["glue2-static-file-computing-share"], glue2DNRegex, glue2QueueRegex)
            managerTable = readDNsAndAttr(config["glue2-static-file-computing-manager"], managerRegex, manAttrRegex)
            
        except Exception, ex:
            sys.stderr.write(str(ex) + '\n')
            sys.exit(1)

    verHandler = LRMSVersionHandler()
    exitCode, errMsg = runCommand("qstat -B -f %s" % config["pbs-host"], verHandler)
    if exitCode <> 0:
        sys.stderr.write(errMsg + '\n')
        sys.exit(1)
        
    cpuInfoHandler = CPUInfoHandler()
    exitCode, errMsg = runCommand("pbsnodes -a -s %s" % config["pbs-host"], cpuInfoHandler)
    if exitCode <> 0:
        sys.stderr.write(errMsg + '\n')
        sys.exit(1)
    

    allQueues = set()
    for queue in glue1QueueTable.values():
        allQueues.add(queue)
    for queue in glue2QueueTable.values():
        allQueues.add(queue)
    
    
    qInfoHandlers = dict()
    for queue in allQueues:
        
        qInfoHandlers[queue] = QueueInfoHandler()
        exitCode, errMsg = runCommand("qstat -Q -f %s\@%s" % (queue, config["pbs-host"]), qInfoHandlers[queue])
        if exitCode <> 0:
            sys.stderr.write(errMsg + '\n')
            sys.exit(1)

    out = sys.stdout
    
    if config['outputformat'] <> "glue2":

        for glue1DN in glue1QueueTable:
        
            queue = glue1QueueTable[glue1DN]
            qInfo = qInfoHandlers[queue]
            
            out.write(glue1DN + '\n')
            if verHandler.version:
                out.write('GlueCEInfoLRMSVersion: %s\n' % verHandler.version)
            
            out.write('GlueCEInfoTotalCPUs: %d\n' % cpuInfoHandler.totalCPU)
            out.write('GlueCEPolicyAssignedJobSlots: %d\n' % cpuInfoHandler.totalCPU)
            out.write('GlueCEStateFreeCPUs: %d\n' % cpuInfoHandler.freeCPU)
            
            if qInfo.maxCPUtime <> -1:
                out.write('GlueCEPolicyMaxCPUTime: %d\n' % (qInfo.maxCPUtime / 60))
                out.write('GlueCEPolicyMaxObtainableCPUTime: %d\n' % (qInfo.maxCPUtime / 60))
                
            if qInfo.maxTotJobs <> -1:
                out.write('GlueCEPolicyMaxTotalJobs: %d\n' % qInfo.maxTotJobs)
                
            if qInfo.policyPriority:
                out.write('GlueCEPolicyPriority: %s\n' % qInfo.policyPriority)
                
            if qInfo.maxRunJobs <> -1:
                out.write('GlueCEPolicyMaxRunningJobs: %d\n' % qInfo.maxRunJobs)
                
            if qInfo.maxTotJobs <> -1 and qInfo.maxRunJobs <> -1:
                out.write('GlueCEPolicyMaxWaitingJobs: %d\n' % (qInfo.maxTotJobs - qInfo.maxRunJobs))
                
            if qInfo.maxWallTime <> -1:
                out.write('GlueCEPolicyMaxWallClockTime: %d\n' % (qInfo.maxWallTime / 60))
                
            if qInfo.obtWallTime <> -1:
                out.write('GlueCEPolicyMaxObtainableWallClockTime: %d\n' % (qInfo.obtWallTime / 60))
                
            if qInfo.maxProcCount <> -1:
                out.write('GlueCEPolicyMaxSlotsPerJob: %d\n' % qInfo.maxProcCount)
                
            out.write('GlueCEStateStatus: %s\n' % qInfo.state)
            out.write('\n')


    if config['outputformat'] <> "glue1":

        for managerDN in managerTable:
        
            out.write(managerDN + '\n')
            if verHandler.version:
                out.write('GLUE2ManagerProductVersion: %s\n' % verHandler.version)
            out.write('GLUE2EntityCreationTime: %s\n' % now)
            out.write('\n')
            
        for glue2DN in glue2QueueTable:
            queue = glue2QueueTable[glue2DN]
            qInfo = qInfoHandlers[queue]
            
            out.write(glue2DN + '\n')
            if qInfo.maxCPUtime <> -1:
                out.write('GLUE2ComputingShareDefaultCPUTime: %d\n' % qInfo.maxCPUtime)
                out.write('GLUE2ComputingShareMaxCPUTime: %d\n' % qInfo.maxCPUtime)
            if qInfo.maxWallTime <> -1:
                out.write('GLUE2ComputingShareDefaultWallTime: %d\n' % qInfo.maxWallTime)
            if qInfo.obtWallTime <> -1:
                out.write('GLUE2ComputingShareMaxWallTime: %d\n' % qInfo.obtWallTime)
            if qInfo.maxProcCount <> -1:
                out.write('GLUE2ComputingShareMaxSlotsPerJob: %d\n' % qInfo.maxProcCount)
                
            out.write('GLUE2ComputingShareServingState: %s\n' % qInfo.state.lower())
            out.write('GLUE2EntityCreationTime: %s\n' % now)
            out.write('\n')



if __name__ == "__main__":
    main()


