#!/usr/bin/env arista-python
# Copyright (c) 2009,2010 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# pylint: disable=consider-using-f-string

import argparse, sys, os, time
import ManagedSubprocess

schedulerProcSig = '%s\0schedule\0' % os.path.basename( sys.argv[ 0 ] )

def listScheduled( args, parser ):
   rowFormat = "%-5s %s"
   print( rowFormat % ( "PID", "ARGS" ) )
   for pid in filter( str.isdigit, os.listdir( '/proc' ) ):
      pid = int( pid )
      try:
         if pid != os.getpid():
            with open( '/proc/%d/cmdline' % pid ) as f:
               l = f.read()
            if schedulerProcSig in l:
               args = " ".join( l.split( '\0' )[ 3: ] )
               print( rowFormat % ( pid, args ) )
      except OSError:
         # This could happen if the process disappeared as we were trying to
         # read its /proc/*/cmdline.
         pass
   
def daemonize():
   if os.getuid() != 0:
      # We need to be root in order to write to /var/log/...
      sys.stderr.write( "Must run as root to daemonize\n" )
      sys.exit( 1 )
   processStdoutDumpDir = "/var/log/ArSchedule"
   if not os.path.isdir( processStdoutDumpDir ):
      os.makedirs( processStdoutDumpDir )
   # do double fork to become a daemon
   pid = os.fork()
   if pid:
      os._exit( 0 ) # pylint: disable=protected-access
   os.setsid()
   pid = os.fork()
   if pid:
      os._exit( 0 ) # pylint: disable=protected-access

   myPid = os.getpid()
   print( "Successfully daemonized with pid %s" % ( myPid ) )
   pmDumpFile = os.path.join( processStdoutDumpDir, "ArSchedule-%d" % myPid )
   os.close( 0 )
   fd = os.open( "/dev/null", os.O_RDONLY )
   if fd != 0:
      os.dup2( fd, 0 )
      os.close( fd )

   f = open( pmDumpFile, "w" ) # pylint: disable=consider-using-with
   os.dup2( f.fileno(), 1 )
   os.dup2( f.fileno(), 2 )
   f.close()

   
def chunk( arr, chunk_size ):
   assert len( arr ) % chunk_size == 0
   return list( zip( *[ iter( arr ) ]*chunk_size ) )

def parseCommands( commands, parser ):
   # commands should be a list of the form:
   #  [ ( delay, command ), ... ]
   if len( commands ) % 2 != 0:
      parser.error( "Incorrect number of arguments (must be even)." )

   delayCommandPairs = chunk( commands, 2 )
   parsedDelayCommandPairs = []
   for delay, command in delayCommandPairs:
      try:
         delay = int( delay )
      except ValueError:
         parser.error( "%s is not a valid delay.\n" % delay )
         
      parsedDelayCommandPairs.append( ( delay, command ) )

   return parsedDelayCommandPairs
   
def schedule( args, parser ):
   delayCommandPairs = parseCommands( args.commands, parser )

   if args.daemonize:
      daemonize()
   # we are now in the child
   startTime = time.time()

   # check that we can be found by listScheduled
   with open( '/proc/%d/cmdline' % os.getpid() ) as f:
      l = f.read()
   assert schedulerProcSig in l
   
   for delay, command in delayCommandPairs:
      delay = delay - ( time.time() - startTime )
      print( 'Waiting %s second(s), then running "%s"' % ( delay, command ) )
      if delay > 0:
         time.sleep( delay )
      argv = [ '/bin/bash', '-c', command ]
      p = ManagedSubprocess.Popen( argv, stdin=sys.stdin, stdout=sys.stdout,
                                   stderr=ManagedSubprocess.STDOUT )
      retVal = p.wait(block=True)
      assert retVal == 0

def main():
   parser = argparse.ArgumentParser()
   subparsers = parser.add_subparsers()

   schedule_parser = subparsers.add_parser( "schedule",
         help="Schedule commands to run at particular time offsets." )
   schedule_parser.add_argument( "-d", "--daemonize",
         action="store_true", default=False )
   schedule_parser.add_argument( 'commands', nargs='+', metavar='offset command' )
   schedule_parser.set_defaults( func=schedule )

   list_parser = subparsers.add_parser( "ls",
         help="List currently running jobs." )
   list_parser.set_defaults( func=listScheduled )

   args = parser.parse_args()
   args.func( args, parser )

if __name__ == "__main__":
   main()
