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

# pylint: disable=raising-format-tuple

"""
Command-line utility for reading from and writing to SflowAccel FPGAs.

This utility can be used in multiple ways:

* Interactive mode: spawns a simple shell which can be used to enter commands to
  read from or write to FPGAs. Simply run the program without any arguments:

    $ sudo sflowaccel

  Once in the shell, use the "fpga" command to select an FPGA to operate on, then run
  the "read" and/or "write" commands.

* Non-interactive mode: accepts a single shell command as argument, executes it,
  prints the result, and exits. Use the -j/--json flag to format the data as JSON.
  Useful for programmatic invocation of the script.

    $ sudo sflowaccel --fpga=SflowAccelFpga0 read fpgaRevision
    0x3

    $ sudo sflowaccel --json fpga
    ["SflowAccelFpga0", "SflowAccelFpga1"]

* Batch file: uses the -b/--batchFile argument to read commands from the provided
  file rather than from the standard input stream.

    $ cat << EOF >> commands.txt
      read sflowVersion
      EOF

    $ sudo sflowaccel --fpga=SflowAccelFpga0 -b commands.txt
      0x1111
"""

# Standard imports
import argparse
from collections import namedtuple
import functools
import os
import sys

# 3rd party imports
import Tac

# Local imports
import SflowAccelFpgaLib
from SflowAccelShellLib import ShellBase, ShellError, arg, cmd, safeEval


# Custom argument types to use when parsing shell command arguments with argparse.

def anyInt( strVal ):
   """
   Return an integer constructed from a string in any base. This is used to accept
   values entered by the user in different radices, e.g. hexadecimal or binary.
   """
   return int( strVal, 0 )

def fpgaHwName( strVal ):
   """
   Verify that the given string is a valid FPGA register or table name.
   """
   if not SflowAccelFpgaLib.validRegOrTableName( strVal ):
      raise ShellError( 'Not a valid FPGA register or table name: {}', strVal )

   return strVal


def requiresFpgaMode( func ):
   """
   Decorator used to mark a "do_*" command as requiring the shell to be in FPGA mode
   (i.e. that the user selected an FPGA).
   """
   @functools.wraps( func )
   def wrapper( shell, line ):
      if not shell.currentFpga:
         raise ShellError( 'Need to be in FPGA mode' )
      return func( shell, line )
   return wrapper

HashTableInfo = namedtuple( 'HashTableInfo',
                            'keyType keyArgsBefore keyArgsExpected keyArgsAfter' )
hashTables = {
   'subintfIn' : HashTableInfo(
      Tac.Type( 'Hardware::SflowAccel::IngressSubIntfMapping' ),
      [ 'Ethernet1' ], [ 'vlanId', 'platIntfIndex', 'isLag' ], [ 0 ] ),
   'subintfOut' : HashTableInfo(
      Tac.Type( 'Hardware::SflowAccel::EgressSubIntfMapping' ),
      [ Tac.Type( 'Ale::HwEgressEncapIndex' )( 0, 0 ) ],
      [ 'platChipId', 'eei' ], [ 0 ] ),
}

def hashTblName( strVal ):
   """
   Verify that the given string is a valid hash table name.
   """
   if strVal not in hashTables:
      raise ShellError( 'Not a valid hash table name: {}', strVal )

   return strVal

class SflowAccelShell( ShellBase ):
   defaultIntro = 'Type "help" to get started, or "exit" to quit.'
   defaultPrompt = 'sflowaccel> '
   fpgaPromptFmt = '{fpgaName}> '

   prompt = defaultPrompt

   def __init__( self, fpgas, fpgaName=None, json=False, batchFile=None ):
      """
      Arguments:
      - fpgas: a "FPGA name -> SflowAccelFpgaLib.Fpga object" dict
         List of usable FPGA objects.
      - fpgaName: string
         If provided, immediately set the shell to operate on the corresponding FPGA.
      - json, batchFile: see the docstring of the ShellBase class.
      """
      if not batchFile:
         self.intro = self.defaultIntro

      ShellBase.__init__( self, json=json, batchFile=batchFile )

      self.fpgas = fpgas
      self.currentFpga = None

      # If we are in interactive mode and there is only one FPGA in the system when
      # the shell starts, go directly into FPGA mode to save time.
      if self.use_rawinput and len( fpgas ) == 1:
         fpgaName = next( iter( self.fpgas ) )

      if fpgaName:
         self.enterFpgaMode( fpgaName )

   def enterFpgaMode( self, fpgaName ):
      """
      Verify that the specified FPGA exists, then set the shell to operate on that
      FPGA and update the prompt to reflect that.
      """
      fpga = self.fpgas.get( fpgaName, None )

      if not fpga:
         raise ShellError( 'Unknown FPGA: {}', fpgaName )

      self.currentFpga = fpga
      self.promptIs( self.fpgaPromptFmt.format( fpgaName=fpgaName ) )

   def leaveFpgaModeOrExit( self ):
      """
      If the shell is in FPGA mode, go back to normal mode to allow the user to
      select a new FPGA. Otherwise, just exit the shell. This is called by "exit" and
      Ctrl-D.
      """
      if not self.currentFpga or len( self.fpgas ) == 1:
         return True

      self.currentFpga = None
      self.promptIs( self.defaultPrompt )
      return False

   def verifyRegisterName( self, name ):
      """
      Verify that the given string is a valid FPGA register name.
      """
      if not SflowAccelFpgaLib.validRegisterName( name ):
         raise ShellError( 'Not a register: {}', name )

   def verifyRegisterInstance( self, name, instanceId ):
      """
      Verify that the given instanceId is valid for this FPGA register.
      """
      if not self.currentFpga.validRegisterInstance( name, instanceId ):
         raise ShellError( 'Invalid instance: {} {}', name, instanceId )

   def verifyTableName( self, name ):
      """
      Verify that the given string is a valid FPGA table name.
      """
      if not SflowAccelFpgaLib.validTableName( name ):
         raise ShellError( 'Not a table: {}', name )

   def verifyTblInstance( self, name, instanceId ):
      """
      Verify that the given instanceId is valid for this FPGA table.
      """
      if not self.currentFpga.validTableInstance( name, instanceId ):
         raise ShellError( 'Invalid instance: {} {}', name, instanceId )

   #
   # Command: exit
   #

   def do_exit( self, _ ):
      """Exit the shell."""
      return self.leaveFpgaModeOrExit()

   #
   # Command: fpga
   #

   @cmd(
      arg( 'fpgaName', nargs='?', help='enter FPGA mode for this FPGA' )
   )
   def do_fpga( self, args ):
      """
      List FPGAs or enter FPGA mode.
      """

      if args.fpgaName:
         return self.enterFpgaMode( args.fpgaName )

      return list( self.fpgas )

   def render_fpga( self, fpgaNames ):
      """
      Print the given list of FPGA names, one per line.
      """
      for fpgaName in fpgaNames:
         print( fpgaName )

   def complete_fpga( self, text, *ignored ):
      """
      Complete the "fpga" command with the names of available FPGAs.
      """
      fpgaNames = set( self.fpgas )

      # If in FPGA mode, remove the corresponding FPGA name from the list of
      # completions.
      if self.currentFpga:
         fpgaNames.remove( self.currentFpga.name )

      return self.completions( text, fpgaNames )

   #
   # Command: read
   #

   @requiresFpgaMode
   @cmd(
      arg( 'regOrTblName', type=fpgaHwName,
           help='name of the register or table to read from' ),
      arg( 'instanceId', type=anyInt,
           help='instance of the table entry to read from' ),
      arg( 'tableIndex', type=anyInt, nargs='?',
           help='index of the table entry to read from' ),
   )
   def do_read( self, args ):
      """
      Read from a register or table entry.
      """
      if args.tableIndex is not None:
         self.verifyTableName( args.regOrTblName )
         self.verifyTblInstance( args.regOrTblName, args.instanceId )
         return self.currentFpga[ ( args.regOrTblName, args.instanceId, ) ][
            args.tableIndex ]
      else:
         self.verifyRegisterName( args.regOrTblName )
         self.verifyRegisterInstance( args.regOrTblName, args.instanceId )
         return self.currentFpga[ ( args.regOrTblName, args.instanceId, ) ].value

   def complete_read( self, text, *ignored ):
      """
      Complete the "read" command with FPGA register and table names.
      """
      return self.completions( text, SflowAccelFpgaLib.regAndTableNames )

   #
   # Command: write
   #

   def parseIndexAndValue( self, args ):
      """
      The "write" command accepts two argument formats:
         write <register> <value...>
         write <table> <index> <value...>

      This means that the second argument is either:
      - the value to write, when writing to a register; or
      - the index to write to, when writing to a table entry.

      In both cases, the value to write is obtained from all remaining arguments
      interpreted as an expression, which is useful to dynamically compute values
      (e.g. "1 << 4 | 1" to set the first and fourth bits).

      This function takes the argument namespace and returns a (index, value) tuple
      by interpreting the arguments depending on whether "regOrTblName" refers to a
      register or table.
      """
      index = args.index
      valueOps = args.value[ : ]

      # When writing to a register, the "index" argument captured by argparse is
      # actually the first operand in the expression to be evaluated as value.
      #
      # e.g. the arguments "<register> 1 << 2" are parsed as follows:
      #   regOrTableName = "<register>"
      #   index = 1
      #   value = ["<<", "2"]
      #
      # We work around this issue by manually prepending the index to the list of
      # value operands, e.g. ["1", "<<", "2"], and then clearing the index itself.
      if SflowAccelFpgaLib.validRegisterName( args.regOrTblName ):
         valueOps.insert( 0, str( index ) )
         index = None

      # Return the index, and the value evaluated as a Python expression.
      return index, safeEval( ' '.join( valueOps ) )

   @requiresFpgaMode
   @cmd(
      arg( 'regOrTblName', type=fpgaHwName,
           help='name of the register or table to write to' ),
      arg( 'instanceId', type=anyInt,
           help='instance of the table entry to write to' ),
      arg( 'index', type=anyInt, nargs='?',
           help='index of the table entry to write to' ),
      arg( 'value', nargs=argparse.REMAINDER,
           help='value to write, as a literal or simple arithmetic expression' )
   )
   def do_write( self, args ):
      """
      Write to a register or table entry.
      """
      index, value = self.parseIndexAndValue( args )

      if index is not None:
         self.verifyTableName( args.regOrTblName )
         self.verifyTblInstance( args.regOrTblName, args.instanceId )
         self.currentFpga[ ( args.regOrTblName, args.instanceId, ) ][ index ] = value
      else:
         self.verifyRegisterName( args.regOrTblName )
         self.verifyRegisterInstance( args.regOrTblName, args.instanceId )
         self.currentFpga[ ( args.regOrTblName, args.instanceId, ) ].value = value

   def complete_write( self, text, *ignored ):
      """
      Complete the "write" command with FPGA register and table names.
      """
      return self.completions( text, SflowAccelFpgaLib.regAndTableNames )

   @cmd(
      arg( 'tblName', type=hashTblName,
           help='name of the hash table' ),
      arg( 'keyArgs', type=anyInt, nargs='+',
           help='arguments for the key to hash' )
   )
   def do_hash( self, args ):
      """
      Get hashes for a given hash table.
      Eg:
         sflowaccel> hash subintfOut 1 16400
         egressSubIntf( (0, 0), 1, 16400 )
         tableId: 0, hash: 0x6f7
         tableId: 1, hash: 0x9a8
         tableId: 2, hash: 0xaa7
      """
      tblInfo = hashTables.get( args.tblName )
      assert tblInfo is not None
      if len( args.keyArgs ) != len( tblInfo.keyArgsExpected ):
         raise ShellError( 'Expected args: {}', ' '.join( tblInfo.keyArgsExpected ) )

      keyArgs = tblInfo.keyArgsBefore + args.keyArgs + tblInfo.keyArgsAfter
      key = tblInfo.keyType( *keyArgs )
      print( key.toStrep() )
      for hashId in range( key.numHashes ):
         print( f'tableId: {hashId}, hash: {key.getHash( hashId ):#x}' )

   def complete_hash( self, text, *ignored ):
      """
      Complete the "hash" command with FPGA table names.
      """
      return self.completions( text, list( hashTables ) )


def parseArgs():
   formatter = argparse.RawDescriptionHelpFormatter
   parser = argparse.ArgumentParser( description=__doc__, formatter_class=formatter )

   parser.add_argument( '-j', '--json', action='store_true',
                        help='display the output in JSON format' )
   parser.add_argument( '-f', '--fpga', metavar='FPGANAME',
                        help='name of the FPGA to operate on' )
   parser.add_argument( '-b', '--batch', metavar='FILE',
                        type=argparse.FileType( 'r' ),
                        help='batch file to read commands from' )
   parser.add_argument( 'cmd', nargs=argparse.REMAINDER,
                        help='command to execute in non-interactive mode' )
   parser.add_argument( '--standalone', action='store_true' )

   return parser.parse_args()


def main():
   args = parseArgs()

   if not args.standalone and os.geteuid() != 0:
      print( 'Please run as root.' )
      sys.exit( 1 )

   # We are in interactive mode if reading from a TTY and no command has been
   # provided as argument.
   interactive = sys.stdin.isatty() and not args.cmd

   # Use the provided batch file, or force stdin if not in interactive mode. This is
   # to allow things like pipes and redirections.
   batchFile = args.batch or ( sys.stdin if not interactive else None )

   if args.standalone:
      fpgas = {}
   else:
      fpgas = SflowAccelFpgaLib.getFpgas()

      if not fpgas:
         print( 'Could not find SflowAccel FPGAs.' )
         sys.exit( 1 )

   shell = SflowAccelShell( fpgas,
                            fpgaName=args.fpga,
                            json=args.json,
                            batchFile=batchFile )

   if args.cmd:
      shell.onecmd( ' '.join( args.cmd ) )
   else:
      shell.cmdloop()


if __name__ == '__main__':
   main()
