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

"""Loads the fragment configuration files from the managed-config directory
at startup.  This script is run by EosInit after it is sure that Sysdb has finished
reading startup-config."""

import argparse
import os
import sys
from tempfile import NamedTemporaryFile

import Ark
from CliCommon import MAX_PRIV_LVL
from FileUrlDefs import MANAGED_CONFIG_DIR
# import Logging
import Tac
import Tracing
from Url import parseUrl
from UrlPlugin.FlashUrl import flashFileUrl

traceHandle = Tracing.defaultTraceHandle()
t0 = traceHandle.trace0
t1 = traceHandle.trace1
t5 = traceHandle.trace5
t9 = traceHandle.trace9

def runLoadFragments( options ):
   # This will only be called by EosInit/SysdbInit *after*
   # startup-config has been loaded.
   errors = []
   # Still need to define log msgs, but as a placeholder:
   Ark.configureLogManager( "LoadFragmentFiles" )
   t0( 'LoadFragmentFiles started' )

   sysname = options.sysname
   timeout = os.environ.get( "EOSINIT_CONFIG_AGENT_TIMEOUT" )
   timeoutArg = [ "--timeout=" + timeout ] if timeout is not None else [ ]
   fragmentStatus = None
   def wfConfigAgent( systemName, targ ):
      try:
         Tac.run( [ sys.executable, "/usr/bin/wfw", "--sysname", systemName,
                    "ConfigAgent" ] + targ )
      except Tac.SystemCommandError as e: # pylint: disable=unused-variable
         sys.stderr.write( "Warning: ConfigAgent failed to warm up; "
                           "cannot load fragment config files.\n" )
         raise
   def loadOneFragment( fragment, systemName, opt ):
      output = ""
      t5( f"Loading fragment {fragment} ..." )
      filename = fragmentConfigDir + "/" + fragment + ".fragment"
      try:
         filenameUrl = parseUrl( filename, None )
         # When --startup-config, is passed, we cannot execute
         # commands from enable mode --- and -c doesn't allow more
         # than one command. So, we need to make sure that any
         # command we want to use is in GlobalConfigMode (see
         # CliFragment/CliPlugin/CliFragmentCli.py) *and* create
         # a temp file with the commands we want to run:
         # pylint: disable-next=consider-using-with
         tempfile = NamedTemporaryFile( mode='w+t' )
         tempfile.write( f'configure fragment prepare {fragment} replace\n' )
         tempfile.write( f'configure fragment copy {filenameUrl} {fragment}\n' )
         tempfile.seek( 0 )
         t9( f"Storing configure fragment commands in {tempfile.name}" )
         cmd = [ sys.executable,
                 '/usr/bin/CliShell',
                 '--disable-aaa', '--disable-guards',
                 '-p', str( MAX_PRIV_LVL ),
                 # This flag lets the CliPlugins know that agents are not yet
                 # running.
                 # There are some Cli commands that special case their
                 # operation when being run as part of the startup-config.
                 # This is a similar situation, so we set the flag. However,
                 # it may not be necessary here because fragment-configs are
                 # being loaded into a session, and not directly into Sysdb,
                 # so they should already be aware that no agent will react
                 # to their configuration.
                 '--startup-config',
                 '--sysname', systemName,
                 tempfile.name
               ]
         if opt.timeout:
            cmd += [ '-T', str( opt.timeout ) ]
         # Wait for fragment file to be loaded into pending session for frag
         output = Tac.run( cmd, stdout=Tac.CAPTURE, stderr=Tac.CAPTURE )
         if output:
            msg = f"Unexpected output loading file for frag {fragment}: {output}"
            t0( msg )
            errors.append( msg )
         # Save it in committed version, even if it had errors
         expr = f"CF.commitFragmentInternal( '{frag}' )"
         result = pc.eval( expr )
         # This eval should never error --- even if there are errors in the
         # config file
         assert not result, f"Expected boolean false, but got {result}"
         loadedFragments.add( frag )
         t5( f"... loaded fragment {frag}" )

      except Tac.SystemCommandError as e:
         t1( f"CliShell ERROR: return code {e.error} output:{e.output}" )
         output = e.output
         if ( "Unable to connect" in output or
              "Cannot connect to" in output ):
            output += '\nCLI could not connect to ConfigAgent'
         msg = f"Error while running shell for {filename}: {output}"
         t1( msg )
         errors.append( msg )
      except Exception as e: # pylint: disable=broad-except
         msg = ( f"Exception while reading {filename}. " +
                 f"error:{e}, of class{ type( e ) }: {str(e)}" )
         t1( msg )
         errors.append( msg )
      finally:
         if tempfile:
            t9( f"Closing {tempfile.name}" )
            tempfile.close()

   try:
      wfConfigAgent( sysname, timeoutArg )
      import PyClient # pylint: disable=import-outside-toplevel
      from PyRpc import Rpc # pylint: disable=import-outside-toplevel
      pc = PyClient.PyClient( sysname, "ConfigAgent",
                              initConnectCmd="import CliFragment as CF",
                              execMode=Rpc.execModeThreadPerConnection )
      t9( "Connected to ConfigAgent" )
      sysdbAgent = 'Sysdb'
      sysdbStatus = pc.root()[ sysname ][ sysdbAgent][ "Sysdb" ][ "status" ]
      sysdbRoot = pc.root()[ sysname ][ 'Sysdb' ]
      # Should not load fragment configs until after startup-config has been
      # processed. But don't wait for sysdbStatus.initialized, to avoid
      # circular dependency
      Tac.waitFor( lambda: ( sysdbStatus.startupConfigStatus == "completed" ),
                   description="startup-config to be loaded" )
      fragmentStatus = sysdbRoot[ "cli" ][ "fragment" ][ "status" ]
      assert fragmentStatus
      # Really, fragmentStatus.initialized should be true if wfw returned
      # for ConfigAgent, but wait(), to confirm that.
      Tac.waitFor( lambda: fragmentStatus.initialized,
                   description=( "CliFragmentMgr for CliFragment's to be " +
                                 "ready in ConfigAgent" ) )

      t9( "Detecting saved fragment config files" )
      fragmentConfigDir = flashFileUrl( MANAGED_CONFIG_DIR )
      fragmentConfigDirUrl = parseUrl( fragmentConfigDir, None )
      if not fragmentConfigDirUrl.exists():
         fragmentStatus.fragmentConfigStatus = "completed"
         msg = f"Nonexistent fragment config dir: {fragmentConfigDirUrl}"
         t1( msg )
         return msg
      fragments = set([])
      loadedFragments = set([])
      fragmentStatus.fragmentConfigStatus = "inProgress"
      for filename in fragmentConfigDirUrl.listdir():
         if filename.endswith( ".fragment" ):
            frag = filename[ 0 : ( - len( ".fragment" ) ) ]
            fragments.add( frag )
      t1( f"Found {len( fragments )} existing fragments ({fragments})" )
      for frag in fragments:
         loadOneFragment( frag, sysname, options )

   except Exception as e: # pylint: disable=broad-except
      msg = f"Exception raised in LoadFragmentFiles: {str(e)}"
      t1( msg )
      errors.append( msg )
   finally:
      if fragments == loadedFragments:
         stat = "completed"
      else:
         stat = "incomplete"
         errors.append( f"expected {fragments}, but loaded {loadedFragments}" )
      if fragmentStatus:
         # Write to cli/fragment/status that we're done loading the
         # startup config
         fragmentStatus.fragmentConfigErrorCount = len( errors )
         fragmentStatus.fragmentConfigStatus = stat
   return "\n".join( errors )

def main():
   parser = argparse.ArgumentParser()
   parser.add_argument( "-s", "--sysname", action="store", default="ar",
                        help="System name (default: %default)" )
   parser.add_argument( "-t", "--timeout", action="store", default=600,
                        type=int, help="Connection timeout (testing only)" )
   options = parser.parse_args()

   runLoadFragments( options )

if __name__ == '__main__':
   main()
