#!/usr/bin/env python

"""
Controls the daemons that represent the components.

The following is a short a set of commands you can give to start
and control the daemons:

# generates a config file based on the root dirs (it looks for DefaultConfig.py files):
# the component sections need a namespace attribute to the class can be loaded and daemonized.
wmcore-new-config --roots=src/python/WMComponent/ErrorHandler/,src/python/WMComponent/DBS3Buffer/ --output=WMCoreConfig.py

#edit the generated config file to fit your needs.

# loads default tables and tables from a specific module. The backend is automatically selected based on the config file.
wmcore-db-init --config=WMCoreConfig.py --create --modules=WMComponent.DBS3Buffer.Database 

# start the daemons from components defined in the config file:
wmcoreD --start --config=WMCoreConfig.py

# prints status of components
wmcoreD --status --config=WMCoreConfig.py

# restarts components
wmcoreD --restart --config=WMCoreConfig.py

# shutdown the daemons:
wmcoreD --shutdown --config=WMCoreConfig.py --cleanup-all

"""

from builtins import str
__revision__ = "$Id: wmcoreD,v 1.16 2010/05/26 21:09:35 sfoulkes Exp $"
__version__ = "$Revision: 1.16 $"
__author__ = "fvlingen@caltech.edu"

import getopt
import os
import subprocess
import sys
import time
import json
import psutil

from Utils.Utilities import extractFromXML
from Utils.ProcFS import processStatus
from WMCore.Agent.Daemon.Details import Details
from WMCore.Configuration import loadConfigurationFile
from WMCore.WMFactory import WMFactory
from WMCore.WMInit import WMInit

def usage():

    msg = """
Usage: wmcoreD <--start|--shutdown|--status> --config <opts>

You must provide either --start OR --shutdown OR --status

You must either set the WMAGENT_CONFIG environment variable or specify the config file with
--config 

--start starts up the components
--shutdown shutsdown the components
--status prints the status of the components
--restart restarts components 

options:
--cleanup-logs purges logs
--cleanup-all purges component dirs

"""
    print(msg)


valid = ['config=', 'start', 'shutdown', 'status', 'restart',
         'components=', 'cleanup-logs', 'cleanup-all']

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

config = None
command = None
doLogCleanup = False
doDirCleanup = False
componentsList = None


for opt, arg in opts:
    if opt == "--config":
        config = arg
    if opt == "--start":
        if command != None:
            print("Command specified twice:\n")
            usage()
            sys.exit(1)
        command = "start"
    if opt == "--shutdown":
        if command != None:
            print("Command specified twice:\n")
            usage()
            sys.exit(1)
        command = "shutdown"
    if opt == "--status":
        if command != None:
            print("Command specified twice:\n")
            usage()
            sys.exit(1)
        command = "status"
    if opt == "--restart":
        if command != None:
            print("Command specified twice:\n")
            usage()
            sys.exit(1)
        command = "restart"
    if opt == "--cleanup-logs":
        doLogCleanup = True
    if opt == "--cleanup-all":
        doDirCleanup = True
    if opt == "--components":
        compList = arg.split(',')
        componentsList = []
        for item in compList:
            if item.strip == "":
                continue
            componentsList.append(item)
           
if command == None:
    msg = "No command specified\n"
    print(msg)
    usage()
    sys.exit(0)
 
if config == None:            
    config = os.environ.get("WMAGENT_CONFIG", None)

    if config == None:
        msg = "No Config file provided\n"
        msg += "provide one with the --config option"
        print(msg)
        usage()
        sys.exit(1)

if not os.path.exists(config):
    print("Can't find config: %s" % config)
    sys.exit(1)

# load the config file here.
cfgObject = loadConfigurationFile(config)
#workingDir = os.path.expandvars(workingDir)

if componentsList != None:
    msg = "Components List Specified:\n"
    msg += str(componentsList).replace('\'', '')
    print(msg)
    components = componentsList
else:    
    components = cfgObject.listComponents_() + cfgObject.listWebapps_()

def connectionTest(configFile):
    """
    _connectionTest_

    Create a DB Connection instance to test the connection specified
    in the config file.

    """
    config = loadConfigurationFile(configFile)
    wmInit = WMInit()    
    print("Checking default database connection...", end=' ')

    if not hasattr(config, "CoreDatabase"):
        print("skipped.")
        return

    dialect, _ = config.CoreDatabase.connectUrl.split(":", 1)
    socket = getattr(config.CoreDatabase, "socket", None)

    try:
        wmInit.setDatabaseConnection(dbConfig = config.CoreDatabase.connectUrl,
                                     dialect = dialect,
                                     socketLoc = socket)
    except Exception as ex:
        msg = "Unable to make connection to using \n"
        msg += "parameters provided in %s\n" % config.CoreDatabase.connectUrl 
        msg += str(ex)
        print(msg)
        raise ex

    print("ok.")
    return

def startup(configFile):
    """
    _startup_

    Start up the component daemons

    """
    print('Starting components: '+str(components))
    config = loadConfigurationFile(configFile)
        
    for component in components:
        print('Starting : '+component)
        if component in config.listWebapps_():
            from WMCore.WebTools.Root import Root            
            webtoolsRoot = Root(config, webApp = component)
            webtoolsRoot.startDaemon(keepParent = True, compName = component)            
        else:
            factory = WMFactory('componentFactory')
            try:
                namespace = config.component_(component).namespace
            except AttributeError:
                print ("Failed to start component: Could not find component named %s in config" % component)
                print ("Aborting")
                return
            componentObject = factory.loadObject(classname = namespace, args = config)
            componentObject.startDaemon(keepParent = True)
            
        print('Waiting 1 seconds, to ensure daemon file is created')
        time.sleep(1)
        compDir = config.section_(component).componentDir
        compDir = os.path.expandvars(compDir)
        daemonXML = os.path.join( compDir, "Daemon.xml")
        if os.path.exists(daemonXML):
            daemon = Details(daemonXML)
            if not daemon.isAlive():         
                print("Error: Component %s Did not start properly..." % component)
                print("Check component log to see why")
                sys.exit(1)
        else:
            print('Path for daemon file does not exist!')
            sys.exit(1)
        # write into component area process status information
        cpath = os.path.join(compDir, "threads.json")
        if os.path.exists(cpath):
            os.remove(cpath)
        cpid = extractFromXML(daemonXML, "ProcessID")
        with open(cpath, 'w', encoding="utf-8") as istream:
            procStatus = processStatus(cpid)
            istream.write(json.dumps(procStatus))
            print("Component %s started with %s threads, see %s\n" % (component, len(procStatus), cpath))

    return
    
def shutdown(configFile):
    """
    _shutdown_

    Shutdown the component daemons

    If cleanup-logs option is specified, wipe out the component logs
    If cleanup-all option is specified, wipe out all component dir
    content and purge the ProdAgentDB

    """
    print('Stopping components: '+str(components))
    config = loadConfigurationFile(configFile)

    for component in components:
        print('Stopping: '+component)
        try:
            compDir = config.section_(component).componentDir
        except AttributeError:
            print ("Failed to shutdown component: Could not find component named %s in config" % component)
            print ("Aborting")
            return
        compDir = os.path.expandvars(compDir)
        daemonXml = os.path.join(compDir, "Daemon.xml")
        if not os.path.exists(daemonXml):
            print("Cannot find Daemon.xml for component:", component)
            print("Unable to shut it down")
        else:
            daemon = Details(daemonXml)
            if not daemon.isAlive():
                print("Component %s with process id %s is not running" % (
                    component, daemon['ProcessID'],
                    ))
                daemon.removeAndBackupDaemonFile()
            else:
                daemon.killWithPrejudice()
        # remove component threads.json file
        cpath = os.path.join(compDir, "threads.json")
        if os.path.exists(cpath):
            os.remove(cpath)
        if doLogCleanup:
            #  //
            # // Log Cleanup
            #//
            msg = "Removing %s/ComponentLog" % compDir
            print(msg)
            try:
                os.remove("%s/ComponentLog" % compDir)
            except Exception as ex:
                msg = "Unable to cleanup Component Log: "
                msg += "%s/ComponentLog\n" % compDir
                msg += str(ex)
                
        if doDirCleanup:
            #  //
            # // Cleanout everything in ComponentDir
            #//  for this component
            print("Removing %s\n" % compDir)
            exitCode = subprocess.call(["rm", "-rf", "%s" % compDir])
            if exitCode:
                msg = "Failed to clean up dir: %s\n" % compDir
                msg += f"with exit code {exitCode}"
                print(msg)

    return


def status(configFile):
    """
    _status_

    Print status of all components in config file

    """
    print('Status components: '+str(components))
    config = loadConfigurationFile(configFile)

    for component in components:
        try:
            compDir = config.section_(component).componentDir
        except AttributeError:
            print ("Failed to check component: Could not find component named %s in config" % component)
            print ("Aborting")
            return
        compDir = config.section_(component).componentDir
        compDir = os.path.expandvars(compDir)
        checkProcessThreads(component, compDir)

    sys.exit(0)

def checkProcessThreads(component, compDir):
    """
    Helper function to check process and its threads for their statuses
    :param component: component name
    :return: prints status of the component process and its threads
    """
    # check if component daemon exists
    daemonXml = os.path.join(compDir, "Daemon.xml")
    if not os.path.exists(daemonXml):
        print("Component:%s Not Running" % component)
        return
    pid = extractFromXML(daemonXml, "ProcessID")

    jsonFile = os.path.join(compDir, "threads.json")
    with open(jsonFile, "r", encoding="utf-8") as istream:
        data = json.load(istream)

    # Extract process and its threads
    threadPids = []

    for entry in data:
        if str(entry["pid"]) == str(pid) and entry["type"] == "process":
            continue
        elif entry["type"] == "thread":
            threadPids.append(entry["pid"])

    # Check if process is running
    processRunning = psutil.pid_exists(int(pid))
    if not processRunning:
        msg = f"Component:{component} with PID={pid} is no longer available on OS"
        print(msg)
        return

    # Check if threads are running
    runningThreads = []
    orphanThreads = []
    process = psutil.Process(int(pid))
    for thread in process.threads():
        if str(thread.id) == str(pid):
            continue  # skip pid thread
        if str(thread.id) in threadPids:
            runningThreads.append(thread.id)
        else:
            orphanThreads.append(thread.id)

    # Output result
    status = "running" if processRunning else "not-running"
    msg = f"Component:{component} {pid} {status} with threads {runningThreads}"
    if status == "running":
        if len(orphanThreads) > 0:
            status = "partially-running" if processRunning else "not-running"
            msg = f"Component:{component} {pid} {status} with threads {runningThreads}, lost {orphanThreads} threads"
    else:
        msg = f"Component:{component} {pid} {status}"
    print(msg)

def restart(config):
    """
    _restart_

    do a shutdown and startup again

    """
    shutdown(config)
    startup(config)
    return

if command == "start":
    connectionTest(config)
    startup(config)
    sys.exit(0)

elif command == "shutdown":
    connectionTest(config)
    shutdown(config)
    sys.exit(0)
elif command == "status":
    connectionTest(config)
    status(config)
    sys.exit(0)
    
elif command == "restart":
    connectionTest(config)
    restart(config)
    sys.exit(0)

