""" Implementation of a step
"""
# pylint: disable=unused-wildcard-import,wildcard-import

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
import os
import time
import traceback
import sys

from DIRAC.Core.Workflow.Parameter import *
from DIRAC.Core.Workflow.Module import *
from DIRAC import S_OK, S_ERROR

__RCSID__ = "$Id$"


class StepDefinition(AttributeCollection):
    def __init__(self, step_type=None, obj=None, parent=None):
        AttributeCollection.__init__(self)
        self.module_instances = None
        # this object can be shared with the workflow
        # to if its =None and workflow!=None we have to
        # pass everything above
        self.module_definitions = None
        self.parent = parent

        # sort out Parameters and class attributes
        if (obj is None) or isinstance(obj, ParameterCollection):
            self.setType("notgiven")
            self.setDescrShort("")
            self.setDescription("")
            self.setOrigin("")
            self.setVersion(0.0)
            self.parameters = ParameterCollection(obj)  # creating copy
            self.module_instances = InstancesPool(self)
            self.module_definitions = DefinitionsPool(self)
        elif isinstance(obj, StepDefinition):
            self.setType(obj.getType())
            self.setDescrShort(obj.getDescrShort())
            self.setDescription(obj.getDescription())
            self.setOrigin(obj.getOrigin())
            self.setVersion(obj.getVersion())
            # copy instances and definitions
            self.parameters = ParameterCollection(obj.parameters)
            self.module_instances = InstancesPool(self, obj.module_instances)
            if obj.module_definitions is not None:
                self.module_definitions = DefinitionsPool(obj.module_definitions)
        else:
            raise TypeError("Can not create object type " + str(type(self)) + " from the " + str(type(obj)))
        if step_type:
            self.setType(step_type)

    def __str__(self):
        ret = str(type(self)) + ":\n" + AttributeCollection.__str__(self) + self.parameters.__str__()
        if self.module_definitions is not None:
            ret = ret + str(self.module_definitions)
        else:
            ret = ret + "Module definitions shared in Workflow\n"
        ret = ret + str(self.module_instances)
        return ret

    def toXML(self):
        ret = "<StepDefinition>\n"
        ret = ret + AttributeCollection.toXML(self)
        ret = ret + self.parameters.toXML()
        if self.module_definitions is not None:
            ret = ret + self.module_definitions.toXML()
        ret = ret + self.module_instances.toXML()
        ret = ret + "</StepDefinition>\n"
        return ret

    def toXMLFile(self, outFile):
        if os.path.exists(outFile):
            os.remove(outFile)
        with open(outFile, "w") as xmlfile:
            xmlfile.write(self.toXML())

    def addModule(self, module):
        # KGG We need to add code to update existing modules
        if self.module_definitions is None:
            self.parent.module_definitions.append(module)
        else:
            self.module_definitions.append(module)
        return module

    def createModuleInstance(self, module_type, name):
        """Creates module instance of type 'type' with the name 'name'"""

        if self.module_definitions[module_type]:
            mi = ModuleInstance(name, self.module_definitions[module_type])
            self.module_instances.append(mi)
            return mi
        else:
            raise KeyError("Can not find ModuleDefinition " + module_type + " to create ModuleInstrance " + name)

    def removeModuleInstance(self, name):
        """Remove module instance specified by its name"""
        self.module_instances.delete(name)

    def compare(self, s):
        """Custom Step comparison operation"""
        return AttributeCollection.compare(self, s) and self.module_instances.compare(s)

    def updateParent(self, parent):
        """ """
        # FIXME: no updateParents for AttributeCollection, What should this be?
        AttributeCollection.updateParents(self, parent)  # pylint: disable=no-member
        self.module_instances.updateParents(self)
        if self.module_definitions is not None:
            self.module_definitions.updateParents(self)

    def createCode(self):
        """Create Step code"""

        str = "class " + self.getType() + ":\n"
        str = str + indent(1) + "def execute(self):\n"
        str = str + self.module_instances.createCode()
        str = str + indent(2) + "# output assignment\n"
        for v in self.parameters:
            if v.isOutput():
                str = str + v.createParameterCode(2, "self")
        str += "\n"
        return str


class StepInstance(AttributeCollection):
    def __init__(self, name, obj=None, parent=None):
        AttributeCollection.__init__(self)
        self.parent = None

        if obj is None:
            self.parameters = ParameterCollection()
        elif isinstance(obj, StepInstance) or isinstance(obj, StepDefinition):
            if name is None:
                self.setName(obj.getName())
            else:
                self.setName(name)
            self.setType(obj.getType())
            self.setDescrShort(obj.getDescrShort())
            self.parameters = ParameterCollection(obj.parameters)
        elif (obj is None) or isinstance(obj, ParameterCollection):
            # set attributes
            self.setName(name)
            self.setType("")
            self.setDescrShort("")
            self.parameters = ParameterCollection(obj)
        elif obj is not None:
            raise TypeError("Can not create object type " + str(type(self)) + " from the " + str(type(obj)))

        self.step_commons = {}
        self.stepStatus = S_OK()

    def resolveGlobalVars(self, step_definitions, wf_parameters):
        """Resolve parameter values defined in the @{<variable>} form"""
        self.parameters.resolveGlobalVars(wf_parameters)
        module_instance_number = 0
        for inst in step_definitions[self.getType()].module_instances:
            module_instance_number = module_instance_number + 1
            if not inst.parameters.find("MODULE_NUMBER"):
                inst.parameters.append(
                    Parameter(
                        "MODULE_NUMBER",
                        "%s" % module_instance_number,
                        "string",
                        "",
                        "",
                        True,
                        False,
                        "ModuleInstance number within the Step",
                    )
                )
            if not inst.parameters.find("MODULE_INSTANCE_NAME"):
                inst.parameters.append(
                    Parameter(
                        "MODULE_INSTANCE_NAME",
                        inst.getName(),
                        "string",
                        "",
                        "",
                        True,
                        False,
                        "Name of the ModuleInstance within the Step",
                    )
                )
            if not inst.parameters.find("MODULE_DEFINITION_NAME"):
                inst.parameters.append(
                    Parameter(
                        "MODULE_DEFINITION_NAME",
                        inst.getType(),
                        "string",
                        "",
                        "",
                        True,
                        False,
                        "Type of the ModuleInstance within the Step",
                    )
                )
            if not inst.parameters.find("JOB_ID"):
                inst.parameters.append(
                    Parameter(
                        "JOB_ID", "", "string", "self", "JOB_ID", True, False, "Job ID within a Production as a string"
                    )
                )
            if not inst.parameters.find("PRODUCTION_ID"):
                inst.parameters.append(
                    Parameter(
                        "PRODUCTION_ID", "", "string", "self", "PRODUCTION_ID", True, False, "Production ID as a string"
                    )
                )
            if not inst.parameters.find("STEP_NUMBER"):
                inst.parameters.append(
                    Parameter(
                        "STEP_NUMBER",
                        "",
                        "string",
                        "self",
                        "STEP_NUMBER",
                        True,
                        False,
                        "Step instance number within the Workflow",
                    )
                )
            if not inst.parameters.find("STEP_ID"):
                inst.parameters.append(
                    Parameter(
                        "STEP_ID", "", "string", "self", "STEP_NUMBER", True, False, "Step ID within the Workflow"
                    )
                )
            inst.resolveGlobalVars(wf_parameters, self.parameters)

    def createCode(self, ind=2):
        """Create the Step code"""
        str = indent(ind) + self.getName() + " = " + self.getType() + "()\n"
        str = str + self.parameters.createParametersCode(ind, self.getName())
        str = str + indent(ind) + self.getName() + ".execute()\n\n"
        return str

    def __str__(self):
        """Step string representation"""
        return str(type(self)) + ":\n" + AttributeCollection.__str__(self) + self.parameters.__str__()

    def toXML(self):
        """Generate the Step XML representation"""
        ret = "<StepInstance>\n"
        ret = ret + AttributeCollection.toXML(self)
        ret = ret + self.parameters.toXML()
        ret = ret + "</StepInstance>\n"
        return ret

    def setWorkflowCommons(self, wf):
        """Add reference to the collection of the common tools"""

        self.workflow_commons = wf

    def execute(self, step_exec_attr, definitions):
        """Step execution method. step_exec_attr is array to hold parameters belong to this Step,
        filled above in the workflow
        """
        print("Executing StepInstance", self.getName(), "of type", self.getType(), list(definitions))
        # Report the Application state if the corresponding tool is supplied
        if "JobReport" in self.workflow_commons:
            if self.parent.workflowStatus["OK"]:
                self.workflow_commons["JobReport"].setApplicationStatus("Executing " + self.getName())

        # Prepare Step statistics evaluation
        self.step_commons["StartTime"] = time.time()
        self.step_commons["StartStats"] = os.times()

        step_def = definitions[self.getType()]
        step_exec_modules = {}
        error_message = ""
        error_code = 0
        for mod_inst in step_def.module_instances:
            mod_inst_name = mod_inst.getName()
            mod_inst_type = mod_inst.getType()

            # print "StepInstance creating module instance ", mod_inst_name, " of type", mod_inst.getType()
            step_exec_modules[mod_inst_name] = step_def.parent.module_definitions[
                mod_inst_type
            ].main_class_obj()  # creating instance

            # Resolve all the linked parameter values
            for parameter in mod_inst.parameters:
                if parameter.preExecute():
                    # print '>>>> Input', parameter
                    if parameter.isLinked():
                        # print ">>>> ModuleInstance", mod_inst_name + '.' + parameter.getName(),
                        # '=', parameter.getLinkedModule() + '.' + parameter.getLinkedParameter()
                        if parameter.getLinkedModule() == "self":
                            # tale value form the step_dict
                            setattr(
                                step_exec_modules[mod_inst_name],
                                parameter.getName(),
                                step_exec_attr[parameter.getLinkedParameter()],
                            )
                        else:

                            setattr(
                                step_exec_modules[mod_inst_name],
                                parameter.getName(),
                                getattr(step_exec_modules[parameter.getLinkedModule()], parameter.getLinkedParameter()),
                            )
                    else:
                        # print ">>>> ModuleInstance", mod_inst_name + '.' + parameter.getName(), '=', parameter.getValue()
                        setattr(step_exec_modules[mod_inst_name], parameter.getName(), parameter.getValue())
                    # print 'Step Input Parameter:', parameter.getName(), getattr(
                    # step_exec_modules[mod_inst_name], parameter.getName() )

            # Set reference to the workflow and step common tools
            setattr(step_exec_modules[mod_inst_name], "workflow_commons", self.parent.workflow_commons)
            setattr(step_exec_modules[mod_inst_name], "step_commons", self.step_commons)
            setattr(step_exec_modules[mod_inst_name], "stepStatus", self.stepStatus)
            setattr(step_exec_modules[mod_inst_name], "workflowStatus", self.parent.workflowStatus)

            try:
                result = step_exec_modules[mod_inst_name].execute()
                if not result["OK"]:
                    if self.stepStatus["OK"]:
                        error_message = result["Message"]
                        error_code = result["Errno"]
                        if "JobReport" in self.workflow_commons:
                            if self.parent.workflowStatus["OK"]:
                                self.workflow_commons["JobReport"].setApplicationStatus(error_message)
                    self.stepStatus = S_ERROR(result["Message"])
                else:
                    for parameter in mod_inst.parameters:
                        if parameter.isOutput():
                            # print '<<<< Output', parameter
                            if parameter.isLinked():
                                # print "ModuleInstance self ." + parameter.getName(), '=',
                                # parameter.getLinkedModule() + '.' + parameter.getLinkedParameter()
                                if parameter.getLinkedModule() == "self":
                                    # this is not supposed to happen
                                    print("Warning! Module OUTPUT attribute", parameter.getName(), end=" ")
                                    print(
                                        "refer to the attribute of the same module",
                                        parameter.getLinkedParameter(),
                                        "=",
                                        getattr(step_exec_modules[mod_inst_name], parameter.getName()),
                                    )
                                    step_exec_attr[parameter.getName()] = getattr(
                                        step_exec_modules[mod_inst_name],
                                        parameter.getLinkedParameter(),
                                        parameter.getValue(),
                                    )
                                    # print "                 OUT", parameter.getLinkedParameter(), '=',
                                    # getattr( step_exec_modules[mod_inst_name], parameter.getName(),
                                    # parameter.getValue() )
                                else:
                                    # print 'Output step_exec_attr', st_parameter.getName(),
                                    # step_exec_modules[st_parameter.getLinkedModule()],
                                    # parameter.getLinkedParameter()
                                    step_exec_attr[parameter.getName()] = getattr(
                                        step_exec_modules[parameter.getLinkedModule()], parameter.getLinkedParameter()
                                    )
                            else:
                                # This also does not make sense - we can give a warning
                                print("Warning! Module OUTPUT attribute ", parameter.getName(), end=" ")
                                print("assigned constant", parameter.getValue())
                                # print "StepInstance self." + parameter.getName(), '=', parameter.getValue()
                                step_exec_attr[parameter.getName()] = parameter.getValue()

                            # print 'Module Output Parameter:', parameter.getName(), step_exec_attr[parameter.getName()]

                    # Get output values to the step_commons dictionary
                    for key in result:
                        if key != "OK":
                            if key != "Value":
                                self.step_commons[key] = result[key]
                            elif isinstance(result["Value"], dict):
                                for vkey in result["Value"].keys():
                                    self.step_commons[vkey] = result["Value"][vkey]

            except Exception as x:
                print("Exception while module execution")
                print("Module", mod_inst_name, mod_inst.getType())
                print(str(x))
                exc = sys.exc_info()
                exc_type = exc[0]
                value = exc[1]
                print(
                    "== EXCEPTION ==\n%s: %s\n\n%s==============="
                    % (exc_type, value, "\n".join(traceback.format_tb(exc[2])))
                )

                print("Step status: ", self.stepStatus)
                print("Workflow status: ", self.parent.workflowStatus)
                if self.stepStatus["OK"]:
                    # This is the error that caused the workflow disruption
                    # report it to the WMS
                    error_message = "Exception while %s module execution: %s" % (mod_inst_name, str(x))
                    error_code = 0
                    if "JobReport" in self.workflow_commons:
                        if self.parent.workflowStatus["OK"]:
                            self.workflow_commons["JobReport"].setApplicationStatus(
                                "Exception in %s module" % mod_inst_name
                            )

                self.stepStatus = S_ERROR(error_code, error_message)

        # now we need to copy output values to the STEP!!! parameters
        for st_parameter in self.parameters:
            if st_parameter.isOutput():
                # print '<< Output', st_parameter
                if st_parameter.isLinked():
                    # print "StepInstance self." + st_parameter.getName(), '=',
                    # st_parameter.getLinkedModule() + '.' + st_parameter.getLinkedParameter()
                    if st_parameter.getLinkedModule() == "self":
                        # this is not supposed to happen
                        print("Warning! Step OUTPUT attribute", st_parameter.getName(), end=" ")
                        print(
                            "refer to the attribute of the same step",
                            st_parameter.getLinkedParameter(),
                            step_exec_attr[st_parameter.getLinkedParameter()],
                        )
                        step_exec_attr[st_parameter.getName()] = step_exec_attr[st_parameter.getLinkedParameter()]

                    else:
                        # print 'Output step_exec_attr', st_parameter.getName(),
                        # step_exec_modules[st_parameter.getLinkedModule()],
                        # st_parameter.getLinkedParameter()
                        step_exec_attr[st_parameter.getName()] = getattr(
                            step_exec_modules[st_parameter.getLinkedModule()], st_parameter.getLinkedParameter()
                        )
                    setattr(self, st_parameter.getName(), step_exec_attr[st_parameter.getName()])
                else:
                    # This also does not make sense - we can give a warning
                    print("Warning! Step OUTPUT attribute ", st_parameter.getName(), end=" ")
                    print("assigned constant", st_parameter.getValue())
                    # print "StepInstance self." + st_parameter.getName(), '=', st_parameter.getValue()
                    step_exec_attr[st_parameter.getName()] = st_parameter.getValue()

                print("Step Output", st_parameter.getName(), "=", step_exec_attr[st_parameter.getName()])

        # Return the result of the first failed module or S_OK
        if not self.stepStatus["OK"]:
            return S_ERROR(error_code, error_message)
        else:
            return S_OK(result["Value"])
