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

import argparse
import AgentDirectory
import FpgaUtil

description = '''Utility to perform FPGA functions using FPGA plugins'''

class Fpgas:
   def __init__( self, checkSmbusRunning=True ):
      # Some plugins may be passing in --raw to smbus commands. Ideally,
      # the utility should be usable even with agents running, but this is
      # not possible until all plugins are cleaned up.
      if checkSmbusRunning:
         assert not AgentDirectory.agent( 'ar', 'Smbus' ), (
            'This utility is currently not supported when the Smbus agent '
            'is running. It is recommended to shut down all agents' )

      with open( '/etc/prefdl', 'rb' ) as f:
         prefdl = f.read()
      flavor = FpgaUtil.translateSwiFlavor()
      self.systemFpgaListObj = FpgaUtil.systemFpgas( prefdl, flavor )

      self.fpgas = { fpga.name() : fpga
                     for fpga in self.systemFpgaListObj.fpgaList() }

   def listFpgas( self ):
      # pylint: disable-next=consider-using-f-string
      print( '{:16s} {:s}'.format( 'Name', 'Version' ) )
      for name in sorted( self.fpgas.keys() ):
         self.printVersion( name )

   def printVersion( self, name ):
      self._verifyName( name )
      fpga = self.fpgas[ name ]
      vers = fpga.hardwareVersion()
      print( f'{fpga.name():16s} {vers}' )
      
   def callFpgaProgram( self, name, imageFilePath,
                         spiImageFilePath=None ):
      self._verifyName( name )
      fpga = self.fpgas[ name ]
      if imageFilePath:
         fpga.imageFilePath = lambda: imageFilePath
         if spiImageFilePath:
            fpga.spiImageFilePath = lambda: spiImageFilePath
      fpga.program()

   def callFpgaFunction( self, funcName, name, imageFilePath,
                         spiImageFilePath=None ):
      self._verifyName( name )
      fpga = self.fpgas[ name ]
      assert hasattr( fpga, funcName ), (
         f'Function {funcName} not available for FPGA {name}' )
      func = getattr( fpga, funcName )
      if imageFilePath:
         fpga.imageFilePath = lambda: imageFilePath
         fpga.spiImageFilePath = lambda: imageFilePath
      func()

   def _verifyName( self, name ):
      assert name in self.fpgas, (
         # pylint: disable-next=consider-using-f-string
         'Unknown FPGA {}. Valid names are {}'.format(
            name, ', '.join( self.fpgas.keys() ) ) )

def handleList( args ):
   fpgasObj = Fpgas( checkSmbusRunning=not args.force )
   fpgasObj.listFpgas()

def handleVersion( args ):
   fpgasObj = Fpgas( checkSmbusRunning=not args.force )
   fpgasObj.printVersion( args.name )

def handleCallFunc( args ):
   fpgasObj = Fpgas( checkSmbusRunning=not args.force )
   fpgasObj.callFpgaFunction( args.funcName, args.name, args.imagePath )

def handleProgramFunc( args ):
   fpgasObj = Fpgas( checkSmbusRunning=not args.force )
   fpgasObj.callFpgaProgram( args.name, args.imagePath, args.spiImagePath )

def main():
   parser = argparse.ArgumentParser( description=description )
   parser.add_argument( '-f', '--force',
                        help='Force operation with agents running',
                        action='store_true' )

   subparsers = parser.add_subparsers( help='FPGA utility' )

   printParser = subparsers.add_parser( 'list', help='List FPGAs' )
   printParser.set_defaults( func=handleList )

   versionParser = subparsers.add_parser( 'version', help='print version' )
   versionParser.add_argument( 'name',
                               help='FPGA name(as listed by print command)' )
   versionParser.set_defaults( func=handleVersion )

   programParser = subparsers.add_parser( 'program', help='call program function' )
   programParser.add_argument( 'name',
                               help='FPGA name(as listed by print command)' )
   programParser.add_argument( 'imagePath', metavar='image-path',
                                help='FPGA image path', nargs='?' )
   programParser.add_argument( '--spiImagePath',
                      help='Use the given file name for the spiImagePath' )
   programParser.set_defaults( func=handleProgramFunc )

   for funcName in [ 'configure', 'erase', 'spiProgram' ]:
      callFuncParser = subparsers.add_parser(
         funcName, help=f'call {funcName} function' )
      callFuncParser.add_argument( 'name',
                                   help='FPGA name (as listed by list command)' )
      callFuncParser.add_argument( 'imagePath', metavar='image-path',
                                   help='FPGA image path', nargs='?' )
      callFuncParser.set_defaults( funcName=funcName )
      callFuncParser.set_defaults( func=handleCallFunc )

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

if __name__ == '__main__':
   main()
