""" StorageManagerHandler is the implementation of the StorageManagementDB in the DISET framework """

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

__RCSID__ = "$Id$"

import six

from DIRAC import gLogger, S_OK
from DIRAC.Core.DISET.RequestHandler import RequestHandler
from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB


class StorageManagerHandler(RequestHandler):
    @classmethod
    def initializeHandler(cls, serviceInfoDict):
        """Initialization of DB object"""

        cls.storageManagementDB = StorageManagementDB()
        return S_OK()

    ######################################################################
    #
    #  Example call back methods
    #

    types_updateTaskStatus = []

    @classmethod
    def export_updateTaskStatus(cls, sourceID, status, successful=[], failed=[]):
        """An example to show the usage of the callbacks."""
        gLogger.info("updateTaskStatus: Received callback information for ID %s" % sourceID)
        gLogger.info("updateTaskStatus: Status = '%s'" % status)
        if successful:
            gLogger.info("updateTaskStatus: %s files successfully staged" % len(successful))
            for lfn, time in successful:
                gLogger.info("updateTaskStatus: %s %s" % (lfn.ljust(100), time.ljust(10)))
        if failed:
            gLogger.info("updateTaskStatus: %s files failed to stage" % len(successful))
            for lfn, time in failed:
                gLogger.info("updateTaskStatus: %s %s" % (lfn.ljust(100), time.ljust(10)))
        return S_OK()

    ######################################################################
    #
    #  Monitoring methods
    #

    types_getTaskStatus = [six.integer_types]

    @classmethod
    def export_getTaskStatus(cls, taskID):
        """Obtain the status of the stage task from the DB."""
        res = cls.storageManagementDB.getTaskStatus(taskID)
        if not res["OK"]:
            gLogger.error("getTaskStatus: Failed to get task status", res["Message"])
        return res

    types_getTaskInfo = [six.integer_types]

    @classmethod
    def export_getTaskInfo(cls, taskID):
        """Obtain the metadata of the stage task from the DB."""
        res = cls.storageManagementDB.getTaskInfo(taskID)
        if not res["OK"]:
            gLogger.error("getTaskInfo: Failed to get task metadata", res["Message"])
        return res

    types_getTaskSummary = [six.integer_types]

    @classmethod
    def export_getTaskSummary(cls, taskID):
        """Obtain the summary of the stage task from the DB."""
        res = cls.storageManagementDB.getTaskSummary(taskID)
        if not res["OK"]:
            gLogger.error("getTaskSummary: Failed to get task summary", res["Message"])
        return res

    types_getTasks = [dict]

    @classmethod
    def export_getTasks(cls, condDict, older=None, newer=None, timeStamp="LastUpdate", orderAttribute=None, limit=None):
        """Get the replicas known to the DB."""
        res = cls.storageManagementDB.getTasks(
            condDict=condDict, older=older, newer=newer, timeStamp=timeStamp, orderAttribute=orderAttribute, limit=limit
        )
        if not res["OK"]:
            gLogger.error("getTasks: Failed to get Cache replicas", res["Message"])
        return res

    types_removeStageRequests = [list]

    @classmethod
    def export_removeStageRequests(cls, replicaIDs):
        res = cls.storageManagementDB.removeStageRequests(replicaIDs)
        if not res["OK"]:
            gLogger.error("removeStageRequests: Failed to remove StageRequests", res["Message"])
        return res

    types_getCacheReplicas = [dict]

    @classmethod
    def export_getCacheReplicas(
        cls, condDict, older=None, newer=None, timeStamp="LastUpdate", orderAttribute=None, limit=None
    ):
        """Get the replcias known to the DB."""
        res = cls.storageManagementDB.getCacheReplicas(
            condDict=condDict, older=older, newer=newer, timeStamp=timeStamp, orderAttribute=orderAttribute, limit=limit
        )
        if not res["OK"]:
            gLogger.error("getCacheReplicas: Failed to get Cache replicas", res["Message"])
        return res

    types_getStageRequests = [dict]

    @classmethod
    def export_getStageRequests(
        cls, condDict, older=None, newer=None, timeStamp="StageRequestSubmitTime", orderAttribute=None, limit=None
    ):
        """Get the replcias known to the DB."""
        res = cls.storageManagementDB.getStageRequests(
            condDict=condDict, older=older, newer=newer, timeStamp=timeStamp, orderAttribute=orderAttribute, limit=limit
        )
        if not res["OK"]:
            gLogger.error("getStageRequests: Failed to get Cache replicas", res["Message"])
        return res

    #
    #                                                Monitoring methods
    #
    ######################################################################

    ####################################################################
    #
    # setRequest is used to initially insert tasks and their associated files. Leaves files in New status.
    #

    types_setRequest = [dict, six.string_types, six.string_types, six.integer_types]

    @classmethod
    def export_setRequest(cls, lfnDict, source, callbackMethod, taskID):
        """This method allows stage requests to be set into the StagerDB"""
        res = cls.storageManagementDB.setRequest(lfnDict, source, callbackMethod, taskID)
        if not res["OK"]:
            gLogger.error("setRequest: Failed to set stage request", res["Message"])
        return res

    ####################################################################
    #
    # The state transition of Replicas method
    #

    types_updateReplicaStatus = [list, six.string_types]

    @classmethod
    def export_updateReplicaStatus(cls, replicaIDs, newReplicaStatus):
        """This allows to update the status of replicas"""
        res = cls.storageManagementDB.updateReplicaStatus(replicaIDs, newReplicaStatus)
        if not res["OK"]:
            gLogger.error("updateReplicaStatus: Failed to update replica status", res["Message"])
        return res

    ####################################################################
    #
    # The state transition of the Replicas from New->Waiting
    #

    types_updateReplicaInformation = [list]

    @classmethod
    def export_updateReplicaInformation(cls, replicaTuples):
        """This method sets the pfn and size for the supplied replicas"""
        res = cls.storageManagementDB.updateReplicaInformation(replicaTuples)
        if not res["OK"]:
            gLogger.error("updateRelicaInformation: Failed to update replica information", res["Message"])
        return res

    ####################################################################
    #
    # The state transition of the Replicas from Waiting->StageSubmitted
    #
    types_getStagedReplicas = []

    @classmethod
    def export_getStagedReplicas(cls):
        """This method obtains the replicas for which all replicas in the task are Staged/StageSubmitted"""
        res = cls.storageManagementDB.getStagedReplicas()
        if not res["OK"]:
            gLogger.error("getStagedReplicas: Failed to obtain Staged/StageSubmitted replicas", res["Message"])
        return res

    types_getWaitingReplicas = []

    @classmethod
    def export_getWaitingReplicas(cls):
        """This method obtains the replicas for which all replicas in the task are Waiting"""
        res = cls.storageManagementDB.getWaitingReplicas()
        if not res["OK"]:
            gLogger.error("getWaitingReplicas: Failed to obtain Waiting replicas", res["Message"])
        return res

    types_getOfflineReplicas = []

    @classmethod
    def export_getOfflineReplicas(cls):
        """This method obtains the replicas for which all replicas in the task are Offline"""
        res = cls.storageManagementDB.getOfflineReplicas()
        if not res["OK"]:
            gLogger.error("getOfflineReplicas: Failed to obtain Offline replicas", res["Message"])
        return res

    types_getSubmittedStagePins = []

    @classmethod
    def export_getSubmittedStagePins(cls):
        """This method obtains the number of files and size of the requests submitted for each storage element"""
        res = cls.storageManagementDB.getSubmittedStagePins()
        if not res["OK"]:
            gLogger.error("getSubmittedStagePins: Failed to obtain submitted request summary", res["Message"])
        return res

    types_insertStageRequest = [dict, [six.integer_types, six.integer_types]]

    @classmethod
    def export_insertStageRequest(cls, requestReplicas, pinLifetime):
        """This method inserts the stage request ID assocaited to supplied replicaIDs"""
        res = cls.storageManagementDB.insertStageRequest(requestReplicas, pinLifetime)
        if not res["OK"]:
            gLogger.error("insertStageRequest: Failed to insert stage request information", res["Message"])
        return res

    ####################################################################
    #
    # The state transition of the Replicas from StageSubmitted->Staged
    #

    types_setStageComplete = [list]

    @classmethod
    def export_setStageComplete(cls, replicaIDs):
        """This method updates the status of the stage request for the supplied replica IDs"""
        res = cls.storageManagementDB.setStageComplete(replicaIDs)
        if not res["OK"]:
            gLogger.error("setStageComplete: Failed to set StageRequest complete", res["Message"])
        return res

    ####################################################################
    #
    # The methods for finalization of tasks
    #
    # Daniela: useless method
    '''types_updateStageCompletingTasks = []
  def export_updateStageCompletingTasks(self):
    """ This method checks whether the file for Tasks in 'StageCompleting' status
    are all Staged and updates the Task status to Staged """
    res = cls.storageManagementDB.updateStageCompletingTasks()
    if not res['OK']:
      gLogger.error('updateStageCompletingTasks: Failed to update StageCompleting tasks.',res['Message'])
    return res
  '''

    types_setTasksDone = [list]

    @classmethod
    def export_setTasksDone(cls, taskIDs):
        """This method sets the status in the Tasks table to Done for the list of supplied task IDs"""
        res = cls.storageManagementDB.setTasksDone(taskIDs)
        if not res["OK"]:
            gLogger.error("setTasksDone: Failed to set status of tasks to Done", res["Message"])
        return res

    types_removeTasks = [list]

    @classmethod
    def export_removeTasks(cls, taskIDs):
        """This method removes the entries from TaskReplicas and Tasks with the supplied task IDs"""
        res = cls.storageManagementDB.removeTasks(taskIDs)
        if not res["OK"]:
            gLogger.error("removeTasks: Failed to remove Tasks", res["Message"])
        return res

    types_removeUnlinkedReplicas = []

    @classmethod
    def export_removeUnlinkedReplicas(cls):
        """This method removes Replicas which have no associated Tasks"""
        res = cls.storageManagementDB.removeUnlinkedReplicas()
        if not res["OK"]:
            gLogger.error("removeUnlinkedReplicas: Failed to remove unlinked Replicas", res["Message"])
        return res

    ####################################################################
    #
    # The state transition of the Replicas from *->Failed
    #

    types_updateReplicaFailure = [dict]

    @classmethod
    def export_updateReplicaFailure(cls, replicaFailures):
        """This method sets the status of the replica to failed with the supplied reason"""
        res = cls.storageManagementDB.updateReplicaFailure(replicaFailures)
        if not res["OK"]:
            gLogger.error("updateRelicaFailure: Failed to update replica failure information", res["Message"])
        return res

    ####################################################################
    #
    # Methods for obtaining Tasks, Replicas with supplied state
    #

    types_getTasksWithStatus = [six.string_types]

    @classmethod
    def export_getTasksWithStatus(cls, status):
        """This method allows to retrieve Tasks with the supplied status"""
        res = cls.storageManagementDB.getTasksWithStatus(status)
        if not res["OK"]:
            gLogger.error("getTasksWithStatus: Failed to get tasks with %s status" % status, res["Message"])
        return res

    types_getReplicasWithStatus = [six.string_types]

    @classmethod
    def export_getReplicasWithStatus(cls, status):
        """This method allows to retrieve replicas with the supplied status"""
        res = cls.storageManagementDB.getCacheReplicas({"Status": status})
        if not res["OK"]:
            gLogger.error("getReplicasWithStatus: Failed to get replicas with %s status" % status, res["Message"])
        return res

    types_getStageSubmittedReplicas = []

    @classmethod
    def export_getStageSubmittedReplicas(cls):
        """This method obtains the replica metadata and the stage requestID for the replicas in StageSubmitted status"""
        res = cls.storageManagementDB.getCacheReplicas({"Status": "StageSubmitted"})
        if not res["OK"]:
            gLogger.error("getStageSubmittedReplicas: Failed to obtain StageSubmitted replicas", res["Message"])
        return res

    types_wakeupOldRequests = [list, six.integer_types]

    @classmethod
    def export_wakeupOldRequests(cls, oldRequests, retryInterval):
        """get only StageRequests with StageRequestSubmitTime older than 1 day AND are still not staged
        delete these requests
        reset Replicas with corresponding ReplicaIDs to Status='New'
        """
        res = cls.storageManagementDB.wakeupOldRequests(oldRequests, retryInterval)
        if not res["OK"]:
            gLogger.error("wakeupOldRequests: Failed to wake up old requests", res["Message"])
        return res

    types_setOldTasksAsFailed = [six.integer_types]

    @classmethod
    def export_setOldTasksAsFailed(cls, daysOld):
        """
        Set Tasks older than "daysOld" number of days to Failed
        These tasks have already been retried every day for staging
        """
        res = cls.storageManagementDB.setOldTasksAsFailed(daysOld)
        if not res["OK"]:
            gLogger.error("setOldTasksAsFailed: Failed to set old Tasks to Failed state. ", res["Message"])
        return res

    types_getAssociatedReplicas = [list]

    @classmethod
    def export_getAssociatedReplicas(cls, replicaIDs):
        """
        Retrieve the list of Replicas that belong to the same Tasks as the provided list
        """
        res = cls.storageManagementDB.getAssociatedReplicas(replicaIDs)
        if not res["OK"]:
            gLogger.error("getAssociatedReplicas: Failed to get Associated Replicas. ", res["Message"])
        return res

    types_killTasksBySourceTaskID = [list]

    @classmethod
    def export_killTasksBySourceTaskID(cls, sourceTaskIDs):
        """Given SourceTaskIDs (jobIDs), this will cancel further staging of files for the corresponding tasks"""
        res = cls.storageManagementDB.killTasksBySourceTaskID(sourceTaskIDs)
        if not res["OK"]:
            gLogger.error("removeTasks: Failed to kill staging", res["Message"])
        return res

    types_getCacheReplicasSummary = []

    @classmethod
    def export_getCacheReplicasSummary(cls):
        """Reports breakdown of file number/size in different staging states across storage elements"""
        res = cls.storageManagementDB.getCacheReplicasSummary()
        if not res["OK"]:
            gLogger.error(" getCacheReplicasSummary: Failed to retrieve summary from server", res["Message"])
        return res
