""" NOTA BENE: This agent should NOT be run alone. Instead, it serves as a base class for extensions.

    The TaskManagerAgentBase is the base class to submit tasks to external systems,
    monitor and update the tasks and file status in the transformation DB.

    This agent is extended in WorkflowTaskAgent and RequestTaskAgent.
    In case you want to further extend it you are required to follow the note on the
    initialize method and on the _getClients method.
"""

__RCSID__ = "$Id$"

import time
import datetime
from Queue import Queue

from DIRAC import S_OK

from DIRAC.Core.Base.AgentModule import AgentModule
from DIRAC.Core.Security.ProxyInfo import getProxyInfo
from DIRAC.Core.Utilities.ThreadPool import ThreadPool
from DIRAC.Core.Utilities.List import breakListIntoChunks
from DIRAC.Core.Utilities.Dictionaries import breakDictionaryIntoChunks
from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getDNForUsername, getUsernameForDN
from DIRAC.FrameworkSystem.Client.MonitoringClient import gMonitor
from DIRAC.TransformationSystem.Client.FileReport import FileReport
from DIRAC.TransformationSystem.Client.TaskManager import WorkflowTasks
from DIRAC.TransformationSystem.Client.TransformationClient import TransformationClient
from DIRAC.TransformationSystem.Agent.TransformationAgentsUtilities import TransformationAgentsUtilities
from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient

AGENT_NAME = 'Transformation/TaskManagerAgentBase'


class TaskManagerAgentBase(AgentModule, TransformationAgentsUtilities):
  """ To be extended. Please look at WorkflowTaskAgent and RequestTaskAgent.
  """

  def __init__(self, *args, **kwargs):
    """ c'tor

        Always call this in the extension agent
    """
    AgentModule.__init__(self, *args, **kwargs)
    TransformationAgentsUtilities.__init__(self)

    self.transClient = None
    self.jobManagerClient = None
    self.transType = []

    self.tasksPerLoop = 50
    self.maxParametricJobs = 20  # will be updated in execute()

    # credentials
    self.shifterProxy = None
    self.credentials = None
    self.credTuple = (None, None, None)

    self.pluginLocation = ''
    self.bulkSubmissionFlag = False

    # for the threading
    self.transQueue = Queue()
    self.transInQueue = []
    self.transInThread = {}

  #############################################################################

  def initialize(self):
    """ Agent initialization.

        The extensions MUST provide in the initialize method the following data members:
        - TransformationClient objects
 (self.transClient),
        - set the shifterProxy if different from the default one set here ('ProductionManager')
        - list of transformation types to be looked (self.transType)
    """

    gMonitor.registerActivity("SubmittedTasks", "Automatically submitted tasks
", "Transformation Monitoring", "Tasks",
                              gMonitor.OP_ACUM)

    self.pluginLocation = self.am_getOption('PluginLocation', 'DIRAC.TransformationSystem.Client.TaskManagerPlugin')

    # Default clients
    self.transClient = TransformationClient()
    self.jobManagerClient = JobManagerClient()

    # Bulk submission flag
    self.bulkSubmissionFlag = self.am_getOption('BulkSubmission', self.bulkSubmissionFlag)

    # Shifter credentials to use, could replace the use of shifterProxy eventually
    self.shifterProxy = self.am_getOption('shifterProxy', self.shifterProxy)
    self.credentials = self.am_getOption('ShifterCredentials', self.credentials)
    resCred = self.__getCredentials()
    if not resCred['OK']:
      return resCred
    # setting up the threading
    maxNumberOfThreads = self.am_getOption('maxNumberOfThreads', 15)
    threadPool = ThreadPool(maxNumberOfThreads, maxNumberOfThreads)
    self.log.verbose("Multithreaded with %d threads" % maxNumberOfThreads)

    for i in xrange(maxNumberOfThreads):
      threadPool.generateJobAndQueueIt(self._execute, [i])

    return S_OK()

  def finalize(self):
    """ graceful finalization
    """
    if self.transInQueue:
      self._logInfo("Wait for threads to get empty before terminating the agent (%d tasks)" %
                    len(self.transInThread))
      self.transInQueue = []
      while self.transInThread:
        time.sleep(2)
      self.log.info("Threads are empty, terminating the agent...")
    return S_OK()

  #############################################################################

  def execute(self):
    """ The TaskManagerBase execution method is just filling the Queues of transformations that need to be processed
    """

    operationsOnTransformationDict = {}
    owner, ownerGroup, ownerDN = None, None, None
    # getting the credentials for submission
    resProxy = getProxyInfo(proxy=False, disableVOMS=False)
    if resProxy['OK']:  # there is a shifterProxy
      proxyInfo = resProxy['Value']
      owner = proxyInfo['username']
      ownerGroup = proxyInfo['group']
      ownerDN = proxyInfo['identity']
      self.log.info("ShifterProxy: Tasks will be submitted with the credentials %s:%s" % (owner, ownerGroup))
    elif self.credentials:
      owner, ownerGroup, ownerDN = self.credTuple
    else:
      self.log.info("Using per Transformation Credentials!")

    # Determine whether the task status is to be monitored and updated
    enableTaskMonitor = self.am_getOption('MonitorTasks', '')
    if not enableTaskMonitor:
      self.log.verbose("Monitoring of tasks is disabled. To enable it, create the 'MonitorTasks' option")
    else:
      # Get the transformations for which the tasks have to be updated
      status = self.am_getOption('UpdateTasksTransformationStatus',
                                 self.am_getOption('UpdateTasksStatus', ['Active', 'Completing', 'Stopped']))
      transformations = self._selectTransformations(transType=self.transType, status=status, agentType=[])
      if not transformations['OK']:
        self.log.warn("Could not select transformations:", transformations['Message'])
      else:
        self._addOperationForTransformations(operationsOnTransformationDict, 'updateTaskStatus', transformations,
                                             owner=owner, ownerGroup=ownerGroup, ownerDN=ownerDN)

    # Determine whether the task files status is to be monitored and updated
    enableFileMonitor = self.am_getOption('MonitorFiles', '')
    if not enableFileMonitor:
      self.log.verbose("Monitoring of files is disabled. To enable it, create the 'MonitorFiles' option")
    else:
      # Get the transformations for which the files have to be updated
      status = self.am_getOption('UpdateFilesTransformationStatus',
                                 self.am_getOption('UpdateFilesStatus', ['Active', 'Completing', 'Stopped']))
      transformations = self._selectTransformations(transType=self.transType, status=status, agentType=[])
      if not transformations['OK']:
        self.log.warn("Could not select transformations:", transformations['Message'])
      else:
        self._addOperationForTransformations(operationsOnTransformationDict, 'updateFileStatus', transformations,
                                             owner=owner, ownerGroup=ownerGroup, ownerDN=ownerDN)

    # Determine whether the checking of reserved tasks is to be performed
    enableCheckReserved = self.am_getOption('CheckReserved', '')
    if not enableCheckReserved:
      self.log.verbose("Checking of reserved tasks is disabled. To enable it, create the 'CheckReserved' option")
    else:
      # Get the transformations for which the check of reserved tasks have to be performed
      status = self.am_getOption('CheckReservedTransformationStatus',
                                 self.am_getOption('CheckReservedStatus', ['Active', 'Completing', 'Stopped']))
      transformations = self._selectTransformations(transType=self.transType, status=status, agentType=[])
      if not transformations['OK']:
        self.log.warn("Could not select transformations:", transformations['Message'])
      else:
        self._addOperationForTransformations(o