#!/usr/bin/env python

"""
wmcore-SysStat

Hardware monitoring utility for launched sensor daemons.
"""
from __future__ import print_function

import sys
import os
import getopt
import subprocess
import time
import signal
import string
from datetime import datetime
from WMCore.Agent.Daemon.Details import Details
from WMCore.Configuration import loadConfigurationFile

########################
### Helper functions ###
########################

def error(msg):
    """
    General error handler.
    """
    
    print("ERROR:", msg) 
    sys.exit(1)
    
def warning(msg):
    """
    General warning handler.
    """
    
    print("WARNING:", msg)

def extractList(targetType, arg):
    """
    Makes python list of comma sepeated string list.
    """
    
    global ITEM_TYPE
    
    targetList = arg.split(",")
    listItems = []
    for item in targetList:
        if len(item.strip()) == 0:
            continue
        if checkItem(targetType, item) == False:
            warning("Ignoring target %s %s" % (ITEM_TYPE[targetType], item))
            continue
        listItems.append(item)    
    return listItems
    
def checkOutput(arg):
    """
    Checks string if it contains the name of supported output format.
    """
    
    global OUT_XML, OUT_TUI, OUT_JSON, OUT_CSV
    
    if arg == "XML":
        return OUT_XML
    elif arg == "TUI":
        return OUT_TUI
    elif arg == "JSON":
        return OUT_JSON
    elif arg == "CSV":
        return OUT_CSV
    else:
        return None

def help(short=True):
    """
    Returns short or long help information.
    """
    
    short_help = \
    """
Usage: wmcore-SysStat --config=<WMCoreConfig.py> <command> {<specified_list>} [optional]
       Where:
       command: --start, --shutdown, --status, --restart, --help, --list=
       specified_list: --wmcomponents=, --services=, --resources=, --disks=
       optional: --ouput=
    """
    
    long_help = \
    """
Help: '--config=' option is obligatory, you must always provide WMCore
      configuration file. It must be the first option passed to
      wmcore-SysStat.
      
      This utility is case-sensitive!
      
      Only one command can be provided for wmcore-SysStat:
        --start - start sensors for all/specified items;
        --shutdown - stop sensors for all/specified items;
        --status - give status of sensors for all/specified items;
        --restart - restarts sensors for all/specified items;
        --help - prints this help information.
        --list=<list> - prints all available targets.
          Possible values: WMComponent, Service, Resource, Disk
        
      You can specify comma separeted lists for different targets:
        --wmcomponents=<list> - for WMComponents;
        --services=<list> - for Services;
        --resources=<list> - for general machine Resources;
        --disks=<list> - for machine Disks;
        
      You can speficy output format for '--status' and '--list=' 
      commands by providing '--ouput=' option. Possible values:
        TUI - Text User Interface (best for terminal);
        XML - Best for third-party applications;
        JSON - Best for web-based applications;
        CSV - Best for machine and human;
        
      <list> - comma separated list, eg. 'GridFTP,mySQL'
      
      You can turn on debug mode on wmcore-sensord, by using '--debug'
      option on wmcore-SysStat.
    """
    
    if short == True:
        return short_help
    else:
        return short_help + long_help
        
def checkItem(targetType, item):
    """
    Checks if items belongs to specific family.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global config, servicesList, resourcesList
    
    if targetType == ITEM_WMCOMPONENT:
        wmComponentsList = config.listComponents_()
        if item not in wmComponentsList:
            return False
    elif targetType == ITEM_SERVICE:
        if item not in servicesList:
            return False
    elif targetType == ITEM_RESOURCE:
        if item not in resourcesList:
            return False
    elif targetType == ITEM_DISK:
        if item not in disks():
            return False
            
    return True

def disks():
    """
    Retuns list of available disks on machine.
    """
    
    disks = os.popen('ls -1 /dev/?d[a-z]').readlines()
    disks = [x.split('/')[2].rstrip() for x in disks]
    return disks

def loadConfiguration():
    """
    Loads WMCore configuration file.
    """
    
    global config, configFile
    
    try:
        if config == None:
            config = loadConfigurationFile(configFile)
            
            if "HWMon" not in dir(config):
                error("No 'HWMon' section found in WMCore configuration file! Please, check your configuration file.")
    except ImportError as ex:
        error("Can no load configuration file! Please, check configuration file (" + os.path.abspath(configFile) + ")")

def prepareList(targetType, arg):
    """
    Parse target list.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global loadDefaults, config
    
    loadDefaults = False
    
    if targetType == ITEM_WMCOMPONENT:
        loadConfiguration() # must load configuratio to confirm WMComponents
        
    if len(arg) != 0:
        listItems = extractList(targetType, arg)
        if len(listItems) != 0:
            specifiedList[targetType] = listItems
    else:
        if targetType == ITEM_WMCOMPONENT:
            specifiedList[ITEM_WMCOMPONENT] = config.listComponents_()
        elif targetType == ITEM_SERVICE:
            specifiedList[ITEM_SERVICE] = list(servicesList)
        elif targetType == ITEM_RESOURCE:
            specifiedList[ITEM_RESOURCE] = resourcesList
        elif targetType == ITEM_DISK:
            specifiedList[ITEM_DISK] = disks()

def isWMComponentRunning(wmcomponent):
    """
    Checks if WMComponent is running.
    """
    
    global config
    
    daemonXml = os.path.abspath("%s/Components/%s/Daemon.xml" % (config.General.workDir, wmcomponent))
    
    if not os.path.isfile(daemonXml):
        return False, 0
        
    daemon = Details(daemonXml)
    if daemon.isAlive():
        return True, int(daemon['ProcessID'])
    else:
        return False, 0
        
def isServiceRunning(service):
    """
    Checks if Service is running.
    """
    
    global servicesList
    
    if service not in servicesList:
        return False, 0
    
    ret = os.popen("ps -C %s wwho pid" % servicesList[service]).read()
    if ret:
        return True, int(ret.strip()) 
    else:
        return False, 0

def isSensorDaemonRunning(target):
    """
    Checks if wmcore-sensord is running for target.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global specifiedList
    
    patterns = {
        ITEM_WMCOMPONENT : "\\-\\-wmcomponent=%s", 
        ITEM_SERVICE     : "\\-\\-service=%s",
        ITEM_RESOURCE    : "\\-\\-resource=%s", 
        ITEM_DISK        : "\\-\\-disk=%s"
    }
    
    pattern = None
    targetTypes = (ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK)
    for ttype in targetTypes:
        if ttype in specifiedList and target in specifiedList[ttype]:
            pattern = patterns[ttype] % target
            break
    
    if pattern == None:
        return False, 0
           
    ret = os.popen("ps -C python wwho pid,cmd | grep -i 'wmcore-sensord' | grep -i '%s'" % pattern).read()
    if ret:
        return True, int(ret.split()[0])
    else:
        return False, 0
    
def isSensorRunning(target):
    """
    Checks if sar/iostat is running for target.
    """
    
    ret = None
    if ITEM_DISK in specifiedList and target in specifiedList[ITEM_DISK]:
        ret = os.popen("ps -C iostat wwho pid,cmd | grep '%s'" % target).read()  
    elif ITEM_RESOURCE in specifiedList and target in specifiedList[ITEM_RESOURCE]:
        patterns = {'CPU': '\\-u', 'MEM': '\\-r', 'SWAP': '\\-W', 'LOAD': '\\-q'}
        ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % patterns[target]).read()
    elif ITEM_WMCOMPONENT in specifiedList and target in specifiedList[ITEM_WMCOMPONENT]:
        stat = isWMComponentRunning(target)
        if stat[0] == True:
            ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % stat[1]).read()
    elif ITEM_SERVICE in specifiedList and target in specifiedList[ITEM_SERVICE]:
        stat = isServiceRunning(target)
        if stat[0] == True:
            ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % stat[1]).read()
    
    if ret:
        return True, int(ret.split()[0])
    else:   
        return False, 0

def kill(pid, signal):
    """
    Sends singal to specified process.
    """
    
    try:
        os.kill(int(pid), signal)
    except OSError as ex:
        err = str(err)
        if err.find("No such process") >= 0:
            pass
        else:
            raise

def formatOutput(caller, data):
    """
    Prints '--status' information in '--output=' specified fromat.
    """
    
    global CMD_TYPE, OUT_TYPE
    global output
    
    def status_XML(data):
        msg = "<?xml version=\"1.0\"?>\n"
        msg += "<targets>\n"
        for ttype in data:
            for target in data[ttype]:
                msg += "<target type=\"%s\" name=\"%s\">\n" % (ttype, target)
                for status in  data[ttype][target]:
                    msg += "<status type=\"%s\" status=\"%s\" pid=\"%s\"/>\n" % \
                           (status, data[ttype][target][status][0], data[ttype][target][status][1])
                msg += "</target>\n"
        msg += "</targets>"
        return msg
    
    # TODO: Other formats
    def status_TUI(data):
        def showPid(format, pid):
            if pid != 0:
                return str(pid)
            else:
                return format[2]
                
        def showStatus(format, status, ttype = "NONE"):
            if ttype in ["DISK", "RESOURCE"]:
                return format[2]
            else:
                return format[status]
                
        status = ["\033[41m OFF  \033[m", "\033[42m  ON  \033[m", "\033[43m  NO  \033[m"]
        
        msg = "\033[30m\033[47m%-11s | %-40s | %-5s | %-6s | %-5s | %-6s | %-5s | %-6s\033[m\n" %\
              ("Target Type", "Target Name", "Target", "-->Pid", "Daemon", "-->Pid", "Sensor", "-->Pid")
        for ttype in data:
            for target in data[ttype]:
                msg += "%-11s | %-40s | %-5s | %-6s | %-5s | %-6s | %-5s | %-6s\n" %\
                       (ttype, target, showStatus(status, data[ttype][target]['TARGET'][0], ttype), showPid(status, data[ttype][target]['TARGET'][1]), \
                       showStatus(status, data[ttype][target]['DAEMON'][0]), showPid(status, data[ttype][target]['DAEMON'][1]), \
                       showStatus(status, data[ttype][target]['SENSOR'][0]), showPid(status, data[ttype][target]['SENSOR'][1]))
                    
        return msg
    
    def list_XML(data):
        msg = "<?xml version=\"1.0\"?>\n"
        msg += "<targets>\n"
        for ttype in data:
            for target in data[ttype]:
                msg += "<target type=\"%s\">%s</target>\n" % (ttype, target)
        msg += "</targets>"
        return msg
        
    def list_TUI(data):
        msg = ""
        msg += "\033[30m\033[47m%-11s | %-40s \033[m\n" % ("Target Type", "Target")
        for ttype in data:
            for target in data[ttype]:
                msg += "%-11s | %-56s\n" % (ttype, target)
        return msg
    
    template = CMD_TYPE[caller].lower() + "_" + OUT_TYPE[output]
    if template in locals():
        return locals()[template](data)
    else:
        return "Can not find template '%s'" % template

def start():
    """
    Starts sensor daemons for specified targets.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK, ITEM_TYPE
    global specifiedList, configFile, debugMode
    
    patterns = {
        ITEM_WMCOMPONENT : "--wmcomponent=%s", 
        ITEM_SERVICE     : "--service=%s",
        ITEM_RESOURCE    : "--resource=%s", 
        ITEM_DISK        : "--disk=%s"
    }
    
    status = ["\033[41m   OFF   \033[m", "\033[42m   ON    \033[m", "\033[43m WARNING \033[m"]
    
    for ttype in specifiedList:
        for target in specifiedList[ttype]:
            statSensorDaemon = isSensorDaemonRunning(target)
            if statSensorDaemon[0]:
                print("%-7s | Sensor daemon (pid: %s) is already running for %s %s." % (status[2], str(statSensorDaemon[1]), ITEM_TYPE[ttype], target))
                continue
            
            #statSensor = isSensorRunning(target)
            #if statSensor[0]:
            #    print "%-7s | Sensor (pid: %s) is already running for %s %s." % (status[2], str(statSensor[1]), ITEM_TYPE[ttype], target)
            #    continue
              
            if ttype == ITEM_WMCOMPONENT and not isWMComponentRunning(target)[0]:
                print("%-7s | %s %s is not running." % (status[2], ITEM_TYPE[ttype], target))
                continue
                
            if ttype == ITEM_SERVICE and not isServiceRunning(target)[0]:
                print("%-7s | %s %s is not running." % (status[2], ITEM_TYPE[ttype], target))
                continue
            
            cmd = ["wmcore-sensord"]
            cmd.append("--config=%s" % configFile)
            cmd.append(patterns[ttype] % target)
            if debugMode == True:
                cmd.append("--debug")
            sensord = subprocess.Popen(cmd)
            time.sleep(1)
            sensord.poll()
            
            pid = 0
            if ttype == ITEM_WMCOMPONENT:
                pid = isWMComponentRunning(target)[1]
            elif ttype == ITEM_SERVICE:
                pid = isServiceRunning(target)[1]
            
            if sensord.returncode == None:
                if pid != 0:
                    print("%-7s | Sensor daemon (pid: %s) for %s %s (pid: %s) started." \
                          % (status[1], str(sensord.pid), ITEM_TYPE[ttype], target, pid))
                else:
                    print("%-7s | Sensor daemon (pid: %s) for %s %s started." \
                          % (status[1], str(sensord.pid), ITEM_TYPE[ttype], target))
            else:                
                if pid != 0:
                    print("%-7s | Sensor daemon for %s %s (pid: %s) did not start." % (status[0], ITEM_TYPE[ttype], target, str(pid)))
                else:
                    print("%-7s | Sensor daemon for %s %s did not start." % (status[0], ITEM_TYPE[ttype], target))
    
def shutdown():
    """
    Shutdowns sensor daemons for specified targets.
    """
    
    global ITEM_TYPE
    
    for ttype in specifiedList:
        for target in specifiedList[ttype]:
            stat = isSensorDaemonRunning(target)
            if stat[0]:
                print(">>> Shutting down sensor daemon (pid: %s) for %s %s" % (stat[1], ITEM_TYPE[ttype], target))
                kill(stat[1], signal.SIGTERM)
    
def restart():
    """
    Restarts sensor daemons for specified targets.
    """
    
    print(">>> Shutting down specified sensor daemons.")
    shutdown()
    print(">>> Please wait 10 sec.")
    startTime = datetime.today()
    while (datetime.today() - startTime).seconds <= 10:
        time.sleep(1)
    print(">>> Starting up specified sensor daemons.")
    start()

def status():
    """
    Gives sensors information.
    """
    
    def showPid(status, pid):    
        if pid != 0:
            return str(pid)
        else:
            return status[2]
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_TYPE
    
    status = ["\033[41m OFF  \033[m", "\033[42m  ON  \033[m", "\033[43m  NO  \033[m"]
    
    if len(specifiedList) == 0:
        print("No information for specified targets.")
        return
    
    data = {}
    
    for ttype in specifiedList:
        data[ITEM_TYPE[ttype]] = {}
        for target in specifiedList[ttype]:
            statSensorDaemon = isSensorDaemonRunning(target)
            statSensor = isSensorRunning(target)
            stat = None
            
            data[ITEM_TYPE[ttype]][target] = {}
            
            if ttype == ITEM_WMCOMPONENT:
                stat = isWMComponentRunning(target)

            elif ttype == ITEM_SERVICE:
                stat = isServiceRunning(target)
            else:
                stat = [False, 0]
                
            data[ITEM_TYPE[ttype]][target]["TARGET"] = [stat[0], stat[1]]
            data[ITEM_TYPE[ttype]][target]["DAEMON"] = [statSensorDaemon[0], statSensorDaemon[1]]
            data[ITEM_TYPE[ttype]][target]["SENSOR"] = [statSensor[0], statSensor[1]]
            
    print(formatOutput(CMD_STATUS, data))

def prepareAvailableList(arg):
    """
    Prepares available WMComponents, Resources, Services, Disks list.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK, ITEM_TYPE
    global config, servicesList, resourcesList
    
    valid = ['WMComponent', 'Service', 'Resource', 'Disk']
    
    listAvail = {}
    
    loadConfiguration()
    if len(arg) == 0:
        listAvail[ITEM_WMCOMPONENT] = config.listComponents_()
        listAvail[ITEM_SERVICE] = list(servicesList)
        listAvail[ITEM_RESOURCE] = resourcesList
        listAvail[ITEM_DISK] = disks()
    else:    
        targetTypeList = arg.split(",")
        for ttype in targetTypeList:
            if len(ttype.strip()) == 0:
                continue
            if ttype in valid:
                targetType = ITEM_TYPE.index(ttype.upper())
                if targetType == ITEM_WMCOMPONENT:
                    listAvail[ITEM_WMCOMPONENT] = config.listComponents_()
                elif targetType == ITEM_SERVICE:
                    listAvail[ITEM_SERVICE] = list(servicesList)
                elif targetType == ITEM_RESOURCE:
                    listAvail[ITEM_RESOURCE] = resourcesList
                elif targetType == ITEM_DISK:
                    listAvail[ITEM_DISK] = disks()
            else:
                warning("Ignoring target type: %s." % ttype)
            
    return listAvail

def available():
    """
    Returns available WMComponents, Resources, Services, Disk string.
    """
    
    global ITEM_TYPE
    global availableList
    
    data = {}
    if len(availableList) == 0:
        return  "No targets for specified target types."
    
    for ttype in availableList:
        data[ITEM_TYPE[ttype]] = availableList[ttype]
    
    return formatOutput(CMD_LIST, data)

##########################
### Global definitions ###
##########################

servicesList = {
    'GridFTP' : 'globus-gridftp-server',
    'mySQL'   : 'mysqld'
}

resourcesList = ['CPU', 'MEM', 'SWAP', 'LOAD']

ITEM_TYPE = ['WMCOMPONENT', 'SERVICE', 'RESOURCE', 'DISK']

CMD_TYPE = ['START', 'SHUTDOWN', 'STATUS', 'RESTART', 'HELP', 'LIST']

OUT_TYPE = ['XML', 'TUI', 'JSON', 'CSV']

ITEM_WMCOMPONENT = 0
ITEM_SERVICE     = 1
ITEM_RESOURCE    = 2
ITEM_DISK        = 3

CMD_START    = 0
CMD_SHUTDOWN = 1
CMD_STATUS   = 2
CMD_RESTART  = 3
CMD_HELP     = 4
CMD_LIST     = 5

OUT_XML  = 0
OUT_TUI  = 1
OUT_JSON = 2
OUT_CSV  = 3

#################
### Main part ###
#################

valid = ['config=', 'wmcomponents=', 'services=', 'resources=', 'disks=', 'list=',
         'start', 'shutdown', 'restart', 'status', 'output=', 'help', 'debug']

try:
    opts, args = getopt.getopt(sys.argv[1:], "", valid)
except getopt.GetoptError as ex:
    print(error(str(ex) + help()))

configFile    = None
config        = None
command       = None
output        = None
specifiedList = {}
loadDefaults  = True
debugMode     = False
availableList = {}

if len(opts) == 0:
    print(help())
    sys.exit(0)

if opts[0][0] != "--config":
    error("First passed option must be '--config='!")

for opt, arg in opts:
    if opt == "--config":
        configFile = arg
    elif opt == "--start":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_START
    elif opt == "--shutdown":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_SHUTDOWN
    elif opt == "--restart":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_RESTART
    elif opt == "--status":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_STATUS
    elif opt == "--help":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_HELP
    elif opt == "--output":
        if output != None:
            error("'--output=' specified twise!" + help())
        output = checkOutput(arg)
        if output == None:
            error("Not support output format! Following output format are posible: 'TUI', 'XML', 'CSV', 'JSON'.")
    elif opt == "--debug":
        if debugMode == True:
            error("'--debug' was specified twise!" + help())
        debugMode = True
    elif opt == "--wmcomponents":
        if ITEM_WMCOMPONENT in specifiedList:
            error("'--wmcomponents=' specified twise!" + help())
        prepareList(ITEM_WMCOMPONENT, arg)
    elif opt == "--services":
        if ITEM_SERVICE in specifiedList:
            error("'--services=' specified twise!" + help())
        prepareList(ITEM_SERVICE, arg)
    elif opt == "--resources":
        if ITEM_RESOURCE in specifiedList:
            error("'--resources=' specified twise!" + help())
        prepareList(ITEM_RESOURCE, arg)
    elif opt == "--disks":
        if ITEM_DISK in specifiedList:
            error("'--disks=' specified twise!" + help())
        prepareList(ITEM_DISK, arg)
    elif opt == "--list":
        if len(availableList) != 0:
            error("'--list=' specified twise!" + help())
        availableList = prepareAvailableList(arg)
        command = CMD_LIST

# Checking curcial variables

if configFile == None:
    error("No configuration file set! Configuration file must be passed via '--config=' option.")
    
if command == None:
    error("No commmand specified! You must specify one of the following commands:\n\
          '--start', '--shutdown', '--status', '--restart', '--help'")

# Loading configuration

loadConfiguration()

# Load defaults if necessary
  
if output == None:
    output = OUT_TUI

if loadDefaults == True:
    specifiedList[ITEM_WMCOMPONENT] = config.listComponents_()
    specifiedList[ITEM_SERVICE] = list(servicesList)
    specifiedList[ITEM_RESOURCE] = resourcesList
    specifiedList[ITEM_DISK] = disks()

# Execute command

if command == CMD_START:
    start()
    sys.exit(0)
elif command == CMD_SHUTDOWN:
    shutdown()
    sys.exit(0)
elif command == CMD_STATUS:
    status()
    sys.exit(0)
elif command == CMD_RESTART:
    restart()
    sys.exit(0)
elif command == CMD_HELP:
    print(help(short=False))
    sys.exit(0)
elif command == CMD_LIST:
    print(available())
    sys.exit(0)