#!/usr/bin/env python
"""
wmagent-workqueue

Utility script for manipulating workqueue.
"""

import os
import sys
from argparse import ArgumentParser

from WMCore.Configuration import loadConfigurationFile
from WMCore.WorkQueue.WorkQueueUtils import queueFromConfig


def createOptionParser():
    """
    _createOptionParser_

    Create an option parser that knows about all the options for manipulating
    and accessing resource control.
    """
    myOptParser = ArgumentParser()
    myOptParser.add_argument("--status", dest="status",
                             default=False, action="store_true",
                             help="Print out element's status.")
    myOptParser.add_argument("--locations", dest="updateLocations",
                             default=False, action="store_true",
                             help="Update data location's.")
    myOptParser.add_argument("--priority", dest="setPriority",
                             default=False, action="store_true",
                             help="Change priority of elements.")
    myOptParser.add_argument("--queue", dest="queueWork",
                             default=False, action="store_true",
                             help="Queue work directly to the WorkQueue.")
    myOptParser.add_argument("--reset", dest="resetWork",
                             default=False, action="store_true",
                             help="Reset elements status.")
    myOptParser.add_argument("--delete", dest="deleteWork",
                             default=False, action="store_true",
                             help="Delete elements.")
    myOptParser.add_argument("--cancel", dest="cancelWork",
                             default=False, action="store_true",
                             help="Cancel request")
    myOptParser.add_argument("--process", dest="process",
                             default=False, action="store_true",
                             help="Process inbound work.")
    myOptParser.add_argument("--clean", dest="clean",
                             default=False, action="store_true",
                             help="Clean queue.")
    myOptParser.add_argument("-i", "--interactive", dest="interactive",
                             default=False, action="store_true",
                             help="Interactive console.")
    myOptParser.add_argument("--reqmgr", dest="requestmanager",
                             default=False, action="store_true",
                             help="Synchronize with ReqMgr2")
    myOptParser.add_argument("--deleteCascade", dest="deleteCascade",
                             default=False, action="store_true",
                             help="delete LQ GQ.")
    myOptParser.add_argument("--resetGQ", dest="resetGQ",
                             default=False, action="store_true",
                             help="Reset GQ elements status.")

    group = myOptParser.add_argument_group("WorkQueue elements selection.",
                                           "Use these to select which elements to perform  "
                                           "the desired operation on.")
    group.add_argument("--workflow", dest="workflow", help="workflow name")
    group.add_argument("--id", dest="id", help="element id")

    group = myOptParser.add_argument_group("General options.",
                                           "General options to direct bahviour")
    group.add_argument("--config", dest="config", help="wm agent config",
                       default=os.environ.get("WMAGENT_CONFIG", None))

    return myOptParser


def createWorkQueue(config):
    """Create a workqueue from wmagent config"""
    # if config has a db sction instantiate a dbi
    if hasattr(config, "CoreDatabase"):
        from WMCore.WMInit import WMInit
        wmInit = WMInit()
        (dialect, junk) = config.CoreDatabase.connectUrl.split(":", 1)
        socket = getattr(config.CoreDatabase, "socket", None)
        wmInit.setDatabaseConnection(dbConfig=config.CoreDatabase.connectUrl,
                                     dialect=dialect,
                                     socketLoc=socket)
    return queueFromConfig(config)


def _deleteWorkflowFromWQ(wqDB, workflow):
    """
    _deleteWorkflowFromCouch_

    If we are asked to delete the workflow from couch, delete it
    to clear up some space.

    Load the document IDs and revisions out of couch by workflowName,
    then order a delete on them.
    """

    elements = wqDB.loadView('WorkQueue', 'elementsByWorkflow',
                             {'key': workflow, 'reduce': False})['rows']
    for j in elements:
        id = j['id']
        wqDB.delete_doc(id=id)
    try:
        wqDB.delete_doc(workflow)
    except Exception as ex:
        print("%s : %s" % (workflow, str(ex)))
    return len(elements)


def deleteWorkflowFromWQ(workqueue, workflowName):
    print("LQ deleted: %s" % _deleteWorkflowFromWQ(workqueue.backend.db, workflowName))
    print("LQ inbox deleted: %s" % _deleteWorkflowFromWQ(workqueue.backend.inbox, workflowName))
    print("GQ deleted: %s" % _deleteWorkflowFromWQ(workqueue.parent_queue.db, workflowName))
    print("GQ inbox deleted: %s" % _deleteWorkflowFromWQ(workqueue.parent_queue.inbox, workflowName))


def main():
    myOptParser = createOptionParser()
    (options, args) = myOptParser.parse_known_args()

    if not options.config or not os.path.exists(options.config):
        msg = "No Config file provided\n"
        msg += "provide one with the --config option"
        print(msg)
        sys.exit(1)

    cfgObject = loadConfigurationFile(options.config)

    workqueue = createWorkQueue(cfgObject)

    if options.status:
        print(workqueue.status())
    elif options.updateLocations:
        workqueue.updateLocationInfo()
        # print out location list
    elif options.setPriority:
        priority = args[0]
        workqueue.setPriority(priority, options.workflow)
    elif options.queueWork:
        wf = args[0]
        workqueue.queueWork(wf)
    elif options.resetWork:
        workqueue.resetWork(options.id)
    elif options.process:
        workqueue.processInboundWork()
    elif options.deleteWork:
        workqueue.deleteWork(options.id)
    elif options.clean:
        workqueue.performQueueCleanupActions()
    elif options.interactive:
        msg = 'Workqueue available in "workqueue" variable'
        import code
        code.interact(banner=msg, local=locals())
    elif options.requestmanager:
        from WMCore.WorkQueue.WorkQueueReqMgrInterface import WorkQueueReqMgrInterface
        reqConfig = getattr(cfgObject.WorkQueueManager, 'reqMgrConfig', {})
        reqMgrInt = WorkQueueReqMgrInterface(**reqConfig)
        reqMgrInt(workqueue)
    elif options.cancelWork:
        workqueue.cancelWork(WorkflowName=options.workflow)
    elif options.deleteCascade:
        deleteWorkflowFromWQ(workqueue, workflowName=options.workflow)
    elif options.resetGQ:
        ids = options.id.split(",")
        workqueue.parent_queue.updateElements(*ids, Status='Available',
                                              ChildQueueUrl=None, WMBSUrl=None)


if __name__ == '__main__':
    main()
