###############################################################################
# (c) Copyright 2019 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "LICENSE".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""Set of functions used by the DMS checking scripts."""
import os
from DIRAC import gLogger, S_OK
from DIRAC.Core.Utilities.List import breakListIntoChunks
from DIRAC.Resources.Storage.StorageElement import StorageElement
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from DIRAC.DataManagementSystem.Utilities.DMSHelpers import DMSHelpers

from LHCbDIRAC.DataManagementSystem.Client.ScriptExecutors import (
    removeFiles,
    removeReplicas,
    registerBK2FC,
    printDMResult,
)


def __removeFile(lfns):
    """Use the ScriptExecutors removeFile method."""
    if isinstance(lfns, str):
        lfns = [lfns]
    removeFiles(lfns)


def __removeReplica(lfnDict):
    """Use the ScriptExecutors removeReplicas method."""
    seLFNs = {}
    for lfn in lfnDict:
        for se in lfnDict[lfn]:
            seLFNs.setdefault(se, []).append(lfn)
    for se, lfns in seLFNs.items():
        removeReplicas(lfns, [se])
    return seLFNs


def __replaceReplica(dm, seLFNs):
    """Re-replicate replicas that had just been removed because they were bad It
    uses the DataManager instance from the ConsistencyCheck."""
    if seLFNs:
        gLogger.notice("Now replicating bad replicas...")
        success = {}
        failed = {}
        for se, lfns in seLFNs.items():
            for lfn in lfns:
                res = dm.replicateAndRegister(lfn, se)
                if res["OK"]:
                    success.update(res["Value"]["Successful"])
                    failed.update(res["Value"]["Failed"])
                else:
                    failed[lfn] = res["Message"]

        failures = 0
        errors = {}
        for lfn, reason in failed.items():
            reason = str(reason)
            errors[reason] = errors.setdefault(reason, 0) + 1
            failures += 1
        gLogger.notice("\t%d success, %d failures%s" % (len(success), failures, ":" if failures else ""))
        for reason in errors:
            gLogger.notice("\tError %s : %d files" % (reason, errors[reason]))


def _dumpErrorAndFiles(title, lfnList, maxFiles, dumpStr, fileName, fp):
    """Print out errors and a (restricted) list of LFNs dump the whole list in a
    file."""
    gLogger.error(title)
    if not gLogger.info("\n".join(sorted(lfnList))):
        if len(lfnList) > maxFiles:
            gLogger.error("First %d files:" % maxFiles)
        gLogger.error("\n".join(sorted(lfnList)[0:maxFiles]))
    if fileName is None:
        return None
    if fp is None:
        fp = open(fileName, "w")
    dumpStr1 = dumpStr.strip().replace("\n", "")
    if not dumpStr.endswith(" "):
        dumpStr += " "
    if not dumpStr.startswith("\n"):
        dumpStr = "\n" + dumpStr
    fp.write(dumpStr)
    fp.write(dumpStr.join([""] + sorted(lfnList)))
    gLogger.notice(f"Whole list available with 'grep {dumpStr1} {fileName}'")
    return fp


def _getUniqueFileName(name):
    """Get a file name that doesn't exist given a prefix."""
    suffix = ""
    nb = 0
    while True:
        fileName = name + f"{suffix}.txt"
        if not os.path.exists(fileName):
            break
        nb += 1
        suffix = "-%d" % nb
    return fileName


def doCheckFC2SE(cc, bkCheck=True, fixIt=False, replace=False, maxFiles=None, fixOption="FixIt"):
    """Method actually calling for the the check using ConsistencyChecks module
    It prints out results and calls corrective actions if required."""
    cc.checkFC2SE(bkCheck)

    if maxFiles is None:
        maxFiles = 20
    fileName = _getUniqueFileName("CheckFC2SE")
    fp = None

    if cc.inSEbutNotInFC:
        gLogger.notice(">>>>")
        gLogger.notice("Some files found in SE but not in FC")
        if fixIt:
            gLogger.notice("Going to register files from BK to FC")
            regResult = {"Successful": [], "Failed": {}}
        for lfn, seList in cc.inSEbutNotInFC.items():
            inBK = lfn in cc.existLFNsBKRepNo or lfn in cc.existLFNsBKRepYes
            gLogger.notice(f"\t{lfn} ({'in BK' if inBK else 'not in BK'}): {','.join(seList)}")
            if fixIt:
                if inBK:
                    res = registerBK2FC([lfn], seList)
                    if res["OK"]:
                        regResult["Successful"] += res["Value"]["Successful"]
                        regResult["Failed"].update(res["Value"]["Failed"])
                    else:
                        regResult["Failed"][lfn] = res["Message"]
                else:
                    regResult["Failed"][lfn] = "Not in BK"

        if fixIt:
            # We can remove these files from the BK-No list as we just set it back
            for lfn in regResult["Successful"]:
                cc.existLFNsBKRepNo.pop(lfn, None)
            printDMResult(S_OK(regResult))
        else:
            gLogger.notice("Use --FixIt to register files to the FC if they are in BK")
        gLogger.notice("<<<<")
    elif cc.seList:
        gLogger.notice("None of the files absent in FC were found at specified SEs")

    if cc.notRegisteredAtSE:
        gLogger.notice(">>>>")
        gLogger.notice("Some files found in SE but not registered in FC")
        if fixIt:
            gLogger.notice("Going to register files in FC")
            registrationProtocol = DMSHelpers().getRegistrationProtocols()
            dm = DataManager()
            replicaTuples = []
            regRepResult = {"Successful": {}, "Failed": {}}
            for lfn, seList in cc.notRegisteredAtSE.items():
                for se in seList:
                    res = StorageElement(se).getURL(lfn, protocol=registrationProtocol)
                    if res["OK"]:
                        pfn = res["Value"]["Successful"][lfn]
                        replicaTuples.append((lfn, pfn, se))
            res = dm.registerReplica(replicaTuples)
            if res["OK"]:
                regRepResult["Successful"].update(res["Value"]["Successful"])
                regRepResult["Failed"].update(res["Value"]["Failed"])
            else:
                regRepResult["Failed"][lfn] = res["Message"]
            printDMResult(S_OK(regRepResult))
        else:
            gLogger.notice("Use --FixIt to register replicas in the FC")
        gLogger.notice("<<<<")

    if cc.existLFNsBKRepNo:
        gLogger.notice(">>>>")
        affectedRuns = {str(run) for run in cc.existLFNsBKRepNo.values() if run}
        title = "%d files are in the FC (or SE) but have replica = NO in BK:\nAffected runs: %s" % (
            len(cc.existLFNsBKRepNo),
            ",".join(sorted(affectedRuns)) if affectedRuns else "None",
        )
        fp = _dumpErrorAndFiles(title, cc.existLFNsBKRepNo, maxFiles, "InFCButBKNo", fileName, fp)
        if fixIt:
            gLogger.notice("Going to fix them, setting the replica flag")
            res = cc.bkClient.addFiles(list(cc.existLFNsBKRepNo))
            if res["OK"]:
                gLogger.notice("\tSuccessfully added replica flag")
            else:
                gLogger.error("Failed to set the replica flag", res["Message"])
        else:
            gLogger.notice(f"Use --{fixOption} to fix it (set the replica flag)")
        gLogger.notice("<<<<")

    elif bkCheck:
        gLogger.notice("No files in FC with replica = NO in BK -> OK!")

    if cc.existLFNsNotInBK:
        gLogger.notice(">>>>")

        title = f"{len(cc.existLFNsNotInBK)} files are in the FC but are NOT in BK"
        fp = _dumpErrorAndFiles(title, cc.existLFNsNotInBK, maxFiles, "InFCButNotInBK ", fileName, fp)
        if fixIt:
            gLogger.notice("Going to fix them, by removing from the FC and storage")
            __removeFile(cc.existLFNsNotInBK)
        else:
            gLogger.notice(f"Use --{fixOption} to fix it (remove from FC and storage)")
        gLogger.notice("<<<<")

    else:
        gLogger.notice("No files in FC not in BK -> OK!")

    seOK = True
    if cc.existLFNsNoSE:
        gLogger.notice(">>>>")

        seOK = False
        nOfExistLFNsNoSE = len(cc.existLFNsNoSE)
        title = f"{nOfExistLFNsNoSE} file{'' if nOfExistLFNsNoSE == 1 else 's'}"
        title += f" {'is' if nOfExistLFNsNoSE == 1 else 'are'} in the BK and FC"
        title += f" but {'is' if nOfExistLFNsNoSE == 1 else 'are'} missing at some SEs"
        fixStr = "removing them from catalogs" if not replace else "re-replicating them"
        fp = _dumpErrorAndFiles(title, cc.existLFNsNoSE, maxFiles, "InFCButNotInSE ", fileName, fp)
        if fixIt:
            gLogger.notice("Going to fix, " + fixStr)
            removeLfns = []
            replicasToRemove = {}
            for lfn, ses in cc.existLFNsNoSE.items():
                if ses == "All":
                    removeLfns.append(lfn)
                else:
                    replicasToRemove.setdefault(lfn, []).extend(ses)
            if removeLfns:
                __removeFile(removeLfns)
            if replicasToRemove:
                seLFNs = __removeReplica(replicasToRemove)
                if replace:
                    __replaceReplica(cc.dataManager, seLFNs)
        else:
            if not replace:
                fixStr += "; use --Replace if you want to re-replicate them"
            gLogger.notice(f"Use --{fixOption} to fix it ({fixStr})")
        gLogger.notice("<<<<")

    else:
        gLogger.notice("No missing replicas at sites -> OK!")

    if cc.existLFNsBadFiles or cc.existLFNsNotExisting:
        gLogger.notice(">>>>")

        seOK = False
        if cc.existLFNsBadFiles:
            title = f"{len(cc.existLFNsBadFiles)} files have a bad checksum"
            fp = _dumpErrorAndFiles(title, cc.existLFNsBadFiles, maxFiles, "BadFiles", fileName, fp)
        if cc.existLFNsNotExisting:
            title = f"{len(cc.existLFNsNotExisting)} files don't exist at any SE"
            fp = _dumpErrorAndFiles(title, cc.existLFNsNotExisting, maxFiles, "LostFiles", fileName, fp)
        toRemove = sorted(set(cc.existLFNsBadFiles) | set(cc.existLFNsNotExisting))
        if fixIt:
            gLogger.notice("Going to fix them, removing files from catalogs and storage")
            __removeFile(toRemove)
        else:
            gLogger.notice(f"Use --{fixOption} to fix it (remove files from SE and catalogs)")
        gLogger.notice("<<<<")

    if cc.existLFNsBadReplicas:
        gLogger.notice(">>>>")

        seOK = False
        title = f"{len(cc.existLFNsBadReplicas)} replicas have a bad checksum"
        badChecksum = [f"{lfn} @ {','.join(sorted(se))}" for lfn, se in cc.existLFNsBadReplicas.items()]
        fp = _dumpErrorAndFiles(title, badChecksum, maxFiles, "BadChecksum", fileName, fp)
        fixStr = "remove replicas from SE and catalogs" if not replace else "re-replicating them"
        if fixIt:
            gLogger.notice("Going to fix, " + fixStr)
            seLFNs = __removeReplica(cc.existLFNsBadReplicas)
            if replace:
                __replaceReplica(cc.dataManager, seLFNs)
        else:
            if not replace:
                fixStr += "; use --Replace if you want to re-replicate them"
            gLogger.notice(f"Use --{fixOption} to fix it ({fixStr})")
        gLogger.notice("<<<<")

    if not cc.existLFNsBadFiles and not cc.existLFNsBadReplicas:
        gLogger.notice("No replicas have a bad checksum -> OK!")
    if seOK:
        gLogger.notice("All files exist and have a correct checksum -> OK!")


def doCheckFC2BK(cc, fixFC=False, fixBK=False, listAffectedRuns=False, checkFC2SE=True):
    """Method actually calling for the the check using ConsistencyChecks module
    It prints out results and calls corrective actions if required."""
    cc.checkFC2BK()

    maxFiles = 10
    suffix = ""
    nb = 0
    baseName = "CheckFC2BK" + (f"-{cc.prod}" if cc.prod else "")
    while True:
        fileName = baseName + f"{suffix}.txt"
        if not os.path.exists(fileName):
            break
        nb += 1
        suffix = "-%d" % nb
    fp = None
    if cc.existLFNsBKRepNo:
        gLogger.notice(">>>>")

        affectedRuns = list({str(run) for run in cc.existLFNsBKRepNo.values()})
        nOfExistLFNsBKRepNo = len(cc.existLFNsBKRepNo)
        errMsg = f"{nOfExistLFNsBKRepNo} file{'' if nOfExistLFNsBKRepNo == 1 else 's'} "
        errMsg += f"{'is' if nOfExistLFNsBKRepNo == 1 else 'are'} in the FC but have replica = NO in BK"
        gLogger.error(errMsg)
        if checkFC2SE:
            from LHCbDIRAC.DataManagementSystem.Client.ConsistencyChecks import ConsistencyChecks

            ccAux = ConsistencyChecks()
            gLogger.notice(f"====== Now checking {nOfExistLFNsBKRepNo} files from FC to SE ======")
            ccAux.lfns = list(cc.existLFNsBKRepNo)
            doCheckFC2SE(ccAux, bkCheck=False, fixIt=fixFC, fixOption="FixFC")
            cc.existLFNsBKRepNo = set(cc.existLFNsBKRepNo)
            cc.existLFNsBKRepNo -= set(ccAux.existLFNsNoSE)
            cc.existLFNsBKRepNo -= set(ccAux.existLFNsNotExisting)
            cc.existLFNsBKRepNo -= set(ccAux.existLFNsBadFiles)
        cc.existLFNsBKRepNo = sorted(cc.existLFNsBKRepNo)
        if cc.existLFNsBKRepNo:
            msg = f"====== Completed, {nOfExistLFNsBKRepNo} file{'s' if nOfExistLFNsBKRepNo == 1 else 's'} "
            msg += f"{'is' if nOfExistLFNsBKRepNo == 1 else 'are'} in the FC and SE but have replica = NO in BK ======"
            gLogger.notice(msg)
            if fp is None:
                fp = open(fileName, "w")
            fp.write("\nInFCButBKNo ".join([""] + sorted(cc.existLFNsBKRepNo)))
            res = cc.bkClient.getFileMetadata(cc.existLFNsBKRepNo)
            if not res["OK"]:
                gLogger.fatal("Unable to get file metadata", res["Message"])
                return
            if res["Value"]["Failed"]:
                gLogger.error("No metadata found for some files", f"{len(res['Value']['Failed'])} files")
            success = res["Value"]["Successful"]
            filesInvisible = {lfn for lfn, meta in success.items() if meta["VisibilityFlag"][0].upper() == "N"}
            filesVisible = set(success) - filesInvisible
            gLogger.notice(f"{len(filesVisible)} files are visible, {len(filesInvisible)} files are invisible")
            # Try and print the whole as INFO (in case --Verbose was used).
            #   If nothing printed, print a limited number of files as ERROR
            if not gLogger.info(
                "\n".join(
                    f"{lfn} : Visi {success.get(lfn, {}).get('VisibilityFlag', '?')}"
                    for lfn in sorted(cc.existLFNsBKRepNo)
                )
            ):
                if len(cc.existLFNsBKRepNo) > maxFiles:
                    gLogger.notice("First %d files:" % maxFiles)
                gLogger.error(
                    "\n".join(
                        f"{lfn} : Visi {success.get(lfn, {}).get('VisibilityFlag', '?')}"
                        for lfn in sorted(cc.existLFNsBKRepNo)[0:maxFiles]
                    )
                )
            if listAffectedRuns:
                gLogger.notice(f"Affected runs: {','.join(affectedRuns)}")
            gLogger.notice(f"Full list of files:    grep InFCButBKNo {fileName}")
            if fixBK:
                gLogger.notice("Going to fix them, setting the replica flag")
                res = cc.bkClient.addFiles(list(success))
                if res["OK"]:
                    gLogger.notice(f"\tSuccessfully added replica flag to {len(success)} files")
                else:
                    gLogger.error("Failed to set the replica flag", res["Message"])
            elif fixFC:
                gLogger.notice("Going to fix them, by removing from the FC and storage")
                __removeFile(success)
            else:
                gLogger.notice(
                    "Use --FixBK to fix it (set the replica flag) or --FixFC (for removing from FC and storage)"
                )
        else:
            gLogger.notice("====== Completed, no files in the FC with replica = NO in BK ======")
        gLogger.notice("<<<<")

    else:
        gLogger.notice("No files in FC with replica = NO in BK -> OK!")

    if cc.existLFNsNotInBK:
        gLogger.notice(">>>>")

        gLogger.error(f"{len(cc.existLFNsNotInBK)} files are in the FC but are NOT in BK:")
        if fp is None:
            fp = open(fileName, "w")
        fp.write("\nInFCNotInBK ".join([""] + sorted(cc.existLFNsNotInBK)))
        if not gLogger.info("\n".join(sorted(cc.existLFNsNotInBK))):
            if len(cc.existLFNsNotInBK) > maxFiles:
                gLogger.notice("First %d files:" % maxFiles)
            gLogger.error("\n".join(sorted(cc.existLFNsNotInBK[0:maxFiles])))
        gLogger.notice(f"Full list of files:    grep InFCNotInBK {fileName}")
        if fixFC:
            gLogger.notice("Going to fix them, by removing from the FC and storage")
            __removeFile(cc.existLFNsNotInBK)
        else:
            gLogger.notice("Use --FixFC to fix it (remove from FC and storage)")
        gLogger.notice("<<<<")

    else:
        gLogger.notice("No files in FC not in BK -> OK!")
    if fp is not None:
        fp.close()


def doCheckBK2FC(cc, checkAll=False, fixIt=False):
    """Method actually calling for the the check using ConsistencyChecks module
    It prints out results and calls corrective actions if required."""
    cc.checkBK2FC(checkAll)
    maxPrint = 10
    chunkSize = 100
    suffix = ""
    nb = 0
    baseName = "CheckBK2FC" + (f"-{cc.prod}" if cc.prod else "")
    while True:
        fileName = baseName + f"{suffix}.txt"
        if not os.path.exists(fileName):
            break
        nb += 1
        suffix = "-%d" % nb
    fp = None

    if checkAll:
        if cc.existLFNsBKRepNo:
            gLogger.notice(">>>>")

            nFiles = len(cc.existLFNsBKRepNo)
            comment = "%d files are in the FC but have replica = NO in BK" % nFiles
            if nFiles > maxPrint:
                comment += " (first %d LFNs) : \n" % maxPrint
            comment += "\n".join(cc.existLFNsBKRepNo[:maxPrint])
            gLogger.error(comment)
            if fp is None:
                fp = open(fileName, "w")
            fp.write("\nInFCButBKNo ".join([""] + sorted(cc.existLFNsBKRepNo)))
            gLogger.notice("To get the full list:")
            gLogger.notice(f"    grep InFCButBKNo {fileName}")
            if fixIt:
                gLogger.notice("Setting the replica flag...")
                nFiles = 0
                for lfnChunk in breakListIntoChunks(cc.existLFNsBKRepNo, chunkSize):
                    res = cc.bkClient.addFiles(lfnChunk)
                    if not res["OK"]:
                        gLogger.error(f"Something wrong: {res['Message']}")
                    else:
                        nFiles += len(lfnChunk)
                gLogger.notice("Successfully set replica flag to %d files" % nFiles)
            else:
                gLogger.notice("Use option --FixIt to fix it (set the replica flag)")
            gLogger.notice("<<<<")

        else:
            gLogger.notice("No LFNs exist in the FC but have replicaFlag = No in the BK -> OK!")

    if cc.absentLFNsBKRepYes:
        gLogger.notice(">>>>")

        nFiles = len(cc.absentLFNsBKRepYes)
        comment = "%d files have replicaFlag = Yes but are not in FC:" % nFiles
        if nFiles > maxPrint:
            comment += " (first %d LFNs) : \n" % maxPrint
        comment += "\n".join(cc.absentLFNsBKRepYes[:maxPrint])
        gLogger.error(comment)
        if fp is None:
            fp = open(fileName, "w")
        fp.write("\nInBKButNotInFC ".join([""] + sorted(cc.absentLFNsBKRepYes)))
        gLogger.notice("To get the full list:")
        gLogger.notice(f"    grep InBKButNotInFC {fileName}")
        if fixIt:
            gLogger.notice("Removing the replica flag...")
            nFiles = 0
            for lfnChunk in breakListIntoChunks(cc.absentLFNsBKRepYes, chunkSize):
                res = cc.bkClient.removeFiles(lfnChunk)
                if not res["OK"]:
                    gLogger.error("Something wrong:", res["Message"])
                else:
                    nFiles += len(lfnChunk)
            gLogger.notice("Successfully removed replica flag to %d files" % nFiles)
        else:
            gLogger.notice("Use option --FixIt to fix it (remove the replica flag)")
        gLogger.notice("<<<<")

    else:
        gLogger.notice("No LFNs have replicaFlag = Yes but are not in the FC -> OK!")
    if fp:
        fp.close()


def doCheckSE(cc, seList, fixIt=False):
    cc.checkSE(seList)

    if cc.absentLFNsInFC:
        gLogger.notice(">>>>")
        gLogger.notice(f"{len(cc.absentLFNsInFC)} files are not in the FC")
        if fixIt:
            __removeFile(cc.absentLFNsInFC)
        else:
            gLogger.notice("Use --FixIt to fix it (remove from catalogs")
        gLogger.notice("<<<<")

    if cc.existLFNsNoSE:
        gLogger.notice("<<<<")
        gLogger.notice(f"{len(cc.existLFNsNoSE)} files are not present at {', '.join(sorted(seList))}")
        gLogger.notice("\n".join(sorted(cc.existLFNsNoSE)))
    else:
        gLogger.notice(f"No LFNs missing at {', '.join(sorted(seList))}")
