""" This object is a wrapper for setting and getting jobs states
"""

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

__RCSID__ = "$Id"

import six
import datetime

from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.WorkloadManagementSystem.Client.JobState.JobManifest import JobManifest
from DIRAC.WorkloadManagementSystem.Client import JobStatus
from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB, singleValueDefFields, multiValueDefFields
from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_GET_INFO, RIGHT_RESCHEDULE
from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_RESET, RIGHT_CHANGE_STATUS


class JobState(object):

  class DBHold(object):

    def __init__(self):
      self.checked = False
      self.reset()

    def reset(self):
      self.jobDB = None
      self.logDB = None
      self.tqDB = None

  __db = DBHold()

  @classmethod
  def checkDBAccess(cls):
    # Init DB if there
    if not JobState.__db.checked:
      JobState.__db.checked = True
      JobState.__db.jobDB = JobDB()
      JobState.__db.logDB = JobLoggingDB()
      JobState.__db.tqDB = TaskQueueDB()

  def __init__(self, jid, source="Unknown"):
    self.__jid = jid
    self.__source = str(source)
    self.checkDBAccess()

  @property
  def jid(self):
    return self.__jid

  def setSource(self, source):
    self.__source = source

  def getManifest(self, rawData=False):
    result = JobState.__db.jobDB.getJobJDL(self.__jid)
    if not result['OK'] or rawData:
      return result
    if not result['Value']:
      return S_ERROR("No manifest for job %s" % self.__jid)
    manifest = JobManifest()
    result = manifest.loadJDL(result['Value'])
    if not result['OK']:
      return result
    return S_OK(manifest)

  def setManifest(self, manifest):
    if not isinstance(manifest, JobManifest):
      manifestStr = manifest
      manifest = JobManifest()
      result = manifest.load(manifestStr)
      if not result['OK']:
        return result
    manifestJDL = manifest.dumpAsJDL()
    return self.__retryFunction(5, JobState.__db.jobDB.setJobJDL, (self.__jid, manifestJDL))

# Execute traces

  def __retryFunction(self, retries, functor, args=False, kwargs=False):
    retries = max(1, retries)
    if not args:
      args = tuple()
    if not kwargs:
      kwargs = {}
    while retries:
      retries -= 1
      result = functor(*args, **kwargs)
      if result['OK']:
        return result
      if retries == 0:
        return result
    return S_ERROR("No more retries")

  right_commitCache = RIGHT_GET_INFO

  def commitCache(self, initialState, cache, jobLog):
    try:
      self.__checkType(initialState, dict)
      self.__checkType(cache, dict)
      self.__checkType(jobLog, (list, tuple))
    except TypeError as excp:
      return S_ERROR(str(excp))
    result = self.getAttributes(list(initialState))
    if not result['OK']:
      return result
    if not result['Value'] == initialState:
      return S_OK(False)
    gLogger.verbose("Job %s: About to execute trace. Current state %s" % (self.__jid, initialState))

    data = {'att': [], 'jobp': [], 'optp': []}
    for key in cache:
      for dk in data:
        if key.find("%s." % dk) == 0:
          data[dk].append((key[len(dk) + 1:], cache[key]))

    if data['att']:
      attN = [t[0] for t in data['att']]
      attV = [t[1] for t in data['att']]
      result = self.__retryFunction(5, JobState.__db.jobDB.setJobAttributes,
                                    (self.__jid, attN, attV), {'update': True})
      if not result['OK']:
        return result

    if data['jobp']:
      result = self.__retryFunction(5, JobState.__db.jobDB.setJobParameters, (self.__jid, data['jobp']))
      if not result['OK']:
        return result

    for k, v in data['optp']:
      result = self.__retryFunction(5, JobState.__db.jobDB.setJobOptParameter, (self.__jid, k, v))
      if not result['OK']:
        return result

    if 'inputData' in cache:
      result = self.__retryFunction(5, JobState.__db.jobDB.setInputData, (self.__jid, cache['inputData']))
      if not result['OK']:
        return result

    gLogger.verbose("Adding logging records for %s" % self.__jid)
    for record, updateTime, source in jobLog:
      gLogger.verbose("Logging records for %s: %s %s %s" % (self.__jid, record, updateTime, source))
      record['date'] = updateTime
      record['source'] = source
      result = self.__retryFunction(5, JobState.__db.logDB.addLoggingRecord, (self.__jid,), record)
      if not result['OK']:
        return result

    gLogger.info("Job %s: Ended trace execution" % self.__jid)
    # We return a new initial state
    return self.getAttributes(list(initialState))
#
# Status
#

  def __checkType(self, value, tList, canBeNone=False):
    """ Raise TypeError if the value does not have one of the expected types

       :param value: the value to test
       :param tList: type or tuple of types
       :param canBeNone: boolean, since there is no type for None to be used with isinstance

    """
    if canBeNone:
      if value is None:
        return
    if not isinstance(value, tList):
      raise TypeError("%s has wrong type. Has to be one of %s" % (value, tList))

  right_setStatus = RIGHT_GET_INFO

  def setStatus(self, majorStatus, minorStatus=None, appStatus=None, source=None, updateTime=None):
    try:
      self.__checkType(majorStatus, six.string_types)
      self.__checkType(minorStatus, six.string_types, canBeNone=True)
      self.__checkType(appStatus, six.string_types, canBeNone=True)
      self.__checkType(source, six.string_types, canBeNone=True)
      self.__checkType(updateTime, datetime.datetime, canBeNone=True)
    except TypeError as excp:
      return S_ERROR(str(excp))
    result = JobState.__db.jobDB.setJobStatus(self.__jid,
                                              status=majorStatus,
                                              minorStatus=minorStatus,
                                              applicationStatus=appStatus)
    if not result['OK']:
      return result
    # HACK: Cause joblogging is crappy
    if not minorStatus:
      minorStatus = 'idem'
    if not source:
      source = self.__source
    return JobState.__db.logDB.addLoggingRecord(self.__jid,
                                                status=majorStatus,
                                                minorStatus=minorStatus,
                                                applicationStatus=appStatus,
                                                date=updateTime, source=source)

  right_getMinorStatus = RIGHT_GET_INFO

  def setMinorStatus(self, minorStatus, source=None, updateTime=None):
    try:
      self.__checkType(minorStatus, six.string_types)
      self.__checkType(source, six.string_types, canBeNone=True)
    except TypeError as excp:
      return S_ERROR(str(excp))
    result = JobState.__db.jobDB.setJobStatus(self.__jid, minorStatus=minorStatus)
    if not result['OK']:
      return result
    if not source:
      source = self.__source
    return JobState.__db.logDB.addLoggingRecord(self.__jid,
                                                minorStatus=minorStatus,
                                                date=updateTime, source=source)

  def getStatus(self):
    result = JobState.__db.jobDB.getJobAttributes(self.__jid, ['Status', 'MinorStatus'])
    if not result['OK']:
      return result
    data = result['Value']
    if data:
      return S_OK((data['Status'], data['MinorStatus']))
    return S_ERROR('Job %d not found in the JobDB' % int(self.__jid))

  right_setAppStatus = RIGHT_GET_INFO

  def setAppStatus(self, appStatus, source=None, updateTime=None):
    try:
      self.__checkType(appStatus, six.string_types)
      self.__checkType(source, six.string_types, canBeNone=True)
    except TypeError as excp:
      return S_ERROR(str(excp))
    result = JobState.__db.jobDB.setJobStatus(self.__jid, applicationStatus=appStatus)
    if not result['OK']:
      return result
    if not source:
      source = self.__source
    return JobState.__db.logDB.addLoggingRecord(self.__jid,
                                                applicationStatus=appStatus,
                                                date=updateTime, source=source)

  right_getAppStatus = RIGHT_GET_INFO

  def getAppStatus(self):
    result = JobState.__db.jobDB.getJobAttributes(self.__jid, ['ApplicationStatus'])
    if result['OK']:
      result['Value'] = result['Value']['ApplicationStatus']
    return result

# Attributes

  right_setAttribute = RIGHT_GET_INFO

  def setAttribute(self, name, value):
    try:
      self.__checkType(name, six.string_types)
      self.__checkType(value, six.string_types)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.setJobAttribute(self.__jid, name, value)

  right_setAttributes = RIGHT_GET_INFO

  def setAttributes(self, attDict):
    try:
      self.__checkType(attDict, dict)
    except TypeError as excp:
      return S_ERROR(str(excp))
    keys = [key for key in attDict]
    values = [attDict[key] for key in keys]
    return JobState.__db.jobDB.setJobAttributes(self.__jid, keys, values)

  right_getAttribute = RIGHT_GET_INFO

  def getAttribute(self, name):
    try:
      self.__checkType(name, six.string_types)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.getJobAttribute(self.__jid, name)

  right_getAttributes = RIGHT_GET_INFO

  def getAttributes(self, nameList=None):
    try:
      self.__checkType(nameList, (list, tuple), canBeNone=True)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.getJobAttributes(self.__jid, nameList)

# OptimizerParameters

  right_setOptParameter = RIGHT_GET_INFO

  def setOptParameter(self, name, value):
    try:
      self.__checkType(name, six.string_types)
      self.__checkType(value, six.string_types)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.setJobOptParameter(self.__jid, name, value)

  right_setOptParameters = RIGHT_GET_INFO

  def setOptParameters(self, pDict):
    try:
      self.__checkType(pDict, dict)
    except TypeError as excp:
      return S_ERROR(str(excp))
    for name in pDict:
      result = JobState.__db.jobDB.setJobOptParameter(self.__jid, name, pDict[name])
      if not result['OK']:
        return result
    return S_OK()

  right_removeOptParameters = RIGHT_GET_INFO

  def removeOptParameters(self, nameList):
    if isinstance(nameList, six.string_types):
      nameList = [nameList]
    try:
      self.__checkType(nameList, (list, tuple))
    except TypeError as excp:
      return S_ERROR(str(excp))
    for name in nameList:
      result = JobState.__db.jobDB.removeJobOptParameter(self.__jid, name)
      if not result['OK']:
        return result
    return S_OK()

  right_getOptParameter = RIGHT_GET_INFO

  def getOptParameter(self, name):
    try:
      self.__checkType(name, six.string_types)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.getJobOptParameter(self.__jid, name)

  right_getOptParameters = RIGHT_GET_INFO

  def getOptParameters(self, nameList=None):
    try:
      self.__checkType(nameList, (list, tuple), canBeNone=True)
    except TypeError as excp:
      return S_ERROR(str(excp))
    return JobState.__db.jobDB.getJobOptParameters(self.__jid, nameList)

# Other

  right_resetJob = RIGHT_RESCHEDULE

  def rescheduleJob(self, source=""):
    result = JobState.__db.tqDB.deleteJob(self.__jid)
    if not result['OK']:
      return S_ERROR("Cannot delete from TQ job %s: %s" % (self.__jid, result['Message']))
    result = JobState.__db.jobDB.rescheduleJob(self.__jid)
    if not result['OK']:
      return S_ERROR("Cannot reschedule in JobDB job %s: %s" % (self.__jid, result['Message']))
    JobState.__db.logDB.addLoggingRecord(self.__jid,
                                         status=JobStatus.RECEIVED,
                                         minorStatus='',
                                         applicationStatus='',
                                         source=source)
    return S_OK()

  right_resetJob = RIGHT_RESET

  def resetJob(self, source=""):
    result = JobState.__db.jobDB.setJobAttribute(self.__jid, "RescheduleCounter", -1)
    if not result['OK']:
      return S_ERROR("Cannot set the RescheduleCounter for job %s: %s" % (self.__jid, result['Message']))
    result = JobState.__db.tqDB.deleteJob(self.__jid)
    if not result['OK']:
      return S_ERROR("Cannot delete from TQ job %s: %s" % (self.__jid, result['Message']))
    result = JobState.__db.jobDB.rescheduleJob(self.__jid)
    if not result['OK']:
      return S_ERROR("Cannot reschedule in JobDB job %s: %s" % (self.__jid, result['Message']))
    JobState.__db.logDB.addLoggingRecord(self.__jid,
                                         status=JobStatus.RECEIVED,
                                         minorStatus='',
                                         applicationStatus='',
                                         source=source)
    return S_OK()

  right_getInputData = RIGHT_GET_INFO

  def getInputData(self):
    return JobState.__db.jobDB.getInputData(self.__jid)

  @classmethod
  def checkInputDataStructure(cls, pDict):
    if not isinstance(pDict, dict):
      return S_ERROR("Input data has to be a dictionary")
    for lfn in pDict:
      if 'Replicas' not in pDict[lfn]:
        return S_ERROR("Missing replicas for lfn %s" % lfn)
        replicas = pDict[lfn]['Replicas']
        for seName in replicas:
          if 'SURL' not in replicas or 'Disk' not in replicas:
            return S_ERROR("Missing SURL or Disk for %s:%s replica" % (seName, lfn))
    return S_OK()

  right_setInputData = RIGHT_GET_INFO

  def set_InputData(self, lfnData):
    result = self.checkInputDataStructure(lfnData)
    if not result['OK']:
      return result
    return JobState.__db.jobDB.setInputData(self.__jid, lfnData)

  right_insertIntoTQ = RIGHT_CHANGE_STATUS

  def insertIntoTQ(self, manifest=None):
    if not manifest:
      result = self.getManifest()
      if not result['OK']:
        return result
      manifest = result['Value']

    reqSection = "JobRequirements"

    result = manifest.getSection(reqSection)
    if not result['OK']:
      return S_ERROR("No %s section in the job manifest" % reqSection)
    reqCfg = result['Value']

    jobReqDict = {}
    for name in singleValueDefFields:
      if name in reqCfg:
        if name == 'CPUTime':
          jobReqDict[name] = int(reqCfg[name])
        else:
          jobReqDict[name] = reqCfg[name]

    for name in multiValueDefFields:
      if name in reqCfg:
        jobReqDict[name] = reqCfg.getOption(name, [])

    jobPriority = reqCfg.getOption('UserPriority', 1)

    result = self.__retryFunction(2, JobState.__db.tqDB.insertJob, (self.__jid, jobReqDict, jobPriority))
    if not result['OK']:
      errMsg = result['Message']
      # Force removing the job from the TQ if it was actually inserted
      result = JobState.__db.tqDB.deleteJob(self.__jid)
      if result['OK']:
        if result['Value']:
          gLogger.info("Job %s removed from the TQ" % self.__jid)
      return S_ERROR("Cannot insert in task queue: %s" % errMsg)
    return S_OK()
