""" I wish who wrote this would have put some doc...

    IIUC this is a wrapper around the JobState object. It basically tries to cache
    everything locally instead of going to the DB.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import six
import copy
import time

from DIRAC.Core.Utilities import Time, DEncode
from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.WorkloadManagementSystem.Client.JobState.JobState import JobState
from DIRAC.WorkloadManagementSystem.Client.JobState.JobManifest import JobManifest


class CachedJobState(object):

  log = gLogger.getSubLogger("CachedJobState")

  def __init__(self, jid, skipInitState=False):
    self.dOnlyCache = False
    self.__jid = jid
    self.__jobState = JobState(jid)
    self.cleanState(skipInitState=skipInitState)

  def cleanState(self, skipInitState=False):
    self.__cache = {}
    self.__jobLog = []
    self.__insertIntoTQ = False
    self.__dirtyKeys = set()
    self.__manifest = False
    self.__initState = None
    self.__lastValidState = time.time()
    if not skipInitState:
      result = self.getAttributes(["Status", "MinorStatus", "LastUpdateTime"])
      if result['OK']:
        self.__initState = result['Value']
      else:
        self.__initState = None

  def recheckValidity(self, graceTime=600):
    now = time.time()
    if graceTime <= 0 or now - self.__lastValidState > graceTime:
      self.__lastValidState = now
      result = self.__jobState.getAttributes(["Status", "MinorStatus", "LastUpdateTime"])
      if not result['OK']:
        return result
      currentState = result['Value']
      if not currentState == self.__initState:
        return S_OK(False)
      return S_OK(True)
    return S_OK(self.valid)

  @property
  def valid(self):
    return self.__initState is not None

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

  def getDirtyKeys(self):
    return set(self.__dirtyKeys)

  def commitChanges(self):
    if self.__initState is None:
      return S_ERROR("CachedJobState( %d ) is not valid" % self.__jid)
    changes = {}
    for k in self.__dirtyKeys:
      changes[k] = self.__cache[k]
    result = self.__jobState.commitCache(self.__initState, changes, self.__jobLog)
    try:
      result.pop('rpcStub')
    except KeyError:
      pass
    if not result['OK']:
      self.cleanState()
      return result
    if not result['Value']:
      self.cleanState()
      return S_ERROR("Initial state was different")
    newState = result['Value']
    self.__jobLog = []
    self.__dirtyKeys.clear()
    # Save manifest
    if self.__manifest and self.__manifest.isDirty():
      result = self.__jobState.setManifest(self.__manifest)
      if not result['OK']:
        self.cleanState()
        for _ in range(5):
          if self.__jobState.rescheduleJob()['OK']:
            break
        return result
      self.__manifest.clearDirty()
    # Insert into TQ
    if self.__insertIntoTQ:
      result = self.__jobState.insertIntoTQ()
      if not result['OK']:
        self.cleanState()
        for _ in range(5):
          if self.__jobState.rescheduleJob()['OK']:
            break
        return result
      self.__insertIntoTQ = False

    self.__initState = newState
    self.__lastValidState = time.time()
    return S_OK()

  def serialize(self):
    if self.__manifest:
      manifest = (self.__manifest.dumpAsCFG(), self.__manifest.isDirty())
    else:
      manifest = None
    return DEncode.encode((self.__jid, self.__cache, self.__jobLog, manifest,
                           self.__initState, self.__insertIntoTQ, tuple(self.__dirtyKeys)))

  @staticmethod
  def deserialize(stub):
    dataTuple, _slen = DEncode.decode(stub)
    if len(dataTuple) != 7:
      return S_ERROR("Invalid stub")
    # jid
    if not isinstance(dataTuple[0], six.integer_types):
      return S_ERROR("Invalid stub 0")
    # cache
    if not isinstance(dataTuple[1], dict):
      return S_ERROR("Invalid stub 1")
    # trace
    if not isinstance(dataTuple[2], list):
      return S_ERROR("Invalid stub 2")
    # manifest
    if dataTuple[3] is not None and (not isinstance(dataTuple[3], tuple) and len(dataTuple[3]) != 2):
      return S_ERROR("Invalid stub 3")
    # initstate
    if not isinstance(dataTuple[4], dict):
      return S_ERROR("Invalid stub 4")
    # Insert into TQ
    if not isinstance(dataTuple[5], bool):
      return S_ERROR("Invalid stub 5")
    # Dirty Keys
    if not isinstance(dataTuple[6], tuple):
      return S_ERROR("Invalid stub 6")
    cjs = CachedJobState(dataTuple[0], skipInitState=True)
    cjs.__cache = dataTuple[1]
    cjs.__jobLog = dataTuple[2]
    dt3 = dataTuple[3]
    if dataTuple[3]:
      manifest = JobManifest()
      result = manifest.loadCFG(dt3[0])
      if not result['OK']:
        return result
      if dt3[1]:
        manifest.setDirty()
      else:
        manifest.clearDirty()
      cjs.__manifest = manifest
    cjs.__initState = dataTuple[4]
    cjs.__insertIntoTQ = dataTuple[5]
    cjs.__dirtyKeys = set(dataTuple[6])
    return S_OK(cjs)

  def __cacheAdd(self, key, value):
    self.__cache[key] = value
    self.__dirtyKeys.add(key)

  def __cacheExists(self, keyList):
    if isinstance(keyList, six.string_types):
      keyList = [keyList]
    for key in keyList:
      if key not in self.__cache:
        return False
    return True

  def __cacheResult(self, cKey, functor, fArgs=None):
    # If it's a string
    if isinstance(cKey, six.string_types):
      if cKey not in self.__cache:
        if self.dOnlyCache:
          return S_ERROR("%s is not cached")
        if not fArgs:
          fArgs = tuple()
        result = functor(*fArgs)
        if not result['OK']:
          return result
        data = result['Value']
        self.__cache[cKey] = data
      return S_OK(self.__cache[cKey])
    # Tuple/List
    elif isinstance(cKey, (list, tuple)):
      if not self.__cacheExists(cKey):
        if self.dOnlyCache:
          return S_ERROR("%s is not cached")
        if not fArgs:
          fArgs = tuple()
        result = functor(*fArgs)
        if not result['OK']:
          return result
        data = result['Value']
        if len(cKey) != len(data):
          gLogger.warn(
              "CachedJobState.__memorize( %s, %s = %s ) doesn't receive the same amount of values as keys" %
              (cKey, functor, data))
          return data
        for i, val in enumerate(cKey):
          self.__cache[val] = data[i]
      # Prepare result
      return S_OK(tuple([self.__cache[cK] for cK in cKey]))
    else:
      raise RuntimeError("Cache key %s does not have a valid type" % cKey)

  def __cacheDict(self, prefix, functor, keyList=None):
    if not keyList or not self.__cacheExists(["%s.%s" % (prefix, key) for key in keyList]):
      result = functor(keyList)
      if not result['OK']:
        return result
      data = result['Value']
      for key in data:
        cKey = "%s.%s" % (prefix, key)
        # If the key is already in the cache. DO NOT TOUCH. User may have already modified it.
        # We update the coming data with the cached data
        if cKey in self.__cache:
          data[key] = self.__cache[cKey]
        else:
          self.__cache[cKey] = data[key]
      return S_OK(data)
    return S_OK(dict([(key, self.__cache["%s.%s" % (prefix, key)]) for key in keyList]))

  def _inspectCache(self):
    return copy.deepcopy(self.__cache)

  def _clearCache(self):
    self.__cache = {}

  @property
  def _internals(self):
    if self.__manifest:
      manifest = (self.__manifest.dumpAsCFG(), self.__manifest.isDirty())
    else:
      manifest = None
    return (self.__jid, self.dOnlyCache, dict(self.__cache),
            list(self.__jobLog), manifest, dict(self.__initState), list(self.__dirtyKeys))

#
# Manifest
#

  def getManifest(self):
    if not self.__manifest:
      result = self.__jobState.getManifest()
      if not result['OK']:
        return result
      self.__manifest = result['Value']
    return S_OK(self.__manifest)

  def setManifest(self, manifest):
    if not isinstance(manifest, JobManifest):
      jobManifest = JobManifest()
      result = jobManifest.load(str(manifest))
      if not result['OK']:
        return result
      manifest = jobManifest
    manifest.setDirty()
    self.__manifest = manifest
    # self.__manifest.clearDirty()
    return S_OK()

# Attributes
#

  def __addLogRecord(self, majorStatus=None, minorStatus=None, appStatus=None, source=None):
    record = {}
    if majorStatus:
      record['status'] = majorStatus
    if minorStatus:
      record['minor'] = minorStatus
    if appStatus:
      record['application'] = appStatus
    if not record:
      return
    if not source:
      source = "Unknown"
    self.__jobLog.append((record, Time.dateTime(), source))

  def setStatus(self, majorStatus, minorStatus=None, appStatus=None, source=None):
    self.__cacheAdd('att.Status', majorStatus)
    if minorStatus:
      self.__cacheAdd('att.MinorStatus', minorStatus)
    if appStatus:
      self.__cacheAdd('att.ApplicationStatus', appStatus)
    self.__addLogRecord(majorStatus, minorStatus, appStatus, source)
    return S_OK()

  def setMinorStatus(self, minorStatus, source=None):
    self.__cacheAdd('att.MinorStatus', minorStatus)
    self.__addLogRecord(minorStatus=minorStatus, source=source)
    return S_OK()

  def getStatus(self):
    return self.__cacheResult(('att.Status', 'att.MinorStatus'), self.__jobState.getStatus)

  def setAppStatus(self, appStatus, source=None):
    self.__cacheAdd('att.ApplicationStatus', appStatus)
    self.__addLogRecord(appStatus=appStatus, source=source)
    return S_OK()

  def getAppStatus(self):
    return self.__cacheResult('att.ApplicationStatus', self.__jobState.getAppStatus)
#
# Attribs
#

  def setAttribute(self, name, value):
    if not isinstance(name, six.string_types):
      return S_ERROR("Attribute name has to be a string")
    self.__cacheAdd("att.%s" % name, value)
    return S_OK()

  def setAttributes(self, attDict):
    if not isinstance(attDict, dict):
      return S_ERROR("Attributes has to be a dictionary and it's %s" % str(type(attDict)))
    for key in attDict:
      self.__cacheAdd("att.%s" % key, attDict[key])
    return S_OK()

  def getAttribute(self, name):
    return self.__cacheResult('att.%s' % name, self.__jobState.getAttribute, (name, ))

  def getAttributes(self, nameList=None):
    return self.__cacheDict('att', self.__jobState.getAttributes, nameList)

# JobParameters --- REMOVED

# Optimizer params

  def setOptParameter(self, name, value):
    if not isinstance(name, six.string_types):
      return S_ERROR("Optimizer parameter name has to be a string")
    self.__cacheAdd('optp.%s' % name, value)
    return S_OK()

  def setOptParameters(self, pDict):
    if not isinstance(pDict, dict):
      return S_ERROR("Optimizer parameters has to be a dictionary")
    for key in pDict:
      self.__cacheAdd('optp.%s' % key, pDict[key])
    return S_OK()

  def getOptParameter(self, name):
    return self.__cacheResult("optp.%s" % name, self.__jobState.getOptParameter, (name, ))

  def getOptParameters(self, nameList=None):
    return self.__cacheDict('optp', self.__jobState.getOptParameters, nameList)

# Other

  def resetJob(self, source=""):
    """ Reset the job!
    """
    return self.__jobState.resetJob(source=source)

  def getInputData(self):
    return self.__cacheResult("inputData", self.__jobState.getInputData)

  def insertIntoTQ(self):
    if self.valid:
      self.__insertIntoTQ = True
      return S_OK()
    return S_ERROR("Cached state is invalid")
