#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
#  =================================================================
#           #     #                 #     #
#           ##    #   ####   #####  ##    #  ######   #####
#           # #   #  #    #  #    # # #   #  #          #
#           #  #  #  #    #  #    # #  #  #  #####      #
#           #   # #  #    #  #####  #   # #  #          #
#           #    ##  #    #  #   #  #    ##  #          #
#           #     #   ####   #    # #     #  ######     #
#
#        ---   The NorNet Testbed for Multi-Homed Systems  ---
#                        https://www.nntb.no
#  =================================================================
#
#  High-Performance Connectivity Tracer (HiPerConTracer)
#  Copyright (C) 2015-2020 by Thomas Dreibholz
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
#  Contact: dreibh@simula.no

import ipaddress
import os
import ripe.atlas.cousteau
import sys
import time

import AtlasMNS
import AtlasMNSLogger


# ###### Schedule RIPE Atlas experiment #####################################
def scheduleRIPEAtlasExperiment(scheduledEntry):
   # ====== Create measurement ==============================================
   AtlasMNSLogger.info('ID #' + str(scheduledEntry['Identifier']) +
                       ': scheduling RIPE Atlas experiment ...')
   ( measurementID, cost, info ) = atlasMNS.createRIPEAtlasTracerouteMeasurement(
      int(scheduledEntry['ProbeID']),
      ipaddress.ip_address(scheduledEntry['AgentFromIP']),
      '托马斯\'s AtlasMNS Traceroute Experiment')

   # ====== Update state ====================================================
   scheduledEntry['ProbeCost'] = cost
   if measurementID != None:
      scheduledEntry['State']              = 'atlas_scheduled'
      scheduledEntry['ProbeMeasurementID'] = measurementID
   else:
      scheduledEntry['State'] = 'failed'
      scheduledEntry['Info']  = info
   atlasMNS.updateScheduledEntry(scheduledEntry)


# ###### Check RIPE Atlas experiment ########################################
def checkRIPEAtlasExperiment(scheduledEntry):
   # ====== Check measurement status ========================================
   (success, results) = atlasMNS.downloadRIPEAtlasMeasurementResults(scheduledEntry['ProbeMeasurementID'])
   if success == True:
      if len(results) > 0:
         # atlasMNS.printRIPEAtlasMeasurementResults(results)

         # ====== Handle results ============================================
         try:
            probeHostIP   = ipaddress.ip_address(results[0]['src_addr'])
            probeFromIP = ipaddress.ip_address(results[0]['from'])
            success = True
            scheduledEntry['ProbeHostIP'] = str(probeHostIP)
            scheduledEntry['ProbeFromIP'] = str(probeFromIP)

         except Exception as e:
            success = False
            scheduledEntry['Info'] = str(e)

         # ====== Update state ==============================================
         if success == True:
            scheduledEntry['State'] = 'agent_scheduled'
            AtlasMNSLogger.info('ID #' + str(scheduledEntry['Identifier']) +
                                ': finished RIPE Atlas Measurement #' + str(scheduledEntry['ProbeMeasurementID']) + ': ' +
                                'Probe #' + str(scheduledEntry['ProbeID']) + ' (' + str(scheduledEntry['ProbeHostIP']) + '/' + str(scheduledEntry['ProbeFromIP']) + ')' +
                                ' -> ' +
                                '(' +  str(scheduledEntry['AgentHostIP']) + '/' + str(scheduledEntry['AgentFromIP']) + ')')

         else:
            scheduledEntry['State'] = 'failed'
            AtlasMNSLogger.info('ID #' + str(scheduledEntry['Identifier']) +
                                ': RIPE Atlas Measurement #' +
                                str(scheduledEntry['ProbeMeasurementID']) + ' failed: ' +
                                str(scheduledEntry['Info']))
         atlasMNS.updateScheduledEntry(scheduledEntry)

      else:
         AtlasMNSLogger.trace('ID #' + str(scheduledEntry['Identifier']) +
                              ': RIPE Atlas Measurement #' +
                              str(scheduledEntry['ProbeMeasurementID']) + ' is still ongoing')


# ###### Finished experiment ################################################
def finished(scheduledEntry):
   # ====== Import RIPE Atlas results into results database =================
   # NOTE: We again download the results here, to avoid caching them after
   #       obtaining them to extract the ProbeHostIP and ProbeFromIP. The
   #       summary was not written before, since there was still no
   #       HiPerConTracer result available.
   (success, results) = atlasMNS.downloadRIPEAtlasMeasurementResults(scheduledEntry['ProbeMeasurementID'])
   if ((success == True) and (len(results) > 0)):
      if atlasMNS.importResults(scheduledEntry, results) == True:
         # ====== Update state ==============================================
         AtlasMNSLogger.info('ID #' + str(scheduledEntry['Identifier']) +
                             ': finished Agent run ' +
                             '(' +  str(scheduledEntry['AgentHostIP']) + '/' + str(scheduledEntry['AgentFromIP']) + ')'
                             ' -> ' +
                             'Probe #' + str(scheduledEntry['ProbeID']) + ' (' + str(scheduledEntry['ProbeHostIP']) + '/' + str(scheduledEntry['ProbeFromIP']) + ')')
         scheduledEntry['State'] = 'finished'
         atlasMNS.updateScheduledEntry(scheduledEntry)

      else:
         AtlasMNSLogger.trace('ID #' + str(scheduledEntry['Identifier']) +
                              ': unable to download results of RIPE Atlas Measurement #' +
                              str(scheduledEntry['ProbeMeasurementID']) + ' -> retrying later!')



# ###### Main program #######################################################

# ====== Initialise =========================================================
if len(sys.argv) == 1:
   configurationFile = os.path.expanduser('~/.atlasmns-database-configuration')
else:
   configurationFile = sys.argv[1]

atlasMNSLogger = AtlasMNSLogger.AtlasMNSLogger(AtlasMNSLogger.TRACE)
atlasMNS = AtlasMNS.AtlasMNS()
if not atlasMNS.loadConfiguration(configurationFile):
   sys.exit(1)

if not atlasMNS.connectToRIPEAtlas():
   sys.exit(1)

if not atlasMNS.connectToSchedulerDB():
   sys.exit(1)

if not atlasMNS.connectToResultsDB():
   sys.exit(1)


# ====== Main loop ==========================================================
AtlasMNSLogger.info('Scheduler is ready!')
while not AtlasMNS.breakDetected:

   # ====== Process schedule ================================================
   schedule = atlasMNS.querySchedule()
   for scheduledEntry in schedule:

      # ====== Check for shutdown ===========================================
      if AtlasMNS.breakDetected:
         break

      # ====== Process schedule entry =======================================
      # print(scheduledEntry)

      # ------ State == 'scheduled' -----------------------------------------
      state = scheduledEntry['State']
      if state == 'scheduled':
         scheduleRIPEAtlasExperiment(scheduledEntry)

      # ------ State == 'atlas_scheduled' -----------------------------------
      elif state == 'atlas_scheduled':
         checkRIPEAtlasExperiment(scheduledEntry)

      # ------ State == 'agent_scheduled' -----------------------------------
      elif state == 'agent_scheduled':
         # Nothing to do here!
         continue

      # ------ State == 'agent_completed' -----------------------------------
      elif state == 'agent_completed':
         finished(scheduledEntry)

      # ------ State == 'failed' --------------------------------------------
      elif state == 'failed':
         # Nothing to do here!
         continue

      # ------ State == 'finished' ------------------------------------------
      elif state == 'finished':
         # Nothing to do here!
         continue

      # ------ Bad State ----------------------------------------------------
      else:
         AtlasMNSLogger.error('Bad state for scheduled entry: ' + str(scheduledEntry))

   # ====== Wait ============================================================
   for i in range(10):
      if AtlasMNS.breakDetected:
         break
      time.sleep(1)


# ====== All done! ==========================================================
AtlasMNSLogger.info('Exiting!')
