########################################################################
# File: RemoveReplica.py
# Author: Krzysztof.Ciba@NOSPAMgmail.com
# Date: 2013/03/25 07:45:06
########################################################################

""" :mod: RemoveReplica

    ===================

    .. module: RemoveReplica

    :synopsis: removeReplica operation handler

    .. moduleauthor:: Krzysztof.Ciba@NOSPAMgmail.com

    RemoveReplica operation handler
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id $"

# #
# @file RemoveReplica.py
# @author Krzysztof.Ciba@NOSPAMgmail.com
# @date 2013/03/25 07:45:17
# @brief Definition of RemoveReplica class.

# # imports
import os

# # from DIRAC
from DIRAC import S_OK
from DIRAC.FrameworkSystem.Client.MonitoringClient import gMonitor
from DIRAC.DataManagementSystem.Agent.RequestOperations.DMSRequestOperationsBase import DMSRequestOperationsBase

from DIRAC.MonitoringSystem.Client.MonitoringReporter import MonitoringReporter

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


class RemoveReplica(DMSRequestOperationsBase):
    """
    .. class:: RemoveReplica

    """

    def __init__(self, operation=None, csPath=None):
        """c'tor

        :param self: self reference
        :param Operation operation: operation to execute
        :param str csPath: CS path for this handler
        """
        # # base class ctor
        DMSRequestOperationsBase.__init__(self, operation, csPath)

    def __call__(self):
        """remove replicas"""

        # The flag  'rmsMonitoring' is set by the RequestTask and is False by default.
        # Here we use 'createRMSRecord' to create the ES record which is defined inside OperationHandlerBase.
        if self.rmsMonitoring:
            self.rmsMonitoringReporter = MonitoringReporter(monitoringType="RMSMonitoring")
        else:
            # # gMonitor stuff
            gMonitor.registerActivity(
                "RemoveReplicaAtt", "Replica removals attempted", "RequestExecutingAgent", "Files/min", gMonitor.OP_SUM
            )
            gMonitor.registerActivity(
                "RemoveReplicaOK", "Successful replica removals", "RequestExecutingAgent", "Files/min", gMonitor.OP_SUM
            )
            gMonitor.registerActivity(
                "RemoveReplicaFail", "Failed replica removals", "RequestExecutingAgent", "Files/min", gMonitor.OP_SUM
            )

        # # prepare list of targetSEs
        targetSEs = self.operation.targetSEList
        # # check targetSEs for removal
        bannedTargets = self.checkSEsRSS(targetSEs, access="RemoveAccess")
        if not bannedTargets["OK"]:
            if self.rmsMonitoring:
                for status in ["Attempted", "Failed"]:
                    self.rmsMonitoringReporter.addRecord(self.createRMSRecord(status, len(self.operation)))
                self.rmsMonitoringReporter.commit()
            else:
                gMonitor.addMark("RemoveReplicaAtt")
                gMonitor.addMark("RemoveReplicaFail")
            return bannedTargets

        if bannedTargets["Value"]:
            return S_OK("%s targets are banned for removal" % ",".join(bannedTargets["Value"]))

        # # get waiting files
        waitingFiles = self.getWaitingFilesList()
        # # and prepare dict
        toRemoveDict = dict((opFile.LFN, opFile) for opFile in waitingFiles)

        self.log.info("Todo: %s replicas to delete from %s SEs" % (len(toRemoveDict), len(targetSEs)))

        if self.rmsMonitoring:
            self.rmsMonitoringReporter.addRecord(self.createRMSRecord("Attempted", len(toRemoveDict)))
        else:
            gMonitor.addMark("RemoveReplicaAtt", len(toRemoveDict) * len(targetSEs))

        # # keep status for each targetSE
        removalStatus = dict.fromkeys(toRemoveDict, None)
        for lfn in removalStatus:
            removalStatus[lfn] = dict.fromkeys(targetSEs, None)

        # # loop over targetSEs
        for targetSE in targetSEs:

            self.log.info("Removing replicas at %s" % targetSE)

            # # 1st step - bulk removal
            bulkRemoval = self._bulkRemoval(toRemoveDict, targetSE)
            if not bulkRemoval["OK"]:
                self.log.error("Bulk replica removal failed", bulkRemoval["Message"])

                if self.rmsMonitoring:
                    self.rmsMonitoringReporter.commit()

                return bulkRemoval

            # # report removal status for successful files
            if self.rmsMonitoring:
                self.rmsMonitoringReporter.addRecord(
                    self.createRMSRecord(
                        "Successful", len(([opFile for opFile in toRemoveDict.values() if not opFile.Error]))
                    )
                )
            else:
                gMonitor.addMark(
                    "RemoveReplicaOK", len([opFile for opFile in toRemoveDict.values() if not opFile.Error])
                )

            # # 2nd step - process the rest again
            toRetry = dict((lfn, opFile) for lfn, opFile in toRemoveDict.items() if opFile.Error)
            for lfn, opFile in toRetry.items():
                self._removeWithOwnerProxy(opFile, targetSE)
                if opFile.Error:
                    if self.rmsMonitoring:
                        self.rmsMonitoringReporter.addRecord(self.createRMSRecord("Failed", 1))
                    else:
                        gMonitor.addMark("RemoveReplicaFail", 1)
                    removalStatus[lfn][targetSE] = opFile.Error
                else:
                    if self.rmsMonitoring:
                        self.rmsMonitoringReporter.addRecord(self.createRMSRecord("Successful", 1))
                    else:
                        gMonitor.addMark("RemoveReplicaOK", 1)

        # # update file status for waiting files
        failed = 0
        for opFile in self.operation:
            if opFile.Status == "Waiting":
                errors = list(set(error for error in removalStatus[opFile.LFN].values() if error))
                if errors:
                    opFile.Error = "\n".join(errors)
                    # This seems to be the only unrecoverable error
                    if "Write access not permitted for this credential" in opFile.Error:
                        failed += 1
                        opFile.Status = "Failed"
                else:
                    opFile.Status = "Done"

        if failed:
            self.operation.Error = "failed to remove %s replicas" % failed

        if self.rmsMonitoring:
            self.rmsMonitoringReporter.commit()

        return S_OK()

    def _bulkRemoval(self, toRemoveDict, targetSE):
        """remove replicas :toRemoveDict: at :targetSE:

        :param dict toRemoveDict: { lfn: opFile, ... }
        :param str targetSE: target SE name
        """
        # Clear the error
        for opFile in toRemoveDict.values():
            opFile.Error = ""
        removeReplicas = self.dm.removeReplica(targetSE, list(toRemoveDict))
        if not removeReplicas["OK"]:
            for opFile in toRemoveDict.values():
                opFile.Error = removeReplicas["Message"]
            return removeReplicas
        removeReplicas = removeReplicas["Value"]
        # # filter out failed
        for lfn, opFile in toRemoveDict.items():
            if lfn in removeReplicas["Failed"]:
                errorReason = str(removeReplicas["Failed"][lfn])
                # If the reason is that the file does not exist,
                # we consider the removal successful
                # TODO: use cmpError once the FC returns the proper error msg corresponding to ENOENT
                if "No such file" not in errorReason:
                    opFile.Error = errorReason
                    self.log.error("Failed removing lfn", "%s:%s" % (lfn, opFile.Error))

        return S_OK()

    def _removeWithOwnerProxy(self, opFile, targetSE):
        """remove opFile replica from targetSE using owner proxy

        :param File opFile: File instance
        :param str targetSE: target SE name
        """
        if "Write access not permitted for this credential" in opFile.Error:
            proxyFile = None
            if "DataManager" in self.shifter:
                # #  you're a data manager - save current proxy and get a new one for LFN and retry
                saveProxy = os.environ["X509_USER_PROXY"]
                try:
                    fileProxy = self.getProxyForLFN(opFile.LFN)
                    if not fileProxy["OK"]:
                        opFile.Error = fileProxy["Message"]
                    else:
                        proxyFile = fileProxy["Value"]
                        removeReplica = self.dm.removeReplica(targetSE, opFile.LFN)
                        if not removeReplica["OK"]:
                            opFile.Error = removeReplica["Message"]
                        else:
                            # Set or reset the error if all OK
                            opFile.Error = removeReplica["Value"]["Failed"].get(opFile.LFN, "")
                finally:
                    if proxyFile:
                        os.unlink(proxyFile)
                    # # put back request owner proxy to env
                    os.environ["X509_USER_PROXY"] = saveProxy
