""" DIRAC Production DB

    Production database is used to collect and serve the necessary information
    in order to automate the task of transformation preparation for high level productions.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# # imports
import six
import json
import threading

from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.Core.Base.DB import DB
from DIRAC.ProductionSystem.Utilities.ProdValidator import ProdValidator
from DIRAC.ProductionSystem.Utilities.ProdTransManager import ProdTransManager
from DIRAC.Core.Utilities.List import intListToString


__RCSID__ = "$Id$"

MAX_ERROR_COUNT = 10

#############################################################################


class ProductionDB(DB):
  """ ProductionDB class
  """

  def __init__(self, dbname=None, dbconfig=None, dbIn=None):
    """ The standard constructor takes the database name (dbname) and the name of the
        configuration section (dbconfig)
    """
    if not dbname:
      dbname = 'ProductionDB'
    if not dbconfig:
      dbconfig = 'Production/ProductionDB'

    if not dbIn:
      DB.__init__(self, dbname, dbconfig)

    self.lock = threading.Lock()

    self.prodValidator = ProdValidator()
    self.ProdTransManager = ProdTransManager()

    self.PRODPARAMS = ['ProductionID',
                       'ProductionName',
                       'Description',
                       'CreationDate',
                       'LastUpdate',
                       'AuthorDN',
                       'AuthorGroup',
                       'Status']

    self.TRANSPARAMS = ['TransformationID',
                        'ProductionID',
                        'LastUpdate',
                        'InsertedTime']

    self.TRANSLINKSPARAMS = ['TransformationID',
                             'ParentTransformationID',
                             'ProductionID']

    self.PRODSTEPSPARAMS = ['StepID',
                            'Name',
                            'Description',
                            'LongDescription',
                            'Body',
                            'Type',
                            'Plugin',
                            'AgentType',
                            'GroupSize',
                            'InputQuery',
                            'OutputQuery',
                            'LastUpdate',
                            'InsertedTime']

    self.statusActionDict = {
        'New': None,
        'Active': 'startTransformation',
        'Stopped': 'stopTransformation',
        'Cleaned': 'cleanTransformation'}

  def addProduction(self, prodName, prodDescription, authorDN, authorGroup, connection=False):
    """ Create new production starting from its description

    :param str prodName: a string with the Production name
    :param str prodDescription: a json object with the Production description
    :param str authorDN: string with the author DN
    :param str authorGroup: string with author group
    """
    connection = self.__getConnection(connection)
    res = self._getProductionID(prodName, connection=connection)
    if res['OK']:
      return S_ERROR("Production with name %s already exists with ProductionID = %d" % (prodName,
                                                                                        res['Value']))
    elif res['Message'] != "Production does not exist":
      return res

    self.lock.acquire()

    req = "INSERT INTO Productions (ProductionName,Description,CreationDate,LastUpdate,\
                                    AuthorDN,AuthorGroup,Status)\
                                VALUES ('%s','%s',UTC_TIMESTAMP(),UTC_TIMESTAMP(),'%s','%s','New');" % \
        (prodName, prodDescription, authorDN, authorGroup)

    res = self._update(req, connection)
    if not res['OK']:
      self.lock.release()
      return res
    prodID = int(res['lastRowId'])
    self.lock.release()

    return S_OK(prodID)

  def getProductions(self, condDict=None, older=None, newer=None, timeStamp='LastUpdate',
                     orderAttribute=None, limit=None, offset=None, connection=False):
    """ Get parameters of all the Productions with support for the web standard structure

    :param dict condDict:
    :param str older:
    :param str newer:
    :param str timeStamp:
    :param str orderAttribute:
    :param int limit:
    :param int offset:
    :param bool connection:
    :return:
    """

    connection = self.__getConnection(connection)
    req = "SELECT %s FROM Productions %s" % (intListToString(self.PRODPARAMS),
                                             self.buildCondition(condDict, older, newer, timeStamp,
                                                                 orderAttribute, limit, offset=offset))
    res = self._query(req, connection)
    if not res['OK']:
      return res

    webList = []
    resultList = []
    for row in res['Value']:
      # Prepare the structure for the web
      rList = [str(item) if not isinstance(item, six.integer_types) else item for item in row]
      prodDict = dict(zip(self.PRODPARAMS, row))
      webList.append(rList)
      resultList.append(prodDict)
    result = S_OK(resultList)
    result['Records'] = webList
    result['ParameterNames'] = self.PRODPARAMS
    return result

  def getProduction(self, prodName, connection=False):
    """ Get the Production definition

    :param str prodName: the Production name or ID
    """
    res = self._getConnectionProdID(connection, prodName)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']
    res = self.getProductions(condDict={'ProductionID': prodID},
                              connection=connection)
    if not res['OK']:
      return res
    if not res['Value']:
      return S_ERROR("Production %s did not exist" % prodName)
    return S_OK(res['Value'][0])

  def getProductionParameters(self, prodName, parameters, connection=False):
    """ Get the requested parameters for a supplied production

    :param str prodName: the Production name or ID
    :param str parameters: any valid production parameter in self.PRODPARAMS
    """
    if isinstance(parameters, six.string_types):
      parameters = [parameters]
    res = self.getProduction(prodName, connection=connection)
    if not res['OK']:
      return res
    prodParams = res['Value']
    paramDict = {}
    for reqParam in parameters:
      if reqParam not in prodParams:
        return S_ERROR("Parameter %s not defined for production" % reqParam)
      paramDict[reqParam] = prodParams[reqParam]
    if len(paramDict) == 1:
      return S_OK(paramDict[reqParam])
    return S_OK(paramDict)

  ###########################################################################
  #
  # These methods manipulate the ProductionSteps table
  #

  def getProductionStep(self, stepID, connection=False):
    """ It returns the ProductionStep corresponding to the stepID

    :param int stepID: the ID of the production Step stored in the ProductionSteps table
    :return: the attributes of Production Step corresponding to the stepID
    """
    connection = self.__getConnection(connection)
    req = "SELECT %s FROM ProductionSteps WHERE StepID = %s" % (intListToString(self.PRODSTEPSPARAMS), str(stepID))
    res = self._query(req, connection)
    if not res['OK']:
      return res
    if not res['Value']:
      return S_ERROR("ProductionStep %s did not exist" % str(stepID))
    return S_OK(res['Value'][0])

  def addProductionStep(self, stepName, stepDescription, stepLongDescription, stepBody, stepType, stepPlugin,
                        stepAgentType, stepGroupSize, stepInputquery, stepOutputquery, connection=False):
    """ Add a Production Step

    :param str stepName: name of the production Step
    :param str stepDescription: description of the production Step
    :param str stepLongDescription: long description of the production Step
    :param str stepBody: body of the production Step
    :param str stepType: type of the production Step
    :param str stepPlugin: plugin to be used for the production Step
    :param str stepAgentType: agent type to be used for the production Step
    :param int stepGroupSize: group size of the production Step
    :param str stepInputquery: InputQuery of the production Step
    :param str stepOutputquery: OutputQuery of the production Step
    :return:
    """
    connection = self.__getConnection(connection)
    self.lock.acquire()
    req = "INSERT INTO ProductionSteps (Name,Description,LongDescription,Body,Type,Plugin,AgentType,GroupSize,\
                                    InputQuery,OutputQuery,LastUpdate,InsertedTime)\
                                VALUES ('%s','%s', '%s', %s, '%s', '%s', '%s', '%s', '%s', '%s',\
                                UTC_TIMESTAMP(),UTC_TIMESTAMP());" % \
          (stepName, stepDescription, stepLongDescription, stepBody, stepType, stepPlugin, stepAgentType, stepGroupSize,
           stepInputquery, stepOutputquery)

    res = self._update(req, connection)
    if not res['OK']:
      self.lock.release()
      return res
    stepID = int(res['lastRowId'])
    self.lock.release()

    return S_OK(stepID)

    ###########################################################################
  #
  # These methods manipulate the ProductionTransformations table
  #

  def getProductionTransformations(self, prodName, condDict=None, older=None, newer=None, timeStamp='CreationTime',
                                   orderAttribute=None, limit=None,
                                   offset=None, connection=False):
    """ Gets the transformations of a given Production

    :param str prodName: the Production name or ID
    :return:
    """
    res = self._getConnectionProdID(connection, prodName)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']
    condDict = {'ProductionID': prodID}

    req = "SELECT %s FROM ProductionTransformations %s" % (intListToString(self.TRANSPARAMS),
                                                           self.buildCondition(condDict, older, newer, timeStamp,
                                                                               orderAttribute, limit, offset=offset))
    res = self._query(req, connection)
    if not res['OK']:
      return res

    webList = []
    resultList = []
    for row in res['Value']:
      # Prepare the structure for the web
      rList = [str(item) if not isinstance(item, six.integer_types) else item for item in row]
      transDict = dict(zip(self.TRANSPARAMS, row))
      webList.append(rList)
      resultList.append(transDict)
    result = S_OK(resultList)
    result['Records'] = webList
    result['ParameterNames'] = self.TRANSPARAMS
    return result

  def __setProductionStatus(self, prodID, status, connection=False):
    """ Set the Status to the Production

    :param int prodID: the ProductionID
    :param str status: the Production status
    """
    req = "UPDATE Productions SET Status='%s', LastUpdate=UTC_TIMESTAMP() WHERE ProductionID=%d" % (status, prodID)
    return self._update(req, connection)

  # This is to be replaced by startProduction, stopProduction etc.
  def setProductionStatus(self, prodName, status, connection=False):
    """ Set the status to the production and to all the associated transformations

    :param str prodName: the Production name or ID
    :param str status: the Production status
    """
    res = self._getConnectionProdID(connection, prodName)
    gLogger.error(res)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']

    res = self.__setProductionStatus(prodID, status, connection=connection)
    if not res['OK']:
      return res

    res = self.ProdTransManager.executeActionOnTransformations(prodID, self.statusActionDict[status])
    if not res['OK']:
      gLogger.error(res['Message'])

    return S_OK()

  def startProduction(self, prodName, connection=False):
    """ Instantiate and start the transformations belonging to the production

    :param str prodName: the Production name or ID
    """

    res = self._getConnectionProdID(connection, prodName)
    gLogger.error(res)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']

    # Instantiate the transformations according to the description
    res = self.getProductionParameters(prodName, 'Description')
    if not res['OK']:
      return res

    prodDescription = json.loads(res['Value'])
    transIDs = []
    for step in sorted(prodDescription):
      res = self.ProdTransManager.addTransformationStep(prodDescription[step]['stepID'], prodID)
      # If the addition of at least one step failes, clean the already added transformation steps from the TS
      if not res['OK']:
        self.ProdTransManager.deleteTransformations(transIDs)
        return S_ERROR(res['Message'])
      transID = res['Value']
      prodDescription[step]['transID'] = transID
      transIDs.append(transID)

    for step in prodDescription:
      transID = prodDescription[step]['transID']
      parentTransIDs = []
      if 'parentStep' in prodDescription[step]:
        for parentStep in prodDescription[step]['parentStep']:
          parentTransID = prodDescription[parentStep]['transID']
          parentTransIDs.append(parentTransID)

      res = self.addTransformationsToProduction(prodID, transID, parentTransIDs=parentTransIDs)
      # If adding the transformations to production fails, clean all the transformation steps from the TS
      if not res['OK']:
        self.ProdTransManager.deleteTransformations(transIDs)
        # Clean the production
        self.deleteProduction(prodID)
        return S_ERROR(res['Message'])

    res = self.__setProductionStatus(prodID, 'Active', connection=connection)
    if not res['OK']:
      return res

    res = self.ProdTransManager.executeActionOnTransformations(prodID, self.statusActionDict['Active'])
    if not res['OK']:
      gLogger.error(res['Message'])

    return S_OK()

  def __deleteProduction(self, prodID, connection=False):
    """ Delete a given production from the Productions table

    :param int prodID: ProductionID
    """
    req = "DELETE FROM Productions WHERE ProductionID=%d;" % prodID
    return self._update(req, connection)

  def __deleteProductionTransformations(self, prodID, connection=False):
    """ Remove all the transformations of the specified production from the TS and from the PS

    :param int prodID: the ProductionID
    """

    # Remove transformations from the TS
    gLogger.notice("Deleting transformations of Production %s from the TS" % prodID)
    res = self.ProdTransManager.deleteProductionTransformations(prodID)
    if not res['OK']:
      gLogger.error("Failed to delete production transformations from the TS", res['Message'])

    # Remove transformations from the PS
    req = "DELETE FROM ProductionTransformationLinks WHERE ProductionID = %d;" % prodID
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("Failed to delete production transformation links from the PS", res['Message'])

    req = "DELETE FROM ProductionTransformations WHERE ProductionID = %d;" % prodID
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("Failed to delete production transformations from the PS", res['Message'])

    return res

  def cleanProduction(self, prodName, author='', connection=False):
    """ Clean the production specified by name or id

    :param str prodName: the Production name or ID
    """
    res = self._getConnectionProdID(connection, prodName)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']
    res = self.__deleteProductionTransformations(prodID, connection=connection)
    if not res['OK']:
      return res

    return S_OK(prodID)

  def deleteProduction(self, prodName, author='', connection=False):
    """ Remove the production specified by name or id

    :param str prodName: the Production name or ID
    """
    res = self._getConnectionProdID(connection, prodName)
    gLogger.error(res)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']
    res = self.cleanProduction(prodID, author=author, connection=connection)
    if not res['OK']:
      return res
    res = self.__deleteProduction(prodID, connection=connection)
    if not res['OK']:
      return res

    return S_OK()

  def addTransformationsToProduction(self, prodName, transIDs, parentTransIDs=None, connection=False):
    """ Check the production validity and the add the transformations to the production

    :param str prodName: the Production name or ID
    :param list transIDs: the list of transformations to be added to the production
    """
    gLogger.info(
        "ProductionDB.addTransformationsToProduction: \
         Attempting to add %s transformations to production: %s" %
        (transIDs, prodName))

    if not transIDs:
      return S_ERROR('Zero length transformation list')
    res = self._getConnectionProdID(connection, prodName)
    if not res['OK']:
      return res
    connection = res['Value']['Connection']
    prodID = res['Value']['ProductionID']

    # Check the production validity
    gLogger.notice('Checking if production is valid')

    # Set transformations in list format
    if parentTransIDs:
      if not isinstance(parentTransIDs, list):
        parentTransIDs = [parentTransIDs]

    if not isinstance(transIDs, list):
      transIDs = [transIDs]

    # Check the status of the transformations (must be 'New')
    for transID in transIDs:
      res = self.prodValidator.checkTransStatus(transID)
      if not res['OK']:
        gLogger.error("Production is not valid:", res['Message'])
        return res

    # Check if transformations are linked to their parent transformations, if they have one(s)
    if parentTransIDs:
      for transID in transIDs:
        for parentTransID in parentTransIDs:
          gLogger.notice(
              'Checking if transformation %s is linked to the parent transformation %s' %
              (transID, parentTransID))
          res = self.prodValidator.checkTransDependency(transID, parentTransID)
          if not res['OK']:
            gLogger.error("Production is not valid:", res['Message'])
            return res

    gLogger.notice('Production %s is valid' % prodName)

    res = self.__addTransformations(prodID, transIDs, connection=connection)
    if not res['OK']:
      msg = "Failed to add transformations %s to production %s: %s" % (transIDs, prodID, res['Message'])
      return S_ERROR(msg)
    # Add the transformation links (tranformation to parent transformation) for the given production
    res = self.__addTransformationLinks(prodID, transIDs, parentTransIDs=parentTransIDs, connection=connection)
    if not res['OK']:
      msg = "Failed to add production transformations links to transformations %s: %s" % (transIDs, res['Message'])
      return S_ERROR(msg)

    # Update the status of the transformation to be in sync with the status of the production
    res = self.getProduction(prodID)
    prodStatus = res['Value']['Status']

    # Execute action on transformations according to the production status
    res = self.ProdTransManager.executeActionOnTransformations(prodID, self.statusActionDict[prodStatus])
    if not res['OK']:
      gLogger.error(res['Message'])

    return S_OK()

  def __addTransformationLinks(self, prodID, transIDs, parentTransIDs=None, connection=False):
    """ Insert the transformations in the ProductionTransformationLinks table

    :param prodName: the Production name or ID
    :param transIDs: the list of transformations to be associated ProductionTransformationLinks table
    """
    req = "INSERT INTO ProductionTransformationLinks \
           (TransformationID,ParentTransformationID,ProductionID) VALUES"

    # Insert transformations and the corresponding parent transformations
    if parentTransIDs:
      for transID in transIDs:
        for parentTransID in parentTransIDs:
          req = "%s (%d,%d,%d)," % (req, transID, parentTransID, prodID)
    # If parent transformations are not defined, just insert transformations and use the parent transformation default
    # value -1
    else:
      req = "INSERT INTO ProductionTransformationLinks (TransformationID,ProductionID) VALUES"
      for transID in transIDs:
        req = "%s (%d,%d)," % (req, transID, prodID)

    gLogger.notice(req)
    req = req.rstrip(',')
    res = self._update(req, connection)
    if not res['OK']:
      return res

    return S_OK()

  def __addTransformations(self, prodID, transIDs, connection=False):
    """ Insert the transformations in the ProductionTransformations table

    :param str prodName: the Production name or ID
    :param list transIDs: the list of transformations to be added to the production
    """
    req = "INSERT INTO ProductionTransformations \
           (ProductionID,TransformationID,LastUpdate,InsertedTime) VALUES"
    for transID in transIDs:
      req = "%s (%d,%d,UTC_TIMESTAMP(),UTC_TIMESTAMP())," % (req, prodID, transID)
      gLogger.notice(req)
    req = req.rstrip(',')
    res = self._update(req, connection)
    if not res['OK']:
      return res

    return S_OK()

  def _getProductionID(self, prodName, connection=False):
    """ Method returns ID of production specified by the prodName

    :param str prodName: the Production name
    """
    try:
      prodName = int(prodName)
      cmd = "SELECT ProductionID from Productions WHERE ProductionID=%d;" % prodName
    except BaseException:
      if not isinstance(prodName, six.string_types):
        return S_ERROR("Production should be ID or name")
      cmd = "SELECT ProductionID from Productions WHERE ProductionName='%s';" % prodName
    res = self._query(cmd, connection)
    if not res['OK']:
      gLogger.error("Failed to obtain production ID for production", "%s: %s" % (prodName, res['Message']))
      return res
    elif not res['Value']:
      gLogger.verbose("Production %s does not exist" % (prodName))
      return S_ERROR("Production does not exist")
    return S_OK(res['Value'][0][0])

  def __getConnection(self, connection):
    """ Get the MySQL connection

    :param bool connection: the DB connection
    """
    if connection:
      return connection
    res = self._getConnection()
    if res['OK']:
      return res['Value']
    gLogger.warn("Failed to get MySQL connection", res['Message'])
    return connection

  def _getConnectionProdID(self, connection, prodName):
    """ Get the Production ID for a given production

    :param bool connection: the DB connection
    :param prodName: the Production name
    """
    connection = self.__getConnection(connection)
    res = self._getProductionID(prodName, connection=connection)
    if not res['OK']:
      gLogger.error("Failed to get ID for production %s: %s" % (prodName, res['Message']))
      return res
    prodID = res['Value']
    resDict = {'Connection': connection, 'ProductionID': prodID}
    return S_OK(resDict)
