#!/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 time
import shlex
import subprocess
import logging
import logging.config
from threading import Thread

from TorqueInfoUtils import CommonUtils
from TorqueInfoUtils import QStatHandler
from TorqueInfoUtils import PBSNodesHandler
from TorqueInfoUtils import NvidiaSMIHandler

from TorqueInfoUtils.QStatHandler import RES_UNDEF

MAX_INT32 = 2**31-1
MAX_UINT32 = 2**32-1
MAX_UINT64 = 2**64-1

def main():
    
    now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

    if len(sys.argv) <> 2:
        sys.stderr.write("Usage: info-dynamic-pbs <config-file>\n")
        sys.exit(1)

    try:
        logging.config.fileConfig(sys.argv[1])
    except Exception, conf_log_err:
        logging.basicConfig(stream=sys.stderr)
    

    glue1QueueTable = dict()
    glue2QueueTable = dict()
    managerTable = None
    allQueues = set()
    freeGPUSlots = 0
    usedGPUSlots = 0
    
    try:
    
        config = CommonUtils.readConfigFile(sys.argv[1])
        
        if config['outputformat'] <> "glue2":
        
            glue1QueueTable = CommonUtils.parseLdif(config["bdii-configfile"], 'GLUE1')
        
            for queue in glue1QueueTable.values():
                allQueues.add(queue)

        if config['outputformat'] <> "glue1":
        
            glue2QueueTable, managerTable = CommonUtils.parseLdif(config["bdii-configfile"], 'GLUE2')

            for queue in glue2QueueTable.values():
                allQueues.add(queue)

        lrmsVer = QStatHandler.parseLRMSVersion(config["pbs-host"])
        
        cpuInfoHandler = PBSNodesHandler.parseCPUInfo(config["pbs-host"])
    
        qInfoHandlers = QStatHandler.parseAllQueuesInfo(allQueues, config["pbs-host"])
            
        if config['enable_glue_2_1']:

            for nodeName in cpuInfoHandler.gpuTable:
            
                gpuStats = cpuInfoHandler.gpuTable[nodeName]
                nodeState = gpuStats['node_state']
                if 'down' in nodeState or 'offline' in nodeState or 'unknown' in nodeState:
                    continue
                else:
                    tmpSlots = gpuStats['total_gpus']

                try:

                    smiHandler = NvidiaSMIHandler.parseGPUInfo(nodeName)
                    for nProcs in smiHandler.num_of_procs.values():
                        if nProcs > 0:
                            tmpSlots -= 1

                    freeGPUSlots += tmpSlots
                    usedGPUSlots += gpuStats['total_gpus'] - tmpSlots

                except Exception, ex:
                    sys.stderr.write(repr(ex) + '\n')
                    freeGPUSlots += gpuStats['free_gpus']
                    usedGPUSlots += gpuStats['total_gpus'] - gpuStats['free_gpus']

    except Exception, ex:
        sys.stderr.write(str(ex) + '\n')
        sys.exit(2)

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

        for glue1DN in glue1QueueTable:
        
            queue = glue1QueueTable[glue1DN]
            qInfo = qInfoHandlers[queue]
            
            out.write(glue1DN + '\n')
            out.write('GlueCEInfoLRMSVersion: %s\n' % lrmsVer)
            
            out.write('GlueCEInfoTotalCPUs: %d\n' % cpuInfoHandler.totalCPU)
            out.write('GlueCEPolicyAssignedJobSlots: %d\n' % cpuInfoHandler.totalCPU)
            out.write('GlueCEStateFreeCPUs: %d\n' % cpuInfoHandler.freeCPU)
            
            if qInfo.defaultCPUtime <> RES_UNDEF:
                out.write('GlueCEPolicyMaxCPUTime: %d\n' % (qInfo.defaultCPUtime / 60))
            else:
                out.write('GlueCEPolicyMaxCPUTime: %d\n' % MAX_INT32)

            if qInfo.maxCPUtime <> RES_UNDEF:
                out.write('GlueCEPolicyMaxObtainableCPUTime: %d\n' % (qInfo.maxCPUtime / 60))
            else:
                out.write('GlueCEPolicyMaxObtainableCPUTime: %d\n' % MAX_INT32)
                
            if qInfo.maxTotJobs <> RES_UNDEF:
                out.write('GlueCEPolicyMaxTotalJobs: %d\n' % qInfo.maxTotJobs)
            else:
                out.write('GlueCEPolicyMaxTotalJobs: %d\n' % MAX_INT32)
                
            if qInfo.policyPriority:
                out.write('GlueCEPolicyPriority: %s\n' % qInfo.policyPriority)
            else:
                out.write('GlueCEPolicyPriority: %s\n' % MAX_INT32)
                
            if qInfo.maxRunJobs <> RES_UNDEF:
                out.write('GlueCEPolicyMaxRunningJobs: %d\n' % qInfo.maxRunJobs)
            else:
                out.write('GlueCEPolicyMaxRunningJobs: %d\n' % MAX_INT32)
                
            if qInfo.maxTotJobs <> RES_UNDEF and qInfo.maxRunJobs <> RES_UNDEF:
                out.write('GlueCEPolicyMaxWaitingJobs: %d\n' % (qInfo.maxTotJobs - qInfo.maxRunJobs))
            else:
                out.write('GlueCEPolicyMaxWaitingJobs: %d\n' % MAX_INT32)
                
            if qInfo.defaultWallTime <> RES_UNDEF:
                out.write('GlueCEPolicyMaxWallClockTime: %d\n' % (qInfo.defaultWallTime / 60))
            else:
                out.write('GlueCEPolicyMaxWallClockTime: %d\n' % MAX_INT32)
                
            if qInfo.maxWallTime <> RES_UNDEF:
                out.write('GlueCEPolicyMaxObtainableWallClockTime: %d\n' % (qInfo.maxWallTime / 60))
            else:
                out.write('GlueCEPolicyMaxObtainableWallClockTime: %d\n' % MAX_INT32)
                
            if qInfo.maxProcCount <> RES_UNDEF:
                out.write('GlueCEPolicyMaxSlotsPerJob: %d\n' % qInfo.maxProcCount)
            else:
                out.write('GlueCEPolicyMaxSlotsPerJob: %d\n' % MAX_INT32)
 
            if CommonUtils.interfaceIsOff(config):
                out.write('GlueCEStateStatus: Draining\n')
            else:
                out.write('GlueCEStateStatus: %s\n' % qInfo.state)
                
            out.write('\n')


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

        for managerDN in managerTable:
        
            out.write(managerDN + '\n')
            out.write('GLUE2ManagerProductVersion: %s\n' % lrmsVer)
            out.write('GLUE2EntityCreationTime: %s\n' % now)
            out.write('GLUE2ComputingManagerTotalAcceleratorSlots: GPU:%d\n' % (freeGPUSlots + usedGPUSlots))
            out.write('GLUE2ComputingManagerUsedAcceleratorSlots: GPU:%d\n' % usedGPUSlots)
            out.write('\n')
            
        for glue2DN in glue2QueueTable:
            queue = glue2QueueTable[glue2DN]
            qInfo = qInfoHandlers[queue]
            
            out.write(glue2DN + '\n')
            if qInfo.defaultCPUtime <> RES_UNDEF:
                out.write('GLUE2ComputingShareDefaultCPUTime: %d\n' % qInfo.defaultCPUtime)
            else:
                out.write('GLUE2ComputingShareDefaultCPUTime: %d\n' % MAX_UINT64)
                
            if qInfo.maxCPUtime <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxCPUTime: %d\n' % qInfo.maxCPUtime)
            else:
                out.write('GLUE2ComputingShareMaxCPUTime: %d\n' % MAX_UINT64)
                
            if qInfo.defaultWallTime <> RES_UNDEF:
                out.write('GLUE2ComputingShareDefaultWallTime: %d\n' % qInfo.defaultWallTime)
            else:
                out.write('GLUE2ComputingShareDefaultWallTime: %d\n' % MAX_UINT64)
                
            if qInfo.maxWallTime <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxWallTime: %d\n' % qInfo.maxWallTime)
            else:
                out.write('GLUE2ComputingShareMaxWallTime: %d\n' % MAX_UINT64)
                
            if qInfo.maxProcCount <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxSlotsPerJob: %d\n' % qInfo.maxProcCount)
            else:
                out.write('GLUE2ComputingShareMaxSlotsPerJob: %d\n' % MAX_UINT32)
                
            # take maxjobs into consideration for max-running and max-waiting as well
            if qInfo.maxTotJobs <> RES_UNDEF:
                maxjobs = qInfo.maxTotJobs
            else:
                maxjobs = MAX_UINT32

            if qInfo.maxRunJobs <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxRunningJobs: %d\n' % qInfo.maxRunJobs)
            else:
                out.write('GLUE2ComputingShareMaxRunningJobs: %d\n' % maxjobs)
            
            #
            # TODO get info per vo (vomaxjobs-*)
            #      for the moment queue-wide values are used
            #    
            out.write('GLUE2ComputingShareMaxTotalJobs: %d\n' % maxjobs)
                
            if qInfo.maxTotJobs <> RES_UNDEF and qInfo.maxRunJobs <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxWaitingJobs: %d\n' % (qInfo.maxTotJobs - qInfo.maxRunJobs))
            else:
                out.write('GLUE2ComputingShareMaxWaitingJobs: %d\n' % maxjobs)
                
            if qInfo.maxMem <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxMainMemory: %d\n' % qInfo.maxMem)
            else:
                out.write('GLUE2ComputingShareMaxMainMemory: %d\n' % MAX_UINT64)
                
            if qInfo.maxVMem <> RES_UNDEF:
                out.write('GLUE2ComputingShareMaxVirtualMemory: %d\n' % qInfo.maxVMem)
            else:
                out.write('GLUE2ComputingShareMaxVirtualMemory: %d\n' % MAX_UINT64)
            
            if CommonUtils.interfaceIsOff(config):    
                out.write('GLUE2ComputingShareServingState: draining\n')
            else:
                out.write('GLUE2ComputingShareServingState: %s\n' % qInfo.state.lower())
                
            out.write('GLUE2EntityCreationTime: %s\n' % now)
            
            if config['enable_glue_2_1']:

                out.write('GLUE2ComputingShareFreeAcceleratorSlots: GPU:%d\n' % freeGPUSlots)
                out.write('GLUE2ComputingShareUsedAcceleratorSlots: GPU:%d\n' % usedGPUSlots)
            
            out.write('\n')



if __name__ == "__main__":
    main()


