"""StageMonitorAgent

This agents queries the storage element about staging requests, to see if files are staged or not.

.. literalinclude:: ../ConfigTemplate.cfg
  :start-after: ##BEGIN StageMonitorAgent
  :end-before: ##END
  :caption: StageMonitorAgent options
  :dedent: 2

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id$"

from DIRAC import gLogger, S_OK, S_ERROR, siteName

from DIRAC.Core.Base.AgentModule import AgentModule
from DIRAC.StorageManagementSystem.Client.StorageManagerClient import StorageManagerClient
from DIRAC.Resources.Storage.StorageElement import StorageElement
from DIRAC.AccountingSystem.Client.Types.DataOperation import DataOperation
from DIRAC.AccountingSystem.Client.DataStoreClient import gDataStoreClient
from DIRAC.Core.Security.ProxyInfo import getProxyInfo

import re

AGENT_NAME = 'StorageManagement/StageMonitorAgent'


class StageMonitorAgent(AgentModule):

  def initialize(self):
    self.stagerClient = StorageManagerClient()
    # This sets the Default Proxy to used as that defined under
    # /Operations/Shifter/DataManager
    # the shifterProxy option in the Configuration can be used to change this default.
    self.am_setOption('shifterProxy', 'DataManager')
    self.storagePlugins = self.am_getOption('StoragePlugins', [])

    return S_OK()

  def execute(self):

    res = getProxyInfo(disableVOMS=True)
    if not res['OK']:
      return res
    self.proxyInfoDict = res['Value']

    return self.monitorStageRequests()

  def monitorStageRequests(self):
    """ This is the third logical task manages the StageSubmitted->Staged transition of the Replicas
    """
    res = self.__getStageSubmittedReplicas()
    if not res['OK']:
      gLogger.fatal(
          "StageMonitor.monitorStageRequests: Failed to get replicas from StorageManagementDB.",
          res['Message'])
      return res
    if not res['Value']:
      gLogger.info("StageMonitor.monitorStageRequests: There were no StageSubmitted replicas found")
      return res
    seReplicas = res['Value']['SEReplicas']
    replicaIDs = res['Value']['ReplicaIDs']
    gLogger.info(
        "StageMonitor.monitorStageRequests: Obtained %s StageSubmitted replicas for monitoring." %
        len(replicaIDs))
    for storageElement, seReplicaIDs in seReplicas.items():
      self.__monitorStorageElementStageRequests(storageElement, seReplicaIDs, replicaIDs)

    gDataStoreClient.commit()

    return S_OK()

  def __monitorStorageElementStageRequests(self, storageElement, seReplicaIDs, replicaIDs):
    terminalReplicaIDs = {}
    oldRequests = []
    stagedReplicas = []

    # Since we are in a given SE, the LFN is a unique key
    lfnRepIDs = {}
    for replicaID in seReplicaIDs:
      lfn = replicaIDs[replicaID]['LFN']
      lfnRepIDs[lfn] = replicaID

    if lfnRepIDs:
      gLogger.info(
          "StageMonitor.__monitorStorageElementStageRequests: Monitoring %s stage requests for %s." %
          (len(lfnRepIDs), storageElement))
    else:
      gLogger.warn("StageMonitor.__monitorStorageElementStageRequests: No requests to monitor for %s." % storageElement)
      return
    oAccounting = DataOperation()
    oAccounting.setStartTime()

    res = StorageElement(storageElement, plugins=self.storagePlugins).getFileMetadata(lfnRepIDs)
    if not res['OK']:
      gLogger.error(
          "StageMonitor.__monitorStorageElementStageRequests: Completely failed to monitor stage requests for replicas",
          res['Message'])
      return
    prestageStatus = res['Value']

    accountingDict = self.__newAccountingDict(storageElement)

    for lfn, reason in prestageStatus['Failed'].items():
      accountingDict['TransferTotal'] += 1
      if re.search('File does not exist', reason):
        gLogger.error("StageMonitor.__monitorStorageElementStageRequests: LFN did not exist in the StorageElement", lfn)
        terminalReplicaIDs[lfnRepIDs[lfn]] = 'LFN did not exist in the StorageElement'
    for lfn, metadata in prestageStatus['Successful'].items():
      if not metadata:
        continue
      staged = metadata.get('Cached', metadata['Accessible'])
      if staged:
        accountingDict['TransferTotal'] += 1
        accountingDict['TransferOK'] += 1
        accountingDict['TransferSize'] += metadata['Size']
        stagedReplicas.append(lfnRepIDs[lfn])
      elif staged is not None:
        oldRequests.append(lfnRepIDs[lfn])  # only ReplicaIDs

    oAccounting.setValuesFromDict(accountingDict)
    oAccounting.setEndTime()
    gDataStoreClient.addRegister(oAccounting)

    # Update the states of the replicas in the database
    if terminalReplicaIDs:
      gLogger.info(
          "StageMonitor.__monitorStorageElementStageRequests: %s replicas are terminally failed." %
          len(terminalReplicaIDs))
      res = self.stagerClient.updateReplicaFailure(terminalReplicaIDs)
      if not res['OK']:
        gLogger.error(
            "StageMonitor.__monitorStorageElementStageRequests: Failed to update replica failures.",
            res['Message'])
    if stagedReplicas:
      gLogger.info(
          "StageMonitor.__monitorStorageElementStageRequests: %s staged replicas to be updated." %
          len(stagedReplicas))
      res = self.stagerClient.setStageComplete(stagedReplicas)
      if not res['OK']:
        gLogger.error(
            "StageMonitor.__monitorStorageElementStageRequests: Failed to updated staged replicas.",
            res['Message'])
      res = self.stagerClient.updateReplicaStatus(stagedReplicas, 'Staged')
      if not res['OK']:
        gLogger.error(
            "StageMonitor.__monitorStorageElementStageRequests: Failed to insert replica status.",
            res['Message'])
    if oldRequests:
      gLogger.info(
          "StageMonitor.__monitorStorageElementStageRequests: %s old requests will be retried." %
          len(oldRequests))
      res = self.__wakeupOldRequests(oldRequests)
      if not res['OK']:
        gLogger.error(
            "StageMonitor.__monitorStorageElementStageRequests: Failed to wakeup old requests.",
            res['Message'])
    return

  def __newAccountingDict(self, storageElement):
    """ Generate a new accounting Dict """

    accountingDict = {}
    accountingDict['OperationType'] = 'Stage'
    accountingDict['User'] = self.proxyInfoDict['username']
    accountingDict['Protocol'] = 'Stager'
    accountingDict['RegistrationTime'] = 0.0
    accountingDict['RegistrationOK'] = 0
    accountingDict['RegistrationTotal'] = 0
    accountingDict['FinalStatus'] = 'Successful'
    accountingDict['Source'] = storageElement
    accountingDict['Destination'] = storageElement
    accountingDict['ExecutionSite'] = siteName()
    accountingDict['TransferTotal'] = 0
    accountingDict['TransferOK'] = 0
    accountingDict['TransferSize'] = 0
    accountingDict['TransferTime'] = self.am_getPollingTime()

    return accountingDict

  def __getStageSubmittedReplicas(self):
    """ This obtains the StageSubmitted replicas from the Replicas table and the RequestID
        from the StageRequests table
    """
    res = self.stagerClient.getCacheReplicas({'Status': 'StageSubmitted'})
    if not res['OK']:
      gLogger.error(
          "StageMonitor.__getStageSubmittedReplicas: Failed to get replicas with StageSubmitted status.",
          res['Message'])
      return res
    if not res['Value']:
      gLogger.debug("StageMonitor.__getStageSubmittedReplicas: No StageSubmitted replicas found to process.")
      return S_OK()
    else:
      gLogger.debug(
          "StageMonitor.__getStageSubmittedReplicas: Obtained %s StageSubmitted replicas(s) to process." %
          len(
              res['Value']))

    seReplicas = {}
    replicaIDs = res['Value']
    for replicaID, info in replicaIDs.items():
      storageElement = info['SE']
      seReplicas.setdefault(storageElement, []).append(replicaID)

    # RequestID was missing from replicaIDs dictionary BUGGY?
    res = self.stagerClient.getStageRequests({'ReplicaID': list(replicaIDs)})
    if not res['OK']:
      return res
    if not res['Value']:
      return S_ERROR('Could not obtain request IDs for replicas %s from StageRequests table' % list(replicaIDs))

    for replicaID, info in res['Value'].items():
      replicaIDs[replicaID]['RequestID'] = info['RequestID']

    return S_OK({'SEReplicas': seReplicas, 'ReplicaIDs': replicaIDs})

  def __wakeupOldRequests(self, oldRequests):
    gLogger.info("StageMonitor.__wakeupOldRequests: Attempting...")
    retryInterval = self.am_getOption('RetryIntervalHour', 2)
    res = self.stagerClient.wakeupOldRequests(oldRequests, retryInterval)
    if not res['OK']:
      gLogger.error("StageMonitor.__wakeupOldRequests: Failed to resubmit old requests.", res['Message'])
      return res
    return S_OK()
