########################################################################
# $HeadURL$
# File :    AgentModule.py
# Author :  Adria Casajus
########################################################################
"""
  Base class for all agent modules
"""
__RCSID__ = "$Id$"

import os
import threading
import types
import time
import DIRAC
from DIRAC import S_OK, S_ERROR, gConfig, gLogger, gMonitor, rootPath
from DIRAC.ConfigurationSystem.Client import PathFinder
from DIRAC.FrameworkSystem.Client.MonitoringClient import MonitoringClient
from DIRAC.Core.Utilities.Shifter import setupShifterProxyInEnv
from DIRAC.Core.Utilities.ReturnValues import isReturnStructure
from DIRAC.Core.Utilities import Time, MemStat

def _checkDir( path ):
  try:
    os.makedirs( path )
  except Exception:
    pass
  if not os.path.isdir( path ):
    raise Exception( 'Can not create %s' % path )

class AgentModule:
  """ Base class for all agent modules

      This class is used by the AgentReactor Class to steer the execution of
      DIRAC Agents.

      For this purpose the following methods are used:
      - am_initialize()      just after instantiated
      - am_getPollingTime()  to set the execution frequency
      - am_getMaxCycles()    to determine the number of cycles
      - am_go()              for the actual execution of one cycle

      Before each iteration, the following methods are used to determine
      if the new cycle is to be started.
      - am_getModuleParam( 'alive' )
      - am_checkStopAgentFile()
      - am_removeStopAgentFile()

      To start new execution cycle the following methods are used
      - am_getCyclesDone()
      - am_setOption( 'MaxCycles', maxCycles )

      At the same time it provides all Agents with common interface.
      All Agent class must inherit from this base class and must implement
      at least the following method:
      - execute()            main method called in the agent cycle

      Additionally they may provide:
      - initialize()         for initial settings
      - finalize()           the graceful exit

      - beginExecution()     before each execution cycle
      - endExecution()       at the end of each execution cycle

      The agent can be stopped either by a signal or by creating a 'stop_agent' file
      in the controlDirectory defined in the agent configuration

  """

  def __init__( self, agentName, loadName, baseAgentName = False, properties = {} ):
    """
      Common __init__ method for all Agents.
      All Agent modules must define:
      __doc__
      __RCSID__
      They are used to populate __codeProperties

      The following Options are used from the Configuration:
      - /LocalSite/InstancePath
      - /DIRAC/Setup
      - Status
      - Enabled
      - PollingTime            default = 120
      - MaxCycles              default = 500
      - ControlDirectory       control/SystemName/AgentName
      - WorkDirectory          work/SystemName/AgentName
      - shifterProxy           ''
      - shifterProxyLocation   WorkDirectory/SystemName/AgentName/.shifterCred

      It defines the following default Options that can be set via Configuration (above):
      - MonitoringEnabled     True
      - Enabled               True if Status == Active
      - PollingTime           120
      - MaxCycles             500
      - ControlDirectory      control/SystemName/AgentName
      - WorkDirectory         work/SystemName/AgentName
      - shifterProxy          False
      - shifterProxyLocation  work/SystemName/AgentName/.shifterCred

      different defaults can be set in the initialize() method of the Agent using am_setOption()

      In order to get a shifter proxy in the environment during the execute()
      the configuration Option 'shifterProxy' must be set, a default may be given
      in the initialize() method.

    """
    if baseAgentName and agentName == baseAgentName:
      self.log = gLogger
      standaloneModule = True
    else:
      self.log = gLogger.getSubLogger( agentName, child = False )
      standaloneModule = False

    self.__basePath = gConfig.getValue( '/LocalSite/InstancePath', rootPath )
    self.__agentModule = None
    self.__codeProperties = {}
    self.__getCodeInfo()

    self.__moduleProperties = { 'fullName' : agentName,
                                'loadName' : loadName,
                                'section' : PathFinder.getAgentSection( agentName ),
                                'loadSection' : PathFinder.getAgentSection( loadName ),
                                'standalone' : standaloneModule,
                                'cyclesDone' : 0,
                                'totalElapsedTime' : 0,
                                'setup' : gConfig.getValue( "/DIRAC/Setup", "Unknown" ),
                                'alive' : True }
    self.__moduleProperties[ 'system' ], self.__moduleProperties[ 'agentName' ] = agentName.split( "/" )
    self.__configDefaults = {}
    self.__configDefaults[ 'MonitoringEnabled'] = True
    self.__configDefaults[ 'Enabled'] = self.am_getOption( "Status", "Active" ).lower() in ( 'active' )
    self.__configDefaults[ 'PollingTime'] = self.am_getOption( "PollingTime", 120 )
    self.__configDefaults[ 'MaxCycles'] = self.am_getOption( "MaxCycles", 500 )
    self.__configDefaults[ 'ControlDirectory' ] = os.path.join( self.__basePath,
                                                                'control',
                                                                *agentName.split( "/" ) )
    self.__configDefaults[ 'WorkDirectory' ] = os.path.join( self.__basePath,
                                                             'work',
                                                             *agentName.split( "/" ) )
    self.__configDefaults[ 'shifterProxy' ] = ''
    self.__configDefaults[ 'shifterProxyLocation' ] = os.path.join( self.__configDefaults[ 'WorkDirectory' ],
                                                                        '.shifterCred' )


    if type( properties ) == types.DictType:
      for key in properties:
        self.__moduleProperties[ key ] = properties[ key ]
      self.__moduleProperties[ 'executors' ] = [ ( self.execute, () ) ]
      self.__moduleProperties[ 'shifterProxy' ] = False

    self.__monitorLastStatsUpdate = -1
    self.monitor = None
    self.__initializeMonitor()
    self.__initialized = False

  def __getCodeInfo( self ):
    versionVar = "__RCSID__"
    docVar = "__doc__"
    try:
      self.__agentModule = __import__( self.__class__.__module__,
                                       globals(),
                                       locals(),
                                       versionVar )
    except Exception:
      self.log.exception( "Cannot load agent module" )
    for prop in ( ( versionVar, "version" ), ( docVar, "description" ) ):
      try:
        self.__codeProperties[ prop[1] ] = getattr( self.__agentModule, prop[0] )
      except Exception:
        self.log.error( "Missing %s" % prop[0] )
        self.__codeProperties[ prop[1] ] = 'unset'
    self.__codeProperties[ 'DIRACVersion' ] = DIRAC.version
    self.__codeProperties[ 'platform' ] = DIRAC.platform

  def am_initialize( self, *initArgs ):
    agentName = self.am_getModuleParam( 'fullName' )
    result = self.initialize( *initArgs )
    if not isReturnStructure( result ):
      return S_ERROR( "initialize must return S_OK/S_ERROR" )
    if not result[ 'OK' ]:
      return S_ERROR( "Error while initializing %s: %s" % ( agentName, result[ 'Message' ] ) )
   
 _checkDir( self.am_getControlDirectory
() )
    workDirectory = self.am_getWorkDirectory()
    _checkDir( workDirectory )
    # Set the work directory in an environment variable available to subprocesses if needed
    os.environ['AGENT_WORKDIRECTORY'] = workDirectory

    self.__moduleProperties[ 'shifterProxy' ] = self.am_getOption( 'shifterProxy' )
    if self.am_monitoringEnabled():
      self.monitor.enable()
    if len( self.__moduleProperties[ 'executors' ] ) < 1:
      return S_ERROR( "At least one executor method has to be defined" )
    if not self.am_Enabled():
      return S_ERROR( "Agent is disabled via 