"""
Accounting agent to consume perfSONAR network metrics received via a message queue.

.. literalinclude:: ../ConfigTemplate.cfg
  :start-after: ##BEGIN NetworkAgent
  :end-before: ##END
  :dedent: 2
  :caption: NetworkAgent options

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

from datetime import datetime

from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.Core.Base.AgentModule import AgentModule
from DIRAC.AccountingSystem.Client.Types.Network import Network
from DIRAC.ConfigurationSystem.Client.Config import gConfig
from DIRAC.ConfigurationSystem.Client.CSAPI import CSAPI
from DIRAC.Resources.MessageQueue.MQCommunication import createConsumer

__RCSID__ = "$Id$"


class NetworkAgent (AgentModule):
  """
  AccountingSystem agent to processes messages containing perfSONAR network metrics.
  Results are stored in the accounting database.
  """

  #: default number of seconds after which network accounting
  #: objects are removed from the temporary buffer
  BUFFER_TIMEOUT = 3600

  def initialize(self):

    self.log = gLogger.getSubLogger('NetworkAgent')

    # API initialization is required to get an up-to-date configuration from the CS
    self.csAPI = CSAPI()
    self.csAPI.initialize()

    # temporary buffer for network accounting objects + some parameters
    self.buffer = {}  # { {addTime: datetime.now(), object: Network() }, ... }
    self.bufferTimeout = self.am_getOption('BufferTimeout', NetworkAgent.BUFFER_TIMEOUT)

    # internal list of message queue consumers
    self.consumers = []

    # host-to-dirac name dictionary
    self.nameDictionary = {}

    # statistics
    self.messagesCount = 0  # number of received messages
    self.messagesCountOld = 0  # previous number of received messages (used to check connection status)

    self.skippedMessagesCount = 0  # number of skipped messages (errors, unsupported metrics, etc.)
    self.PLRMetricCount = 0  # number of received packet-loss-rate metrics
    self.OWDMetricCount = 0  # number of received one-way-delay metrics
    self.skippedMetricCount = 0  # number of skipped metrics (errors, invalid data, etc.)
    self.insertedCount = 0  # number of properly inserted accounting objects
    self.removedCount = 0  # number of removed accounting objects (missing data)

    return S_OK()

  def finalize(self):
    '''
    Gracefully close all consumer connections and commit last records to the DB.
    '''

    for consumer in self.consumers:
      consumer.close()

    self.commitData()

    return S_OK()

  def execute(self):
    '''
    During each cycle update the internal host-to-dirac name dictionary,
    check the consumers status (restart them if necessary),
    commit data stored in the buffer and show statistics.
    '''

    self.updateNameDictionary()
    self.checkConsumers()
    self.commitData()
    self.showStatistics()

    return S_OK()

  def updateNameDictionary(self):
    '''
    Update the internal host-to-dirac name dictionary.
    '''

    result = gConfig.getConfigurationTree('/Resources/Sites', 'Network/', '/Enabled')
    if not result['OK']:
      self.log.error("getConfigurationTree() failed with message: %s" % result['Message'])
      return S_ERROR('Unable to fetch perfSONAR endpoints from CS.')

    tmpDict = {}
    for path, value in result['Value'].items():
      if value == 'True':
        elements = path.split('/')
        diracName = elements[4]
        hostName = elements[6]
        tmpDict[hostName] = diracName

    self.nameDictionary = tmpDict

  def checkConsumers(self):
    '''
    Check whether consumers exist and work properly.
    (Re)create consumers if needed.
    '''

    # recreate consumers if there are any problems
    if not self.consumers or self.messagesCount == self.messagesCountOld:

      for consumer in self.consumers:
        consumer.close()

      for uri in self.am_getOption('MessageQueueURI', '').replace(" ", "").split(','):
        result = createConsumer(uri, self.processMessage)
        if not result['OK']:
          self.log.error('Failed to create a consumer from URI: %s' % uri)
          continue
        else:
          self.log.info('Successfully created a consumer from URI: %s' % uri)

        self.consumers.append(result['Value'])

      if self.consumers:
        return S_OK('Successfully created at least one consumer')
      return S_ERROR('Failed to create at least one consumer')

    # if everything is OK just update the counter
    else:
      self.messagesCountOld = self.messagesCount

  def processMessage(self, headers, body):
    '''
    Process a message containing perfSONAR data and store the result in the Accounting DB.
    Supports packet-loss-rate and one-way-delay metrics send in raw data streams.

    Function is designed to be an MQConsumer callback function.
    '''

    self.messagesCount += 1
    metadata = {
        'SourceIP': body['meta']['source'],
        'SourceHostName': body['meta']['input_source'],
        'DestinationIP': body['meta']['destination'],
        'DestinationHostName': body['meta']['input_destination'],
    }

    try:
      metadata['Source'] = self.nameDictionary[body['meta']['input_source']]
      metadata['Destination'] = self.nameDictionary[body['meta']['input_destination']]
    except KeyError as error:
      # messages with unsupported source or destination host name can be safely skipped
      self.skippedMessagesCount += 1
      self.log.debug('Host "%s" does not exist in the host-to-dirac name dictionary (message skipped)' % error)
      return S_OK()

    metadataKey = ''
    for value in metadata.values():
      metadataKey += value

    timestamps = sorted(body['datapoints'])
    for timestamp in timestamps:
      try:
        date = datetime.utcfromtimestamp(float(timestamp))

        # create a key that allows to join packet-loss-rate and one-way-delay
        # metrics in one network accounting record
        networkAccountingObjectKey = "%s%s" % (metadataKey, str(date))

        # use existing or create a new temporary accounting
        # object to store the data in DB
        if networkAccountingObjectKey in self.buffer:
          net = self.buffer[networkAccountingObjectKey]['object']

          timeDifference = datetime.now() - self.buffer[networkAccountingObjectKey]['addTime']
          if timeDifference.total_seconds() > 60:
            self.log.warn('Object was taken from buffer after %s' % (timeDifference))
        else:
          net = Network()
          net.setStartTime(date)
          net.setEndTime(date)
          net.setValuesFromDict(metadata)

        # get data stored in metric
        metricData = body['datapoints'][timestamp]

        # look for supported event types
        if headers['event-type'] == 'packet-loss-rate':
          self.PLRMetricCount += 1
          if metricData < 0 or metricData > 1:
            raise Exception('Invalid PLR metric (%s)' % (metricData))

          net.setValueByKey('PacketLossRate', metricData * 100)
        elif headers['event-type'] == 'histogram-owdelay':
          self.OWDMetricCount += 1

          # calculate statistics from histogram
          OWDMin = 999999
          OWDMax = 0
          total = 0
          count = 0
          for value, items in metricData.items():
            floatValue = float(value)
            total += floatValue * items
            count += items
            OWDMin = min(OWDMin, floatValue)
            OWDMax = max(OWDMax, floatValue)
          OWDAvg = float(total) / count

          # skip metrics with invalid data
          if OWDAvg < 0 or OWDMin < 0 or OWDMax < 0:
            raise Exception('Invalid OWD metric (%s, %s, %s)' %
                            (OWDMin, OWDAvg, OWDMax))
          else:
            # approximate jitter value
            net.setValueByKey('Jitter', OWDMax - OWDMin)
            net.setValueByKey('OneWayDelay', OWDAvg)

        else:
          self.skippedMetricCount += 1
          continue

        self.buffer[networkAccountingObjectKey] = {'addTime': datetime.now(),
                                                   'object': net}

      # suppress all exceptions to protect the listener thread
      except Exception as e:
        self.skippedMetricCount += 1
        self.log.warn('Metric skipped because of an exception: %s' % e)

    return S_OK()

  def commitData(self):
    '''
    Iterates through all object in the temporary buffer and commit objects to DB
    if both packet-loss-rate and one-way-delay values are set.

    Objects in the buffer older than self.bufferTimeout seconds which still have
    missing data are removed permanently (a warning is issued).
    '''

    now = datetime.now()
    removed = False

    for key, value in self.buffer.items():
      result = value['object'].checkValues()
      if not result['OK']:
        if (now - value['addTime']).total_seconds() > self.bufferTimeout:
          del self.buffer[key]
          self.removedCount += 1
          removed = True
      else:
        value['object'].delayedCommit()
        del self.buffer[key]
        self.insertedCount += 1

    if removed:
      self.log.warn('Network accounting object(s) has been removed because of missing data')

    return S_OK()

  def showStatistics(self):
    ''' Display different statistics as info messages in the log file.
    '''

    self.log.info("\tReceived messages:           %s" % self.messagesCount)
    self.log.info("\tSkipped messages:            %s" % self.skippedMessagesCount)
    self.log.info("\tPacket-Loss-Rate metrics:    %s" % self.PLRMetricCount)
    self.log.info("\tOne-Way-Delay metrics:       %s" % self.OWDMetricCount)
    self.log.info("\tSkipped metrics:             %s" % self.skippedMetricCount)
    self.log.info("")
    self.log.info("\tObjects in the buffer:       %s" % len(self.buffer))
    self.log.info("\tObjects inserted to DB:      %s" % self.insertedCount)
    self.log.info("\tPermanently removed objects: %s" % self.removedCount)

    return S_OK()
