########################################################################
# File :   ExecutorReactor.py
# Author : Adria Casajus
########################################################################
"""
  DIRAC class to execute Executors

  Executors are an active part of DIRAC.

  All DIRAC executors must inherit from the basic class ExecutorModule

  In the most common case, DIRAC Executors are executed using the dirac-executor command.
  dirac-execuot accepts a list positional arguments.

  dirac-executo then:
  - produces a instance of ExecutorReactor

  Executor modules must be placed under the Executor directory of a DIRAC System.
  DIRAC Systems are called XXXSystem where XXX is the [DIRAC System Name], and
  must inherit from the base class ExecutorModule

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

import time
import threading
from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.Core.DISET.MessageClient import MessageClient
from DIRAC.ConfigurationSystem.Client import PathFinder
from DIRAC.Core.Base.private.ModuleLoader import ModuleLoader
from DIRAC.Core.Base.ExecutorModule import ExecutorModule

__RCSID__ = "$Id$"


class ExecutorReactor(object):
    class AliveLock(object):
        def __init__(self):
            self.__alive = 0
            self.__cond = threading.Condition(threading.Lock())

        def alive(self):
            self.__cond.acquire()
            self.__alive += 1
            self.__cond.release()

        def dead(self):
            self.__cond.acquire()
            self.__alive -= 1
            self.__cond.notify()
            self.__cond.release()

        def lockUntilAllDead(self):
            self.__cond.acquire()
            while True:
                if self.__alive < 1:
                    break
                self.__cond.wait(1)
            self.__cond.release()

    class MindCluster(object):
        def __init__(self, mindName, aliveLock):
            self.__mindName = mindName
            self.__modules = {}
            self.__maxTasks = 1
            self.__reconnectSleep = 1
            self.__reconnectRetries = 10
            self.__extraArgs = {}
            self.__instances = {}
            self.__instanceLock = threading.Lock()
            self.__aliveLock = aliveLock

        def updateMaxTasks(self, mt):
            self.__maxTasks = max(self.__maxTasks, mt)

        def addModule(self, name, exeClass):
            self.__modules[name] = exeClass
            self.__maxTasks = max(self.__maxTasks, exeClass.ex_getOption("MaxTasks", 0))
            self.__reconnectSleep = max(self.__reconnectSleep, exeClass.ex_getOption("ReconnectSleep", 0))
            self.__reconnectRetries = max(self.__reconnectRetries, exeClass.ex_getOption("ReconnectRetries", 0))
            self.__extraArgs[name] = exeClass.ex_getExtraArguments()

        def connect(self):
            self.__msgClient = MessageClient(self.__mindName)
            self.__msgClient.subscribeToMessage("ProcessTask", self.__processTask)
            self.__msgClient.subscribeToDisconnect(self.__disconnected)
            result = self.__msgClient.connect(
                executorTypes=list(self.__modules), maxTasks=self.__maxTasks, extraArgs=self.__extraArgs
            )
            if result["OK"]:
                self.__aliveLock.alive()
                gLogger.info("Connected to %s" % self.__mindName)
            return result

        def __disconnected(self, msgClient):
            retryCount = 0
            while True:
                gLogger.notice("Trying to reconnect to %s" % self.__mindName)
                result = self.__msgClient.connect(
                    executorTypes=list(self.__modules), maxTasks=self.__maxTasks, extraArgs=self.__extraArgs
                )

                if result["OK"]:
                    if retryCount >= self.__reconnectRetries:
                        self.__aliveLock.alive()
                    gLogger.notice("Reconnected to %s" % self.__mindName)
                    return S_OK()
                retryCount += 1
                if retryCount == self.__reconnectRetries:
                    self.__aliveLock.alive()
                gLogger.info("Connect error failed: %s" % result["Message"])
                gLogger.notice("Failed to reconnect. Sleeping for %d seconds" % self.__reconnectSleep)
                time.sleep(self.__reconnectSleep)

        def __storeInstance(self, modName, modObj):
            self.__instanceLock.acquire()
            try:
                self.__instances[modName].append(modObj)
            finally:
                self.__instanceLock.release()

        def __getInstance(self, moduleName):
            self.__instanceLock.acquire()
            try:
                if moduleName not in self.__instances:
                    self.__instances[moduleName] = []
                try:
                    return S_OK(self.__instances[moduleName].pop(0))
                except IndexError:
                    pass
            finally:
                self.__instanceLock.release()
            try:
                modObj = self.__modules[moduleName]
            except KeyError:
                return S_ERROR("Unknown %s executor")
            modInstance = modObj()
            return S_OK(modInstance)

        def __sendExecutorError(self, eType, taskId, errMsg):
            result = self.__msgClient.createMessage("ExecutorError")
            if not result["OK"]:
                return result
            msgObj = result["Value"]
            msgObj.taskId = taskId
            msgObj.errorMsg = errMsg
            msgObj.eType = eType
            return self.__msgClient.sendMessage(msgObj)

        def __processTask(self, msgObj):
            eType = msgObj.eType
            taskId = msgObj.taskId
            taskStub = msgObj.taskStub

            result = self.__moduleProcess(eType, taskId, taskStub)
            if not result["OK"]:
                return self.__sendExecutorError(eType, taskId, result["Message"])
            msgName, taskStub, extra = result["Value"]

            result = self.__msgClient.createMessage(msgName)
            if not result["OK"]:
                return self.__sendExecutorError(
                    eType, taskId, "Can't generate %s message: %s" % (msgName, result["Message"])
                )
            gLogger.verbose("Task %s: Sending %s" % (str(taskId), msgName))
            msgObj = result["Value"]
            msgObj.taskId = taskId
            msgObj.taskStub = taskStub
            if msgName == "TaskError":
                msgObj.errorMsg = extra
                msgObj.eType = eType
            elif msgName == "TaskFreeze":
                msgObj.freezeTime = extra
            return self.__msgClient.sendMessage(msgObj)

        def __moduleProcess(self, eType, taskId, taskStub, fastTrackLevel=0):
            result = self.__getInstance(eType)
            if not result["OK"]:
                return result
            modInstance = result["Value"]
            try:
                result = modInstance._ex_processTask(taskId, taskStub)
            except Exception as excp:
                gLogger.exception("Error while processing task %s" % taskId, lException=excp)
                return S_ERROR("Error processing task %s: %s" % (taskId, excp))

            self.__storeInstance(eType, modInstance)

            if not result["OK"]:
                return S_OK(("TaskError", taskStub, "Error: %s" % result["Message"]))
            taskStub, freezeTime, fastTrackType = result["Value"]
            if freezeTime:
                return S_OK(("TaskFreeze", taskStub, freezeTime))
            if fastTrackType:
                if fastTrackLevel < 10 and fastTrackType in self.__modules:
                    gLogger.notice("Fast tracking task %s to %s" % (taskId, fastTrackType))
                    return self.__moduleProcess(fastTrackType, taskId, taskStub, fastTrackLevel + 1)
                else:
                    gLogger.notice("Stopping %s fast track. Sending back to the mind" % (taskId))

            return S_OK(("TaskDone", taskStub, True))

    #####
    # Start of ExecutorReactor
    #####

    def __init__(self):
        self.__aliveLock = self.AliveLock()
        self.__executorModules = {}
        self.__codeModules = {}
        self.__minds = {}
        self.__loader = ModuleLoader("Executor", PathFinder.getExecutorSection, ExecutorModule)

    def loadModules(self, modulesList, hideExceptions=False):
        """
        Load all modules required in moduleList
        """
        result = self.__loader.loadModules(modulesList, hideExceptions=hideExceptions)
        if not result["OK"]:
            return result
        self.__executorModules = self.__loader.getModules()
        return S_OK()

    # Go!
    def go(self):
        for name in self.__executorModules:
            exeClass = self.__executorModules[name]["classObj"]
            result = exeClass._ex_initialize(name, self.__executorModules[name]["loadName"])
            if not result["OK"]:
                return result
            mind = exeClass.ex_getMind()
            if mind not in self.__minds:
                self.__minds[mind] = self.MindCluster(mind, self.__aliveLock)
            mc = self.__minds[mind]
            mc.addModule(name, exeClass)
        for mindName in self.__minds:
            gLogger.info("Trying to connect to %s" % mindName)
            result = self.__minds[mindName].connect()
            if not result["OK"]:
                return result
        self.__aliveLock.lockUntilAllDead()
        return S_OK()
