"""
  The InputData Optimizer Executor queries the file catalog for specified job input data and adds the
  relevant information to the job optimizer parameters to be used during the scheduling decision.
"""

import pprint
import time

from DIRAC                                                           import S_OK, S_ERROR
from DIRAC.WorkloadManagementSystem.Executor.Base.OptimizerExecutor  import OptimizerExecutor
from DIRAC.Resources.Storage.StorageElement                          import StorageElement
from DIRAC.Resources.Catalog.FileCatalog                             import FileCatalog
from DIRAC.DataManagementSystem.Utilities.DMSHelpers                 import DMSHelpers
from DIRAC.DataManagementSystem.Client.DataManager                   import DataManager
from DIRAC.ConfigurationSystem.Client.Helpers.Operations             import Operations
from DIRAC.Core.Utilities.Proxy                                      import executeWithUserProxy

__RCSID__ = "$Id$"


class InputData( OptimizerExecutor ):
  """
      The specific Optimizer must provide the following methods:
        - initializeOptimizer() before each execution cycle
        - optimizeJob() - the main method called for each job
  """

  @classmethod
  def initializeOptimizer( cls ):
    """Initialize specific parameters for InputData executor.
    """
    cls.failedMinorStatus = cls.ex_getOption( '/FailedJobStatus', 'Input Data Not Available' )
    # this will ignore failover SE files
    cls.checkFileMetadata = cls.ex_getOption( 'CheckFileMetadata', True )
    # flag to require Input Data lookup with a user proxy
    cls.checkWithUserProxy = cls.ex_getOption( 'CheckWithUserProxy', False )

    cls.__dataManDict = {}
    cls.__fcDict = {}
    cls.__SEToSiteMap = {}
    cls.__lastCacheUpdate = 0
    cls.__cacheLifeTime = 600

    # Note: this is a default, that right now is generically the default for user jobs, at least for main DIRAC users
    # (since this now doesn't run for production jobs)
    # But this should probably be replaced by what the job actually request.
    # The problem is that the InputDataPolicy is not easy to get (a JDL parameter).
    # This may be used but clear how now
    # cls.__connectionLevel = 'PROTOCOL'

    return S_OK()

  def __getDataManager( self, vo ):
    if vo in self.__dataManDict:
      return self.__dataManDict[vo]
    else:
      try:
        self.__dataManDict[vo] = DataManager( vo = vo )
      except Exception as _e:
        msg = 'Failed to create DataManager'
        self.log.exception( msg )
        return None
      return self.__dataManDict[vo]

  def __getFileCatalog( self, vo ):
    if vo in self.__fcDict:
      return self.__fcDict[vo]
    else:
      try:
        self.__fcDict[vo] = FileCatalog( vo = vo )
      except Exception as _e:
        msg = 'Failed to create FileCatalog'
        self.log.exception( msg )
        return None
      return self.__fcDict[vo]

  def optimizeJob( self, jid, jobState ):
    """ This is the method that needs to be implemented by each and every Executor

        This optimizer will run if and only if it is needed:
          - it will run only if there are input files
          - for production jobs this can be skipped,
            since the logic is already applied by the transformation system, via the TaskManagerPlugins
    """
    # Is it a production job?
    result = jobState.getAttribute( "JobType" )
    if not result['OK']:
      return S_ERROR( "Could not retrieve job type" )
    jobType = result['Value']
    if jobType in Operations().getValue( 'Transformations/DataProcessing', [] ):
      self.jobLog.info( "Skipping optimizer, since this is a Production job" )
      return self.setNextOptimizer()

    # Is there input data or not?
    result = jobState.getInputData()
    if not result[ 'OK' ]:
      self.jobLog.error( "Cannot retrieve input data: %s" % result[ 'Message' ] )
      return S_ERROR( "Cannot retrieve input data" )
    if not result[ 'Value' ]:
      self.jobLog.notice( "No input data. Skipping." )
      return self.setNextOptimizer()

    # From now on we know that it is a user job with input data
    inputData = result[ 'Value' ]

    # Check if we already executed this Optimizer and the input data is resolved
    result = self.retrieveOptimizerParam( self.ex_getProperty( 'optimizerName' ) )
    if result['OK'] and result['Value']:
      self.jobLog.info( "InputData optimizer ran already" )
    else:
      self.jobLog.info( "Processing input data" )
      if self.checkWithUserProxy:
        result = jobState.getAttribute( "Owner" )
        if not result['OK']:
          return S_ERROR( "Could not retrieve job owner" )
        userName = result['Value']
        result = jobState.getAttribute( "OwnerGroup" )
        if not result['OK']:
          return S_ERROR( "Could not retrieve job owner group" )
        userGroup = result['Value']
        result = self._resolveInputData( jobState, inputData, proxyUserName = userName, proxyUserGroup = userGroup, executionLock = True )  # pylint: disable=unexpected-keyword-arg
      else:
        result = self._resolveInputData( jobState, inputData )
      if not result['OK']:
        self.jobLog.warn( result['Message'] )
        return result

    return self.setNextOptimizer()

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

  @executeWithUserProxy
  def _resolveInputData( self, jobState, inputData ):
    """ This method checks the file catalog for replica information.
    """
    lfns = inputData

    result = jobState.getManifest()
    if not result['OK']:
      return result
    manifest = result['Value']
    vo = manifest.getOption( 'VirtualOrganization' )
    startTime = time.time()
    dm = self.__getDataManager( vo )
    if dm is None:
      return S_ERROR( 'Failed to instantiate DataManager for vo %s' % vo )
    else:
      # This will return already active replicas, excluding banned SEs, and removing tape replicas if there are disk replicas
      result = dm.getReplicasForJobs( lfns )
    self.jobLog.info( 'Catalog replicas lookup time: %.2f seconds ' % ( time.time() - startTime ) )
    if not result['OK']:
      self.log.warn( result['Message'] )
      return result

    replicaDict = result['Value']

    self.jobLog.verbose( "REPLICA DICT: %s" % replicaDict )

    result = self.__checkReplicas( jobState, replicaDict, vo )

    if not result['OK']:
      self.jobLog.error( "Failed to check replicas", result['Message'] )
      return result
    siteCandidates = result[ 'Value' ]

    if self.ex_getOption( 'CheckFileMetadata', True ):
      result = jobState.getManifest()
      if not result['OK']:
        return result
      manifest = result['Value']
      vo = manifest.getOption( 'VirtualOrganization' )
      fc = self.__getFileCatalog( vo )
      if fc is None:
        return S_ERROR( 'Failed to instantiate FileCatalog for vo %s' % vo )
      else:
        guidDict = fc.getFileMetadata( lfns )
      self.jobLog.info( 'Catalog Metadata Lookup Time: %.2f seconds ' % ( time.time() - startTime ) )

      if not guidDict['OK']:
        self.log.warn( guidDict['Message'] )
        return guidDict

      failed = guidDict['Value']['Failed']
      if failed:
        self.log.warn( 'Failed to establish some GUIDs' )
        self.log.warn( failed )

      for lfn in replicaDict['Successful']:
        replicas = replicaDict['Successful'][ lfn ]
        guidDict['Value']['Successful'][lfn].update( replicas )

    resolvedData = {}
    resolvedData['Value'] = guidDict
    resolvedData['SiteCandidates'] = siteCandidates
    self.jobLog.verbose( "Storing:\n%s" % pprint.pformat( resolvedData ) )
    result = self.storeOptimizerParam( self.ex_getProperty( 'optimizerName' ), resolvedData )
    if not result['OK']:
      self.log.warn( result['Message'] )
      return result
    return S_OK( resolvedData )

  #############################################################################
  def __checkReplicas( self, jobState, replicaDict, vo ):
    """Check that all input lfns have valid replicas and can all be found at least in one single site.
    """
    badLFNs = []

    if 'Successful' not in replicaDict:
      return S_ERROR( 'No replica Info available' )

    okReplicas = replicaDict['Successful']
    for lfn in okReplicas:
      if not okReplicas[ lfn ]:
        badLFNs.append( 'LFN:%s -> No replicas available' % ( lfn ) )

    if 'Failed' in replicaDict:
      errorReplicas = replicaDict[ 'Failed' ]
      for lfn in errorReplicas:
        badLFNs.append( 'LFN:%s -> %s' % ( lfn, errorReplicas[ lfn ] ) )

    if badLFNs:
      errorMsg = "\n".join( badLFNs )
      self.jobLog.info( 'Found %s problematic LFN(s):\n%s' % ( len( badLFNs ), errorMsg ) )
      result = jobState.setParameter( self.ex_getProperty( 'optimizerName' ), errorMsg )
      if not result['OK']:
        self.log.error( 'Failed to set job parameter', result['Message'] )
      return S_ERROR( 'Input data not available' )

    return self.__getSiteCandidates( okReplicas, vo )

  #############################################################################
  def __getSitesForSE( self, seName ):
    """ Returns a list of sites having the given SE as a local one.
        Uses the local cache of the site-se information
    """

    # Empty the cache if too old
    now = time.time()
    if ( now - self.__lastCacheUpdate ) > self.__cacheLifeTime:
      self.log.verbose( 'Resetting the SE to site mapping cache' )
      self.__SEToSiteMap = {}
      self.__lastCacheUpdate = now

    if seName not in self.__SEToSiteMap:
      result = DMSHelpers().getSitesForSE( seName )
      if not result['OK']:
        return result
      self.__SEToSiteMap[ seName ] = list( result['Value'] )
    return S_OK( self.__SEToSiteMap[ seName ] )

  #############################################################################
  def __getSiteCandidates( self, okReplicas, vo ):
    """ This method returns a list of possible site candidates based on the job input data requirement.

        For each site candidate, the number of files on disk and tape is resolved.
    """

    lfnSEs = {}
    for lfn in okReplicas:
      replicas = okReplicas[ lfn ]
      siteSet = set()
      for seName in replicas:
        result = self.__getSitesForSE( seName )
        if not result['OK']:
          self.jobLog.warn( "Could not get sites for SE %s: %s" % ( seName, result[ 'Message' ] ) )
          return result
        siteSet.update( result['Value'] )
      lfnSEs[ lfn ] = siteSet

    if not lfnSEs:
      return S_ERROR( "No candidate sites available" )

    # This makes an intersection of all sets in the dictionary and returns a set with it
    siteCandidates = set.intersection( *[ lfnSEs[ lfn ] for lfn in lfnSEs ] )

    if not siteCandidates:
      return S_ERROR( 'No candidate sites available' )

    # In addition, check number of files on tape and disk for each site
    # for optimizations during scheduling
    sitesData = {}
    for siteName in siteCandidates:
      sitesData[ siteName ] = { 'disk': set(), 'tape': set() }

    # Loop time!
    seDict = {}
    for lfn in okReplicas:
      replicas = okReplicas[ lfn ]
      # Check each SE in the replicas
      for seName in replicas:
        # If not already "loaded" the add it to the dict
        if seName not in seDict:
          result = self.__getSitesForSE( seName )
          if not result['OK']:
            self.jobLog.warn( "Could not get sites for SE %s: %s" % ( seName, result[ 'Message' ] ) )
            continue
          siteList = result[ 'Value' ]
          seObj = StorageElement( seName, vo = vo )
          seStatus = seObj.getStatus()
          if not seStatus['OK']:
            return seStatus
          seDict[ seName ] = { 'Sites': siteList, 'Status': seStatus['Value'] }
        # Get SE info from the dict
        seData = seDict[ seName ]
        siteList = seData[ 'Sites' ]
        seStatus = seData[ 'Status' ]
        for siteName in siteList:
          # If not a candidate site then skip it
          if siteName not in siteCandidates:
            continue
          # Add the LFNs to the disk/tape lists
          diskLFNs = sitesData[ siteName ][ 'disk' ]
          tapeLFNs = sitesData[ siteName ][ 'tape' ]
          if seStatus[ 'DiskSE' ]:
            # Sets contain only unique elements, no need to check if it's there
            diskLFNs.add( lfn )
            if lfn in tapeLFNs:
              tapeLFNs.remove( lfn )
          if seStatus[ 'TapeSE' ]:
            if lfn not in diskLFNs:
              tapeLFNs.add( lfn )

    for siteName in sitesData:
      sitesData[siteName]['disk'] = len( sitesData[siteName]['disk'] )
      sitesData[siteName]['tape'] = len( sitesData[siteName]['tape'] )
    return S_OK( sitesData )

# EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#EOF#
