#!/usr/bin/env python3
# Copyright (c) 2021 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

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

import argparse
import fileinput
import os
import sys

from ProcMgrLib import reloadProcMgr
import PyClient
from Tac import Timeout
import Tracing

traceHandle = Tracing.Handle( "SetHeartbeatPeriod" )
t0 = traceHandle.trace0

CONFIG_DIR = "/etc/ProcMgr.d/run"

def error( msg ):
   print( msg, file=sys.stderr )
   exit( 1 ) # pylint: disable=consider-using-sys-exit

def parseArgs():
   parser = argparse.ArgumentParser( \
      description="Change the heartbeat period of the given agent" )
   parser.add_argument( "agentName", metavar="AGENT", \
      help="name of agent to modify" )
   parser.add_argument( "period", metavar="PERIOD", type=float, \
      help="new heartbeat period (in seconds)" )
   parser.add_argument( "-s", "--sysname", metavar="SYSNAME", required=False, \
      default="ar", help="current sysname (default ar)" )
   parser.add_argument( "-d", "--configDir", required=False, \
      default="/etc/ProcMgr.d/run", dest="dir", \
      help="dir where ProcMgr stores config files (default /etc/ProcMgr.d/run)" )
   args = parser.parse_args()
   return ( args.agentName, args.period, args.sysname, args.dir )

def managedApproach( agentName, period, sysname ):
   """Tries to modify the AgentConfig entry for the agent in Sysdb. Returns True on
   success, False if the agent doesn't seem to be managed by Launcher."""

   try:
      pc = PyClient.PyClient( sysname, "Launcher", connectTimeout=10 )
      agentConfig = pc.root()[ sysname ][ "Sysdb" ][ "sys" ][ "config" ]\
         [ "agentConfigDir" ].agent[ agentName ]
      agentConfig.stable = False
      agentConfig.heartbeatPeriod = period
      agentConfig.stable = True
      return True
   except Timeout:
      print( "Unable to PyClient into Launcher. Trying unamanaged approach." )
      return False
   except PyClient.RpcError:
      t0( "Agent not managed by Launcher. Trying unmanaged approach." )
      return False

def unmanagedApproach( agentName, period ):
   """Directly modifies the agent config file under /etc/ProcMgr.d/run, if it exists.
   Triggers a warmstart in ProcMgr to get it to pick up the changes. Returns True on
   success and False if the agent doesn't have a config file."""

   nonLauncherConfigPath = os.path.join( CONFIG_DIR, agentName )
   if not os.path.exists( nonLauncherConfigPath ):
      # Agent doesn't have a config anywhere
      return False

   # Update config file
   updateFlag = False
   try:
      f = fileinput.input( nonLauncherConfigPath, inplace=True )
      for line in f:
         # Note that inplace=True temporarily redirects stdout to the open file
         if line.startswith( "heartbeatPeriod = " ):
            print( f"heartbeatPeriod = {period}" )
            updateFlag = True
         else:
            print( line, end="" )
   except OSError:
      error( "Permission to config file denied. Run `sudo {}` instead." \
         .format( " ".join( sys.argv ) ) )
   finally:
      f.close()

   if not updateFlag:
      error( "Found config file, but it didn't contain heartbeatPeriod." )

   # Reload ProcMgr to get it to pick up the new changes from the config file
   reloadProcMgr()
   return True

def main():
   agentName, period, sysname, cDir = parseArgs()
   global CONFIG_DIR
   CONFIG_DIR = cDir

   # Assume agent is managed by Launcher. Attempt to modify in Sysdb
   if managedApproach( agentName, period, sysname ):
      print( "Successfully set heartbeat period for {} to {}." \
         .format( agentName, period ) )
      return

   # Agent isn't managed by Launcher. Try to brute force the change
   if unmanagedApproach( agentName, period ):
      print( "Successfully set heartbeat period for {} to {}." \
         .format( agentName, period ) )
      return

   print( "{} doesn't have a config file anywhere. It may not be running." \
      .format( agentName ) )

if __name__ == "__main__":
   main()
