import threading
# Because eval(valenc) might require it
import datetime
from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.Core.Utilities import DEncode, List
from DIRAC.Core.Utilities.ReturnValues import isReturnStructure
from DIRAC.Core.Base.ExecutorModule import ExecutorModule
from DIRAC.WorkloadManagementSystem.Client.JobState.CachedJobState import CachedJobState
from DIRAC.WorkloadManagementSystem.Client.JobState.OptimizationTask import OptimizationTask

class OptimizerExecutor( ExecutorModule ):


  class FreezeTask( Exception ):
    def __init__( self, msg, secs = 0 ):
      self.__msg = msg
      self.__secs = secs
    @property
    def freezeTime( self ):
      return self.__secs
    @property
    def msg( self ):
      return self.__msg

  class TQTask( Exception ):
    def __init__( self, msg ):
      self.__msg = msg
    @property
    def msg( self ):
      return self.__msg

  class SplitTask( Exception ):
    def __init__( self, manifests ):
      self.__manifests = manifests
    @property
    def manifests( self ):
      return self.__manifests


  class JobLog:

    class LogWrap:

      def __init__( self, log, jid, funcName ):
        self.__log = log
        self.__jid = jid
        self.__funcName = funcName

      def __call__( self, msg, varMsg = "" ):
        try:
          funcObj = getattr( self.__log, self.__funcName )
        except AttributeError:
          raise AttributeError( "Logger does not have %s method" % self.__funcName )
        msg = "\n".join( "[JID %s] %s" % ( self.__jid, line ) for line in msg.split( "\n" ) )
        funcObj( msg, varMsg )

    def __init__( self, log, jid ):
      self.__jid = jid
      self.__log = log

    def __nonzero__( self ):
      return True

    def __getattr__( self, name ):
      return self.LogWrap( self.__log, self.__jid, name )

  @classmethod
  def initialize( cls ):
    opName = cls.ex_getProperty( 'fullName' )
    opName = "/".join( opName.split( "/" )[1:] )
    if opName.find( "Agent" ) == len( opName ) - 5:
      opName = opName[ :-5]
    cls.__optimizerName = opName
    maxTasks = cls.ex_getOption( 'Tasks', 1 )
    cls.__jobData = threading.local()
    cls.__jobData.jobState = None
    cls.__jobData.jobLog = None
    cls.ex_setProperty( 'optimizerName', cls.__optimizerName )
    cls.ex_setOption( "FailedStatus", cls.__optimizerName )
    try:
      result = cls.initializeOptimizer()
      if not result[ 'OK' ]:
        return result
    except Exception, excp:
      cls.log.exception( "Error while initializing optimizer" )
      return S_ERROR( "Error initializing: %s" % str( excp ) )
    cls.ex_setMind( "WorkloadManagement/OptimizationMind" )
    return S_OK()

  @classmethod
  def ex_optimizerName( cls ):
    return cls.__optimizerName

  @classmethod
  def initializeOptimizer( cls ):
    return S_OK()


  def setManifestsFromParametric( self, manifestList ):
    self.disableFastTrackForTask()
    self.__jobData.task.setParametricManifests( manifestList )

  def processTask( self, jid, taskObj ):
    jobState = taskObj.jobState
    self.__jobData.task = taskObj
    self.__jobData.jobState = jobState
    self.__jobData.jobLog = self.JobLog( self.log, jid )
    try:
      self.jobLog.info( "Processing" )
      try:
        optResult = self.optimizeJob( jid, jobState )
      except self.FreezeTask, excp:
        fT = excp.freezeTime
        if fT == 0:
          fT = self.ex_getOption( "HoldTime", 600 )
        self.freezeTask( fT )
        self.jobLog.info( "On hold -> %s" % excp.msg )
        jobState.setAppStatus( excp.msg, source = self.ex_optimizerName() )
        optResult = S_OK()
      except self.TQTask:
        taskObj.setTQReady()
        optResult = S_OK()
      except self.SplitTask, excp:
        taskObj.splitJobInto( excp.manifests )
        optResult = S_OK()
      if not isReturnStructure( optResult ):
        raise RuntimeError( "Executor does not return S_OK/S_ERROR!" )
      # Did it go as expected? If not Failed!
      if not optResult[ 'OK' ]:
        self.jobLog.info( "Set to Failed/%s" % optResult[ 'Message' ] )
        minorStatus = minorStatus = "%s optimizer" % self.ex_optimizerName()
        return jobState.setStatus( "Failed", minorStatus = minorStatus,
                                   appStatus = optResult[ 'Message' ],
                                   source = self.ex_optimizerName() )
      #If the manifest is dirty, update it!
      result = jobState.getManifest()
      if not result[ 'OK' ]:
        return result
      manifest = result[ 'Value' ]
      if manifest.isDirty():
        manifest.expand()
        jobState.setManifest( manifest )
      return S_OK()
    finally:
      self.__jobData.jobState = None
      self.__jobData.jobLog = None

  def optimizeJob( self, jid, jobState ):
    raise Exception( "You need to overwrite this method to optimize the job!" )

  def sendJobToTQ( self ):
    raise self.TQTask( "Sending job to TaskQueue" )

  def splitJob( self, manifests ):
    raise self.SplitTask( manifests )

  def setNextOptimizer( self, jobState = None ):
    if not jobState:
      jobState = self.__jobData.jobState
    result = jobState.getOptParameter( 'OptimizerChain' )
    if not result['OK']:
      return result
    opChain = List.fromChar( result[ 'Value' ], "," )
    opName = self.__optimizerName
    try:
      opIndex = opChain.index( opName )
    except ValueError:
      return S_ERROR( "Optimizer %s is not in the chain!" % opName )
    chainLength = len( opChain )
    if opIndex == chainLength - 1:
      return self.sendJobToTQ()
    # Keep optimizing!
    nextOp = opChain[ opIndex + 1 ]
    self.jobLog.info( "Set to Checking/%s" % nextOp )
    return jobState.setStatus( "Checking", nextOp, source = opName )

  def storeOptimizerParam( self, name, value ):
    if not self.__jobData.jobState:
      return S_ERROR( "This function can only be called inside the optimizeJob function" )
    valenc = DEncode.encode( value )
    return self.__jobData.jobState.setOptParameter( name, valenc )

  def retrieveOptimizerParam( self, name ):
    if not self.__jobData.jobState:
      return S_ERROR( "This function can only be called inside the optimizeJob function" )
    result = self.__jobData.jobState.getOptParameter( name )
    if not result[ 'OK' ]:
      return result
    valenc = result[ 'Value' ]
    try:
      value, encLength = DEncode.decode( valenc )
      if encLength == len( valenc ):
        return S_OK( value )
    except Exception:
      self.jobLog.warn( "Opt param %s doesn't seem to be dencoded %s" % ( name, valenc ) )
    return S_OK( eval( valenc ) )


  @property
  def jobLog( self ):
    try:
      if not self.__jobData.jobLog:
        raise RuntimeError( "jobLog can only be invoked inside the optimizeJob function" )
    except AttributeError:
      return gLogger
    return self.__jobData.jobLog

  def deserializeTask( self, taskStub ):
    return OptimizationTask.deserialize( taskStub )

  def serializeTask( self, taskObj ):
    return S_OK( taskObj.serialize() )

  def fastTrackDispatch( self, jid, taskObj ):
    jobState = taskObj.jobState
    result = jobState.getStatus()
    if not result[ 'OK' ]:
      return S_ERROR( "Could not retrieve job status for %s: %s" % ( jid, result[ 'Message' ] ) )
    if taskObj.splitManifests:
      self.log.info( "Job has been split. Avoid fast track" )
      return S_OK()
    status, minorStatus = result[ 'Value' ]
    if status != "Checking":
      self.log.info( "[JID %s] Not in checking state. Avoid fast track" % jid )
      return S_OK()
    result = jobState.getOptParameter( "OptimizerChain" )
    if not result[ 'OK' ]:
      return S_ERROR( "Could not retrieve OptimizerChain for job %s: %s" % ( jid, result[ 'Message' ] ) )
    optChain = result[ 'Value' ]
    if minorStatus not in optChain:
      self.log.info( "[JID %s] End of chain for job" % jid )
      return S_OK()
    self.log.info( "[JID %s] Fast track possible to %s" % ( jid, minorStatus ) )
    return S_OK( "WorkloadManagement/%s" % minorStatus )
