#!/usr/bin/env python
"""
wmagent-resource-control

Utility script for manipulating resource control.
"""
from __future__ import print_function

import os
import sys
from argparse import ArgumentParser

from WMCore.Configuration import loadConfigurationFile
from WMCore.ResourceControl.ResourceControl import ResourceControl
from WMCore.Services.CRIC.CRIC import CRIC
from WMCore.WMInit import connectToDB


def getTaskTypes(tier0=False):
    """
    _getTaskTypes_

    Get the list of task types we are currently using with
    priorities.
    """
    taskTypes = ["Merge",
                 "Cleanup",
                 "Harvesting",
                 "LogCollect",
                 "Skim",
                 "Production",
                 "Processing"]

    if tier0:
        taskTypes.append("Repack")
        taskTypes.append("Express")

    return taskTypes


def createOptionParser():
    """
    _createOptionParser_

    Create an option parser that knows about all the options for manipulating
    and accessing resource control.
    """
    myOptParser = ArgumentParser()
    myOptParser.add_argument("-p", "--thresholds", dest="printThresholds",
                             default=False, action="store_true",
                             help="Print out all known thresholds and site information.")
    myOptParser.add_argument("--add-Test", dest="addTest", default=False, action="store_true",
                             help="Add a fake test site to resource control.")
    myOptParser.add_argument("--add-T1s", dest="addT1s", default=False, action="store_true",
                             help="Add all of the CMS T1 sites to resource control.")
    myOptParser.add_argument("--add-T2s", dest="addT2s", default=False, action="store_true",
                             help="Add all of the CMS T2 sites to resource control.")
    myOptParser.add_argument("--add-T3s", dest="addT3s", default=False, action="store_true",
                             help="Add all of the CMS T3 sites to resource control.")
    myOptParser.add_argument("--add-all-sites", dest="addAllSites", default=False,
                             action="store_true",
                             help="Add all of the CMS sites to resource control.")
    myOptParser.add_argument("--add-one-site", dest="addOneSite",
                             help="Specify site name you want to add")
    myOptParser.add_argument("--site-name", dest="siteName",
                             help="Specify the unique name of the location")
    myOptParser.add_argument("--pending-slots", dest="pendingSlots",
                             help="Specify the maximum number of pending slots to use at the site or task type")
    myOptParser.add_argument("--running-slots", dest="runningSlots",
                             help="Specify the maximum number of running slots to use at the site or task type")
    myOptParser.add_argument("--apply-to-all-tasks", dest="allTasks", default=False, action="store_true",
                             help="Apply site running/pending threshold to all task types")
    myOptParser.add_argument("--ce-name", dest="ceName",
                             help="Specify the CEName for the site")
    myOptParser.add_argument("--pnn", dest="pnn",
                             help="Specify the PNN for the site")
    myOptParser.add_argument("--cms-name", dest="cmsName",
                             help="Specify the name of the site in CRIC")
    myOptParser.add_argument("--task-type", dest="taskType",
                             help="Specify the name of the task to add/modify")
    myOptParser.add_argument("--priority", dest="priority", default=None,
                             help="Specify the priority of the task across all sites (higher is better)")
    myOptParser.add_argument("--plugin", dest="plugin",
                             help="Specify submit plugin to use for specific site")
    myOptParser.add_argument("--drain", dest="state", action="store_const",
                             const='Draining', help="Drain the site.")
    myOptParser.add_argument("--abort", dest="state", action="store_const",
                             const='Aborted', help="Abort jobs at the site.")
    myOptParser.add_argument("--down", dest="state", action="store_const",
                             const='Down', help="Stop submission to the site.")
    myOptParser.add_argument("--normal", dest="state", action="store_const",
                             const='Normal', help="Normal operation for the site.")
    myOptParser.add_argument("--add-disk", dest="addDisk", default=True, action="store_false",
                             help="Use it to explicitly add the Disk endpoints together with the other T1s")

    return myOptParser


def getCMSSiteInfo(pattern):
    """
    _getCMSSiteInfo_

    Query CRIC for the site and SE names matching the pattern.  Return a
    dictionary keyed by site name.
    """
    cricObj = CRIC()

    mapping = cricObj.PSNtoPNNMap(pattern)
    print("Retrieved %i maps from %s" % (len(mapping), cricObj['endpoint']))
    return mapping


def addSites(resourceControl, allSites, ceName, plugin, pendingSlots, runningSlots,
             tier0Option=False, cmsName=None, state=None):
    """
    _addSites_

    Add the given sites to resource control and add tasks as well.
    """
    if pendingSlots is None:
        pendingSlots = 25
    else:
        pendingSlots = pendingSlots
    if runningSlots is None:
        runningSlots = 25
    else:
        runningSlots = runningSlots

    for siteName in allSites.keys():
        updateSiteInfo(resourceControl, siteName, pendingSlots, runningSlots,
                       ceName or siteName, allSites[siteName], plugin, cmsName or siteName)

        tasks = getTaskTypes(tier0Option)
        for task in tasks:
            updateThresholdInfo(resourceControl, siteName, task, runningSlots, pendingSlots)

    if state:
        for siteName in allSites.keys():
            setSiteState(resourceControl, siteName, state)

    return


def addPNNs(resourceControl, pattern):
    """
    _addPNNs_

    Add the given sites to resource control and add tasks as well.
    """
    cricObj = CRIC()

    pnns = cricObj.getAllPhEDExNodeNames(pattern, excludeBuffer=True)
    print("Retrieved %i PNNs from %s" % (len(pnns), cricObj['endpoint']))
    resourceControl.insertPNNs(pnns)

    return


def printThresholds(resourceControl, desiredSiteName):
    """
    _printThresholds_

    Print out the current resource control thresholds.
    """
    print("Thresholds and current status for all sites:\n")

    thresholds = resourceControl.listThresholdsForSubmit()

    for siteName in thresholds.keys():
        if desiredSiteName and desiredSiteName != siteName:
            continue
        siteThresholds = thresholds[siteName]

        if len(siteThresholds) < 1:
            # No thresholds
            print("No thresholds for site %s" % siteName)
            continue

        pendingSlots = siteThresholds["total_pending_slots"]
        runningSlots = siteThresholds["total_running_slots"]
        pendingJobs = siteThresholds["total_pending_jobs"]
        runningJobs = siteThresholds["total_running_jobs"]
        state = siteThresholds["state"]
        stateMsg = ", Site is %s" % state
        print("%s - %d running, %d pending, %d running slots total, %d pending slots total%s:" % (
            siteName, runningJobs, pendingJobs, runningSlots, pendingSlots, stateMsg))

        for task, thres in siteThresholds['thresholds'].iteritems():
            print("  %s - %d running, %d pending, %d max running, %d max pending, priority %s" %
                  (task, thres["task_running_jobs"], thres['task_pending_jobs'], thres["max_slots"],
                   thres["pending_slots"], str(thres.get("priority", 1))))
            if thres["task_running_jobs"] > 0:
                taskWorkloads = resourceControl.listWorkloadsForTaskSite(task, siteName)
                for t in taskWorkloads:
                    print("    %d - %s" % (t["running"], t["task"]))

        print("")


def updateSiteInfo(resourceControl, siteName, pendingSlots=None,
                   runningSlots=None, ceName=None, pnn=None,
                   plugin=None, cmsName=None, allTasks=False):
    """
    _updateSiteInfo_

    Add a site to the resource control database if it doesn't exist.  Update
    information about sites in the database if it already exists.
    """
    if resourceControl.listSiteInfo(siteName) is None:
        if pendingSlots is None or runningSlots is None or ceName is None or pnn is None or plugin is None:
            print("You must specify the number of pending or running slots, the SE name,")
            print("the CE name and a plugin when adding a site.")
            print("Type --help for more information.")
            sys.exit(1)

        print("Adding %s to the resource control db..." % siteName)
        if isinstance(pnn, (list, set, tuple)):
            for singlePNN in pnn:
                resourceControl.insertSite(siteName=siteName, pendingSlots=int(pendingSlots),
                                           runningSlots=int(runningSlots),
                                           pnn=singlePNN, ceName=ceName,
                                           plugin=plugin, cmsName=cmsName)
        else:
            resourceControl.insertSite(siteName=siteName, pendingSlots=int(pendingSlots),
                                       runningSlots=int(runningSlots),
                                       pnn=pnn, ceName=ceName,
                                       plugin=plugin, cmsName=cmsName)

        return

    if pendingSlots is not None or runningSlots is not None:
        resourceControl.setJobSlotsForSite(siteName=siteName,
                                           pendingJobSlots=pendingSlots,
                                           runningJobSlots=runningSlots)
        if allTasks:
            tasksWithPriorities = getTaskTypes()
            for task in tasksWithPriorities:
                updateThresholdInfo(resourceControl, siteName, task, runningSlots, pendingSlots)

    if pnn is not None:
        print("It's not possible to change a site's SE name after the site has")
        print("been added to the database.")
        sys.exit(1)

    if ceName is not None:
        print("It's not possible to change a site's CE name after the site has")
        print("been added to the database.")
        sys.exit(1)

    return


def updateThresholdInfo(resourceControl, siteName, taskType, maxSlots=None, pendingSlots=None):
    """
    _updateThresholdInfo_

    Add or update a task threshold in the database.
    """
    if resourceControl.listSiteInfo(siteName) is None:
        print("You must add the site to the database before you can add")
        print("thresholds.")
        print("Type --help for more information.")
        sys.exit(1)

    if maxSlots is None and pendingSlots is None:
        print("You must provide either pending or running slots for the task.")
        print("Type --help for more information.")
        sys.exit(1)

    resourceControl.insertThreshold(siteName=siteName, taskType=taskType,
                                    maxSlots=maxSlots, pendingSlots=pendingSlots)
    return


def updateTaskPriority(resourceControl, taskType, priority):
    """
    _updateTaskPriority_

    Update the priority of a task in the database, if task
    doesn't exist then add it.
    """
    if taskType is None:
        print("Sub-task type must be specified when changing priority.")
        myOptParser.print_help()
        sys.exit(1)
    resourceControl.changeTaskPriority(taskType, priority)
    return


def setSiteState(resourceControl, siteName, state):
    """
    _setSiteState_

    Set an specific state to a site
    """
    if resourceControl.listSiteInfo(siteName) is None:
        print("You must add the site to the database before you can change its state")
        print("Type --help for more information.")
        sys.exit(1)

    return resourceControl.changeSiteState(siteName, state)


myOptParser = createOptionParser()
options = myOptParser.parse_args()

connectToDB()
wmConfig = loadConfigurationFile(os.environ['WMAGENT_CONFIG'])
myResourceControl = ResourceControl(config=wmConfig)

if options.printThresholds:
    printThresholds(myResourceControl, options.siteName)
    sys.exit(0)
else:
    if options.priority:
        updateTaskPriority(myResourceControl, options.taskType, options.priority)
    elif options.addTest:
        testSite = {'CERN': ['T0_CH_CERN_Disk', 'T2_CH_CERN']}
        addSites(myResourceControl, testSite, "CERN", "SimpleCondorPlugin",
                 500, 1, state=options.state)
    elif options.addT1s:
        pattern = "T1.*"
        t1Sites = getCMSSiteInfo(pattern)
        addSites(myResourceControl, t1Sites, options.ceName, options.plugin,
                 options.pendingSlots, options.runningSlots, state=options.state)
        addPNNs(myResourceControl, pattern=pattern)
    elif options.addT2s:
        pattern = "T2.*"
        t2Sites = getCMSSiteInfo(pattern)
        addSites(myResourceControl, t2Sites, options.ceName, options.plugin,
                 options.pendingSlots, options.runningSlots, state=options.state)
        addPNNs(myResourceControl, pattern=pattern)
    elif options.addT3s:
        pattern = "T3.*"
        t3Sites = getCMSSiteInfo(pattern)
        addSites(myResourceControl, t3Sites, options.ceName, options.plugin,
                 options.pendingSlots, options.runningSlots, state=options.state)
        addPNNs(myResourceControl, pattern=pattern)
    elif options.addAllSites:
        pattern = ".*"
        allSites = getCMSSiteInfo(pattern)
        addSites(myResourceControl, allSites, options.ceName, options.plugin,
                 options.pendingSlots, options.runningSlots, state=options.state)
        addPNNs(myResourceControl, pattern=pattern)
    elif options.addOneSite:
        oneSite = getCMSSiteInfo(options.addOneSite)
        addSites(myResourceControl, oneSite, options.ceName, options.plugin,
                 options.pendingSlots, options.runningSlots, state=options.state)
    elif options.siteName is None:
        print("You must specify a site name.")
        myOptParser.print_help()
        sys.exit(1)
    elif options.state is not None:
        setSiteState(myResourceControl, options.siteName, options.state)
        sys.exit(0)
    elif options.taskType is None:
        pnn = options.pnn
        if options.pnn:
            pnn = options.pnn.strip().split(',')
        updateSiteInfo(myResourceControl, options.siteName, options.pendingSlots,
                       options.runningSlots, options.ceName, pnn,
                       options.plugin, options.cmsName, options.allTasks)
        sys.exit(0)
    else:
        taskType = options.taskType.strip().split(',')
        updateThresholdInfo(myResourceControl, options.siteName,
                            taskType, options.runningSlots,
                            options.pendingSlots)
        sys.exit(0)
