""" Priority corrector for the group and in-group shares
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id$"

from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
from DIRAC.WorkloadManagementSystem.private.correctors.BaseCorrector import BaseCorrector


class SharesCorrector(object):
    def __init__(self, opsHelper):
        if not opsHelper:
            opsHelper = Operations()
        self.__opsHelper = opsHelper
        self.__log = gLogger.getSubLogger("SharesCorrector")
        self.__shareCorrectors = {}
        self.__correctorsOrder = []
        self.__baseCS = "JobScheduling/ShareCorrections"

    def __getCSValue(self, path, defaultValue=""):
        return self.__opsHelper.getValue("%s/%s" % (self.__baseCS, path), defaultValue)

    def __getCorrectorClass(self, correctorName):
        baseImport = "WorkloadManagementSystem.private.correctors"
        fullCN = "%s.%sCorrector" % (baseImport, correctorName)
        result = ObjectLoader().getObjects(baseImport, ".*Corrector", parentClass=BaseCorrector)
        if not result["OK"]:
            return result
        data = result["Value"]
        if fullCN not in data:
            return S_ERROR("Can't find corrector %s" % fullCN)
        return S_OK(data[fullCN])

    def instantiateRequiredCorrectors(self):
        correctorsToStart = self.__getCSValue("ShareCorrectorsToStart", [])
        self.__correctorsOrder = correctorsToStart
        self.__log.info("Correctors requested: %s" % ", ".join(correctorsToStart))
        for corrector in self.__shareCorrectors:
            if corrector not in correctorsToStart:
                self.__log.info("Stopping corrector %s" % corrector)
                del self.__shareCorrectors[corrector]
        for corrector in correctorsToStart:
            if corrector not in self.__shareCorrectors:
                self.__log.info("Starting corrector %s" % corrector)
                result = self.__opsHelper.getSections("%s/%s" % (self.__baseCS, corrector))
                if not result["OK"]:
                    self.__log.error(
                        "Cannot get list of correctors to instantiate",
                        " for corrector type %s: %s" % (corrector, result["Message"]),
                    )
                    continue
                groupCorrectors = result["Value"]
                self.__shareCorrectors[corrector] = {}
                result = self.__getCorrectorClass(corrector)
                if not result["OK"]:
                    self.__log.error("Cannot instantiate corrector", "%s %s" % (corrector, result["Message"]))
                    continue
                correctorClass = result["Value"]
                for groupCor in groupCorrectors:
                    groupPath = "%s/%s/Group" % (corrector, groupCor)
                    groupToCorrect = self.__getCSValue(groupPath, "")
                    if groupToCorrect:
                        groupKey = "gr:%s" % groupToCorrect
                    else:
                        groupKey = "global"
                    self.__log.info(
                        "Instantiating group corrector %s (%s) of type %s" % (groupCor, groupToCorrect, corrector)
                    )
                    if groupKey in self.__shareCorrectors[corrector]:
                        self.__log.error(
                            "There are two group correctors defined",
                            " for %s type (group %s)" % (corrector, groupToCorrect),
                        )
                    else:
                        groupCorPath = "%s/%s/%s" % (self.__baseCS, corrector, groupCor)
                        correctorObj = correctorClass(self.__opsHelper, groupCorPath, groupToCorrect)
                        result = correctorObj.initialize()
                        if not result["OK"]:
                            self.__log.error(
                                "Could not initialize corrector %s for %s: %s"
                                % (corrector, groupKey, result["Message"])
                            )
                        else:
                            self.__shareCorrectors[corrector][groupKey] = correctorObj
        return S_OK()

    def updateCorrectorsKnowledge(self):
        for corrector in self.__shareCorrectors:
            for groupTC in self.__shareCorrectors[corrector]:
                self.__shareCorrectors[corrector][groupTC].updateHistoryKnowledge()

    def update(self):
        self.instantiateRequiredCorrectors()
        self.updateCorrectorsKnowledge()

    def correctShares(self, shareDict, group=""):
        if group:
            groupKey = "gr:%s" % group
        else:
            groupKey = "global"
        for corrector in self.__shareCorrectors:
            if groupKey in self.__shareCorrectors[corrector]:
                shareDict = self.__shareCorrectors[corrector][groupKey].applyCorrection(shareDict)
        return shareDict
