"""
This is the Data Integrity Client which allows the simple reporting of
problematic file and replicas to the IntegrityDB and their status
correctly updated in the FileCatalog.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import six
from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from DIRAC.Resources.Storage.StorageElement import StorageElement
from DIRAC.Resources.Catalog.FileCatalog import FileCatalog
from DIRAC.Core.Utilities.ReturnValues import returnSingleResult
from DIRAC.Core.Base.Client import Client, createClient

__RCSID__ = "$Id$"


@createClient("DataManagement/DataIntegrity")
class DataIntegrityClient(Client):
    """Client exposing the DataIntegrity Service."""

    def __init__(self, **kwargs):

        super(DataIntegrityClient, self).__init__(**kwargs)
        self.setServer("DataManagement/DataIntegrity")
        self.dm = DataManager()
        self.fc = FileCatalog()

    def setFileProblematic(self, lfn, reason, sourceComponent=""):
        """This method updates the status of the file in the FileCatalog and the IntegrityDB

        lfn - the lfn of the file
        reason - this is given to the integrity DB and should reflect the problem observed with the file

        sourceComponent is the component issuing the request.
        """
        if isinstance(lfn, list):
            lfns = lfn
        elif isinstance(lfn, six.string_types):
            lfns = [lfn]
        else:
            errStr = "DataIntegrityClient.setFileProblematic: Supplied file info must be list or a single LFN."
            gLogger.error(errStr)
            return S_ERROR(errStr)
        gLogger.info("DataIntegrityClient.setFileProblematic: Attempting to update %s files." % len(lfns))
        fileMetadata = {}
        for lfn in lfns:
            fileMetadata[lfn] = {"Prognosis": reason, "LFN": lfn, "PFN": "", "SE": ""}
        res = self.insertProblematic(sourceComponent, fileMetadata)
        if not res["OK"]:
            gLogger.error("DataIntegrityClient.setReplicaProblematic: Failed to insert problematics to integrity DB")
        return res

    def reportProblematicReplicas(self, replicaTuple, se, reason):
        """Simple wrapper function around setReplicaProblematic"""
        gLogger.info("The following %s files had %s at %s" % (len(replicaTuple), reason, se))
        for lfn, _pfn, se, reason in sorted(replicaTuple):
            if lfn:
                gLogger.info(lfn)
        res = self.setReplicaProblematic(replicaTuple, sourceComponent="DataIntegrityClient")
        if not res["OK"]:
            gLogger.info("Failed to update integrity DB with replicas", res["Message"])
        else:
            gLogger.info("Successfully updated integrity DB with replicas")

    def setReplicaProblematic(self, replicaTuple, sourceComponent=""):
        """This method updates the status of the replica in the FileCatalog and the IntegrityDB
        The supplied replicaDict should be of the form {lfn :{'PFN':pfn,'SE':se,'Prognosis':prognosis}

        lfn - the lfn of the file
        pfn - the pfn if available (otherwise '')
        se - the storage element of the problematic replica (otherwise '')
        prognosis - this is given to the integrity DB and should reflect the problem observed with the file

        sourceComponent is the component issuing the request.
        """
        if isinstance(replicaTuple, tuple):
            replicaTuple = [replicaTuple]
        elif isinstance(replicaTuple, list):
            pass
        else:
            errStr = (
                "DataIntegrityClient.setReplicaProblematic: Supplied replica info must be a tuple or list of tuples."
            )
            gLogger.error(errStr)
            return S_ERROR(errStr)
        gLogger.info("DataIntegrityClient.setReplicaProblematic: Attempting to update %s replicas." % len(replicaTuple))
        replicaDict = {}
        for lfn, pfn, se, reason in replicaTuple:
            replicaDict[lfn] = {"Prognosis": reason, "LFN": lfn, "PFN": pfn, "SE": se}
        res = self.insertProblematic(sourceComponent, replicaDict)
        if not res["OK"]:
            gLogger.error("DataIntegrityClient.setReplicaProblematic: Failed to insert problematic to integrity DB")
            return res
        for lfn in replicaDict.keys():
            replicaDict[lfn]["Status"] = "Problematic"

        res = self.fc.setReplicaStatus(replicaDict)
        if not res["OK"]:
            errStr = "DataIntegrityClient.setReplicaProblematic: Completely failed to update replicas."
            gLogger.error(errStr, res["Message"])
            return res
        failed = res["Value"]["Failed"]
        successful = res["Value"]["Successful"]
        resDict = {"Successful": successful, "Failed": failed}
        return S_OK(resDict)

    ##########################################################################
    #
    # This section contains the resolution methods for various prognoses
    #

    def __updateCompletedFiles(self, prognosis, fileID):
        gLogger.info("%s file (%d) is resolved" % (prognosis, fileID))
        return self.setProblematicStatus(fileID, "Resolved")

    def __returnProblematicError(self, fileID, res):
        self.incrementProblematicRetry(fileID)
        gLogger.error("DataIntegrityClient failure", res["Message"])
        return res

    def __updateReplicaToChecked(self, problematicDict):
        lfn = problematicDict["LFN"]
        fileID = problematicDict["FileID"]
        prognosis = problematicDict["Prognosis"]
        problematicDict["Status"] = "Checked"

        res = returnSingleResult(self.fc.setReplicaStatus({lfn: problematicDict}))

        if not res["OK"]:
            return self.__returnProblematicError(fileID, res)
        gLogger.info("%s replica (%d) is updated to Checked status" % (prognosis, fileID))
        return self.__updateCompletedFiles(prognosis, fileID)

    def resolveCatalogPFNSizeMismatch(self, problematicDict):
        """This takes the problematic dictionary returned by the integrity DB and resolved the CatalogPFNSizeMismatch prognosis"""
        lfn = problematicDict["LFN"]
        se = problematicDict["SE"]
        fileID = problematicDict["FileID"]

        res = returnSingleResult(self.fc.getFileSize(lfn))
        if not res["OK"]:
            return self.__returnProblematicError(fileID, res)
        catalogSize = res["Value"]
        res = returnSingleResult(StorageElement(se).getFileSize(lfn))
        if not res["OK"]:
            return self.__returnProblematicError(fileID, res)
        storageSize = res["Value"]
        bkKCatalog = FileCatalog(["BookkeepingDB"])
        res = returnSingleResult(bkKCatalog.getFileSize(lfn))
        if not res["OK"]:
            return self.__returnProblematicError(fileID, res)
        bookkeepingSize = res["Value"]
        if bookkeepingSize == catalogSize == storageSize:
            gLogger.info("CatalogPFNSizeMismatch replica (%d) matched all registered sizes." % fileID)
            return self.__updateReplicaToChecked(problematicDict)
        if catalogSize == bookkeepingSize:
            gLogger.info("CatalogPFNSizeMismatch replica (%d) found to mismatch the bookkeeping also" % fileID)
            res = returnSingleResult(self.fc.getReplicas(lfn))
            if not res["OK"]:
                return self.__returnProblematicError(fileID, res)
            if len(res["Value"]) <= 1:
                gLogger.info("CatalogPFNSizeMismatch replica (%d) has no other replicas." % fileID)
                return S_ERROR("Not removing catalog file mismatch since the only replica")
            else:
                gLogger.info("CatalogPFNSizeMismatch replica (%d) has other replicas. Removing..." % fileID)
                res = self.dm.removeReplica(se, lfn)
                if not res["OK"]:
                    return self.__returnProblematicError(fileID, res)
                return self.__updateCompletedFiles("CatalogPFNSizeMismatch", fileID)
        if (catalogSize != bookkeepingSize) and (bookkeepingSize == storageSize):
            gLogger.info("CatalogPFNSizeMismatch replica (%d) found to match the bookkeeping size" % fileID)
            res = self.__updateReplicaToChecked(problematicDict)
            if not res["OK"]:
                return self.__returnProblematicError(fileID, res)
            return self.changeProblematicPrognosis(fileID, "BKCatalogSizeMismatch")
        gLogger.info("CatalogPFNSizeMismatch replica (%d) all sizes found mismatch. Updating retry count" % fileID)
        return self.incrementProblematicRetry(fileID)

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

    def _reportProblematicFiles(self, lfns, reason):
        """Simple wrapper function around setFileProblematic"""
        gLogger.info("The following %s files were found with %s" % (len(lfns), reason))
        for lfn in sorted(lfns):
            gLogger.info(lfn)
        res = self.setFileProblematic(lfns, reason, sourceComponent="DataIntegrityClient")
        if not res["OK"]:
            gLogger.info("Failed to update integrity DB with files", res["Message"])
        else:
            gLogger.info("Successfully updated integrity DB with files")
