"""
    Workflow class is the main container of Steps and Modules
"""

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

import os
import xml.sax
import six

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

__RCSID__ = "$Id$"


class Workflow(AttributeCollection):

  def __init__(self, obj=None, name=None):
    """ Be aware that 1-st param is an obj not a name!!!!
        obj can me a string with XML representation or with filename
        also obj can be a Workflow or ParameterCollections
    """
    AttributeCollection.__init__(self)
    if (obj is None) or isinstance(obj, ParameterCollection):
      self.setName('notgiven')
      self.setType('')
      self.setDescrShort('')
      self.setDescription('')
      self.setOrigin('')
      self.setVersion(0.0)
      self.parameters = ParameterCollection(obj)
      self.step_instances = InstancesPool(self)
      self.step_definitions = DefinitionsPool(self)
      self.module_definitions = DefinitionsPool(self)

    elif isinstance(obj, Workflow):
      self.fromWorkflow(obj)
    elif isinstance(obj, six.string_types):
      self.parameters = ParameterCollection(None)
      self.step_instances = InstancesPool(self)
      self.step_definitions = DefinitionsPool(self)
      self.module_definitions = DefinitionsPool(self)
      # if obj is an XML string
      if obj.startswith('<'):
        fromXMLString(obj, self)
      else:
        fromXMLFile(obj, self)
    elif obj is not None:
      raise TypeError('Can not create object type ' + str(type(self)) + ' from the ' + str(type(obj)))

    if name:
      self.setName(name)
    self.workflow_commons = {}
    self.workflowStatus = S_OK()

  def fromWorkflow(self, obj):
    self.setName(obj.getName())
    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_definitions = DefinitionsPool(self, obj.module_definitions)
    self.step_instances = InstancesPool(self, obj.step_instances)
    self.step_definitions = DefinitionsPool(self, obj.step_definitions)

  def __str__(self):
    """Creates a string representation of itself
    """
    ret = str(self.getName()) + ':\n' + AttributeCollection.__str__(self) + self.parameters.__str__()
    ret = ret + str(self.step_definitions)
    ret = ret + str(self.step_instances)
    ret = ret + str(self.module_definitions)
    return ret

  def toXML(self):
    """Creates an XML representation of itself
    """
    # THIS is very important that Definitions should be written before instances
    ret = '<Workflow>\n'
    ret = ret + AttributeCollection.toXML(self)
    ret = ret + self.parameters.toXML()
    ret = ret + self.module_definitions.toXML()
    ret = ret + self.step_definitions.toXML()
    ret = ret + self.step_instances.toXML()
    ret = ret + '</Workflow>\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 addTool(self, name, tool):
    """ Add an object that will be available in all the modules to perform some operations.
        For example, a state reporting facility.
    """
    self.workflow_commons[name] = tool

  def addStep(self, step):
    # this is a VERY important piece of code
    # we have to join all Modules definition from all added steps in the single dictionary
    # and we have to share this dictionary between all included steps
    # we also have to check versions of the modules and instances
    for type_o in step.module_definitions:
      if type_o not in self.module_definitions:
        self.module_definitions.append(step.module_definitions[type_o])
    self.step_definitions.append(step)
    del step.module_definitions  # we need to clean up all unwanted definitions
    step.module_definitions = None
    return step

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

  def createStepInstance(self, type_o, name):
    """ Creates step instance of type 'type_o' with the name 'name'
    """
    if type_o in self.step_definitions:
      stepi = StepInstance(name, self.step_definitions[type_o])
      self.step_instances.append(stepi)
      return stepi
    else:
      raise KeyError('Can not find StepDefinition ' + type + ' to create StepInstrance ' + name)

  def removeStepInstance(self, name):
    self.step_instances[name].setParent(None)
    self.step_instances.delete(name)

  def updateParents(self):
    self.module_definitions.updateParents(self)
    self.step_instances.updateParents(self)
    self.step_definitions.updateParents(self)

  def resolveGlobalVars(self):
    """ This method will create global parameter list and then will resolve all instances of @{VARNAME}
    Be aware that parameters of that type are GLOBAL!!! are string and can not be dynamically change
    The scope: the resolution of that parameters apply from lower to upper object, for example if
    parameter use in module, then it checks module, then step, then workflow

    Comment: If varible linked it should not be used in a global list"""

    # reenforced global parameters on the level of Workflow
    if not self.parameters.find("PRODUCTION_ID"):
      self.parameters.append(Parameter("PRODUCTION_ID", "00000000", "string", "", "", True, False,
                                       "Transformation ID taken from the ProductionManager"))
    if not self.parameters.find("JOB_ID"):
      self.parameters.append(Parameter("JOB_ID", "00000000", "string", "", "", True, False,
                                       "Job ID within Tranformation taken from the Transformation Manager"))

    self.parameters.resolveGlobalVars()
    step_instance_number = 0
    for inst in self.step_instances:
      # for each step instance we can define STEP_NUMBER
      step_instance_number = step_instance_number + 1
      if not inst.parameters.find("STEP_NUMBER"):
        inst.parameters.append(Parameter("STEP_NUMBER", "%s" % step_instance_number, "string", "", "", True, False,
                                         "Number of the StepInstance within the Workflow"))
      if not inst.parameters.find("STEP_ID"):
        prod_ID = self.parameters.find("PRODUCTION_ID").getValue()
        job_ID = self.parameters.find("JOB_ID").getValue()
        inst.parameters.append(Parameter("STEP_ID", "%s_%s_%d" % (prod_ID, job_ID, step_instance_number),
                                         "string", "", "", True, False, "Step instance ID"))
      if not inst.parameters.find("STEP_INSTANCE_NAME"):
        inst.parameters.append(Parameter("STEP_INSTANCE_NAME", inst.getName(), "string", "", "", True, False,
                                         "Name of the StepInstance within the Workflow"))
      if not inst.parameters.find("STEP_DEFINITION_NAME"):
        inst.parameters.append(Parameter("STEP_DEFINITION_NAME", inst.getType(), "string", "", "", True, False,
                                         "Type of the StepInstance within the Workflow"))
      if not inst.parameters.find("JOB_ID"):
        inst.parameters.append(Parameter("JOB_ID", "", "string", "self", "JOB_ID", True, False,
                                         "Job ID within Tranformation taken from the Transformation Manager"))
      if not inst.parameters.find("PRODUCTION_ID"):
        inst.parameters.append(Parameter("PRODUCTION_ID", "", "string", "self", "PRODUCTION_ID", True, False,
                                         "Type of the StepInstance within the Workflow"))
      inst.resolveGlobalVars(self.step_definitions, self.parameters)

  def createCode(self, combine_steps=False):
    self.resolveGlobalVars()
    str = ''
    str = str + self.module_definitions.createCode()
    str = str + self.step_definitions.createCode()
    str = str + "\nclass job:\n"
    str = str + indent(1) + 'def execute(self):\n'
    # str=str+indent(2)+'# flush self.step_instances\n'
    str = str + self.step_instances.createCode()
    # it seems we do not need it on this level
    str = str + indent(2) + '# output assignment\n'
    for v in self.parameters:
      if v.isOutput():
        str = str + v.createParameterCode(2, 'self')

    str = str + '\nj=job()\n'
    str = str + self.parameters.createParametersCode(0, 'j')
    str = str + 'j.execute()'
    return str

  def execute(self):
    self.resolveGlobalVars()
    # define workflow attributes
    wf_exec_attr = {}  # dictionary with the WF attributes, used to resolve links to self.attrname
    for wf_parameter in self.parameters:
      # parameters shall see objects in the current scope order to resolve links
      if wf_parameter.preExecute():  # for parm which not just outputs
        # print 'Input', wf_parameter
        if wf_parameter.isLinked():
          # print "Workflow self." + wf_parameter.getName(), '=',
          # wf_parameter.getLinkedModule() + '.' + wf_parameter.getLinkedParameter()
          if wf_parameter.getLinkedModule() == 'self':
            # this is not suppose to happen
            print(
                "Warning! Job attribute ",
                wf_parameter.getName(),
                "refers to the attribute of the same workflow",
                wf_parameter.getLinkedParameter())
            wf_exec_attr[wf_parameter.getName()] = wf_exec_attr[wf_parameter.getLinkedParameter()]
          else:
            wf_exec_attr[wf_parameter.getName()] = wf_exec_attr[wf_parameter.getLinkedModule()
                                                                ][wf_parameter.getLinkedParameter()]
        else:
          # print "Workflow self." + wf_parameter.getName(), '=', wf_parameter.getValue()
          wf_exec_attr[wf_parameter.getName()] = wf_parameter.getValue()

      # Put all the workflow parameters into the workflow_commons dictionary
      self.workflow_commons[wf_parameter.getName()] = wf_parameter.getValue()

    self.module_definitions.loadCode()  # loading Module classes into current python scope

    # wf_exec_steps will be dictionary of dictionaries [step instance name][parameter name]
    # used as dictionary of step instances to carry parameters
    wf_exec_steps = {}
    # print 'Executing Workflow',self.getType()
    error_message = ''
    step_result = ''
    for step_inst in self.step_instances:
      step_inst_name = step_inst.getName()
      step_inst_type = step_inst.getType()
      wf_exec_steps[step_inst_name] = {}
      # print "WorkflowInstance creating Step instance ",step_inst_name," of type", step_inst_type
      for parameter in step_inst.parameters:
        if parameter.preExecute():
          # print '>> Input', parameter
          if parameter.isLinked():
            # print ">> StepInstance", step_inst_name + '.' + parameter.getName(),
            # '=', parameter.getLinkedModule() + '.' + parameter.getLinkedParameter()
            if parameter.getLinkedModule() == 'self':
              # tale value form the step_dict
              wf_exec_steps[step_inst_name][parameter.getName()] = wf_exec_attr[parameter.getLinkedParameter()]
            else:
              # print wf_exec_steps[parameter.getLinkedModule()].keys()
              wf_exec_steps[step_inst_name][parameter.getName()] = wf_exec_steps[parameter.getLinkedModule()
                                                                                 ][parameter.getLinkedParameter()]
          else:
            # print ">> StepInstance", step_inst_name + '.' + parameter.getName(), '=', parameter.getValue()
            wf_exec_steps[step_inst_name][parameter.getName()] = parameter.getValue()

        # In the step_commons all parameters are added, both Input and Output ones.
        step_inst.step_commons[parameter.getName()] = parameter.getValue()

      resolveVariables(wf_exec_steps[step_inst_name])
      # Set proper values for all Input Parameters
      for key, value in wf_exec_steps[step_inst_name].items():
        step_inst.step_commons[key] = value

      step_inst.setParent(self)
      step_inst.setWorkflowCommons(self.workflow_commons)

      result = step_inst.execute(wf_exec_steps[step_inst_name], self.step_definitions)
      if not result['OK']:
        if self.workflowStatus['OK']:
          error_message = result['Message']
        self.workflowStatus = S_ERROR(result['Message'])
        self.workflowStatus['Errno'] = result['Errno']

      step_result = result.get('Value', step_result)

    # now we need to copy output values to the STEP!!! parameters
    # print "WorkflowInstance output assignment"
    for wf_parameter in self.parameters:
      if wf_parameter.isOutput():
        if wf_parameter.isLinked():
          # print "WorkflowInstance  self." + wf_parameter.getName(), '=',
          # wf_parameter.getLinkedModule() + '.' + wf_parameter.getLinkedParameter()
          if wf_parameter.getLinkedModule() == 'self':
            # this is not suppose to happen
            print(
                "Warning! Workflow OUTPUT attribute ",
                wf_parameter.getName(),
                "refer on the attribute of the same workflow",
                wf_parameter.getLinkedParameter())
            wf_exec_attr[wf_parameter.getName()] = wf_exec_attr[wf_parameter.getLinkedParameter()]
          else:
            wf_exec_attr[wf_parameter.getName()] = wf_exec_steps[wf_parameter.getLinkedModule()
                                                                 ][wf_parameter.getLinkedParameter()]
        else:
          # it is also does not make sense - we can produce warning
          print(
              "Warning! Workflow OUTPUT attribute",
              wf_parameter.getName(),
              "assigned constant",
              wf_parameter.getValue())
          # print "WorkflowInstance  self."+ wf_parameter.getName(),'=',wf_parameter.getValue()
          wf_exec_attr[wf_parameter.getName()] = wf_parameter.getValue()
        setattr(self, wf_parameter.getName(), wf_exec_attr[wf_parameter.getName()])
    # Return the result of the first failed step or S_OK
    if not self.workflowStatus['OK']:
      # return S_ERROR( error_message )
      return self.workflowStatus
    return S_OK(step_result)


from DIRAC.Core.Workflow.WorkflowReader import WorkflowXMLHandler


def fromXMLString(xml_string, obj=None):
  # KGG !!! We need to reset Workflow if it exists
  handler = WorkflowXMLHandler(obj)
  xml.sax.parseString(xml_string, handler)
  return handler.root


def fromXMLFile(xml_file, obj=None):
  # KGG !!! We need to reset Workflow if it exists
  handler = WorkflowXMLHandler(obj)
  xml.sax.parse(xml_file, handler)
  return handler.root
