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

# APIs defined in this file can be run by image with different version
# from image this file belongs to. The APIs must be consistent with
# previous and future images.

# There could be a case where ASU/SSU is triggered to upgrade the switch from an
# old release that only supports python2, and this file will be extracted from the
# destination image and get triggered in a python2 environment. Thus, we need to
# maintain the python2 compatibility in this file to avoid any ASU/SSU upgrade
# issues.

from __future__ import absolute_import, division, print_function

import os
import subprocess
import sys
import sysconfig

import Tac
import EosVersion
import SimpleConfigFile

def pstoreUpgradePolicy( prevDict, currDict, errOut=None ):
   '''PStore manifest is a dictionary with feature name as key and
   list of supported keys per feature as value. Upgrade policy checks
   for each feature in prevDict if key set for the feature is subset
   of supported key set in currDict. It returns True if policy check
   passes. The error string is optionally written to errOut.
   '''
   if not isinstance( prevDict, dict ) or not isinstance( currDict, dict ):
      if errOut:
         errOut.write( "Invalid input type\n" )
         errOut.flush()
      return False

   for feature, featureKeys in prevDict.items():
      if feature not in currDict:
         if errOut:
            errOut.write( 
               "PStore manifest check fail: feature %s not in manifest\n" %
               repr( feature ) )
            errOut.flush()
         return False

      currKeys = currDict[ feature ]
      diffset = set( featureKeys ).difference( set( currKeys ) )
      if diffset:
         if errOut:
            errOut.write( 
               "PStore manifest check fail: keys in feature %s not supported\n" %
               repr( feature ) )
            diffSetStr = "    " + ", ".join( str( k ) for k in diffset )
            errOut.write( diffSetStr )
            errOut.write( "\n" )
            errOut.flush()
         return False
   return True

class RunCmd( object ):
   def run( self, cmd ):
      return Tac.run( cmd.split( ' ' ), stdout=Tac.CAPTURE, stderr=Tac.CAPTURE,
                      asRoot=True )

class BootConfigParser( object ):
   def __init__( self ):
      # Hack since BootConfigParser depends on Url creating a circular dependency
      # Cli -> SysdbAgent -> Asu -> Cli
      # This pacher is only intended for Delhi.C2 so using this dynamic import
      # workaround.
      try:
         self.urlImport_ = __import__( 'Url' )
      except: # pylint: disable-msg=W0702
         self.urlImport_ = None

   def _bootConfigFilename( self ):
      url = self.urlImport_.parseUrl(
         "flash:/boot-config", self.urlImport_.Context( None, False ) )
      return url.realFilename_ # pylint: disable-msg=E1103

   def _bootConfig( self ):
      return SimpleConfigFile.SimpleConfigFileDict(
         self._bootConfigFilename(), createIfMissing=True, autoSync=True )

   def _bootSwiUrl( self ):
      return self._bootConfig()[ 'SWI' ]

   def bootSwiFilename( self ):
      if not self.urlImport_:
         return None

      url = self.urlImport_.parseUrl(
         self._bootSwiUrl(), self.urlImport_.Context( None, False ),
         lambda fs: fs.fsType == 'flash' )
      return url.realFilename_ # pylint: disable-msg=E1103

class AsuPolicyPatcher( object ):
   def __init__( self, srcVersion ):
      self.runCmd_ = RunCmd()
      self.swiFile_ = BootConfigParser().bootSwiFilename()
      self.srcVersion_ = srcVersion

      self.filesToExtract_ = [
         'AsuPolicyPatch/AsuNeighAdv.py',
      ]
      self.extractPath_ = '/tmp/'

   def _runCmd( self, cmd ):
      return self.runCmd_.run( cmd )

   def _logPatchInfo( self ):
      dstVersion = None
      if self.swiFile_:
         dstVersion = EosVersion.swiVersion( self.swiFile_ ).version()

      self.log( 'srcVerion: %s dstSwi: %s dstVersion: %s' % (
         self.srcVersion_, self.swiFile_, dstVersion ) )

   def _extractPatches( self ):
      cmd = 'unzip -o %s %s -d %s' % (
         self.swiFile_, ' '.join( self.filesToExtract_ ), self.extractPath_ )
      self.log( 'Extracting: %s' % cmd )
      self._runCmd( cmd )

   def _installPatches( self ):
      srcFiles = ' '.join(
         [ '%s%s' % ( self.extractPath_, x ) for x in self.filesToExtract_ ] )
      dstPath = sysconfig.get_path( 'purelib' )
      cmd = 'cp -f %s %s' % ( srcFiles, dstPath )
      self.log( "Installing: %s" % cmd )
      self._runCmd( cmd )

   def _execNeighAdvScript( self ):
      cmd = [ 'nohup', 'sudo',
              sys.executable, '-m', 'AsuNeighAdv.py' ]
      null = open(os.devnull, 'wb')
      _ = subprocess.Popen( cmd, shell=False, stdin=None, stdout=null,
                               stderr=null, close_fds=True)

   def log( self, message ):
      #self.logger_.log( message )
      print(message)

   def doPatch( self ):
      try:
         self._logPatchInfo()
      except IOError:
         # this happens with FastBootCliTest btest. ignore the error.
         return
      if self.swiFile_:
         self._extractPatches()
         self._installPatches()
         self._execNeighAdvScript()

def defaultUpgradePolicy( version, errOut=None ):
   '''A catch-all policy that can be used to fail ASU policy check.
   '''
   AsuPolicyPatcher( version ).doPatch()

   # No extra check now, return success (True).
   return True

