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

from __future__ import print_function
import sys
import Pci
import ScdRegisters
from ScdRegisters import PciId
import argparse

pciAddressHelp = "[[DDDD:]BB:]SS[.F] D=Domain, B=BusNumber, S=Slot, F=Function"
addrHelp = "a register address within the SCD." \
   " Can be decimal, binary (0b...) or hex (0x...)"
dataHelp = "a 32bit value. Can be decimal, binary or hex"
scdIdHelp = "integer valued scd id."
moduleTypeHelp = "integer valued scd module type."
exampleList = \
"""  e.g., scd list
      list the details of scds in the system
  e.g., scd 0000:bb:00.0 list
      list the details of the scd with specific pci address
"""
exampleWrite = \
"""  e.g., scd write 0x5000 0b101101
      write scd address 5000 with 101 101 (turning on two leds).
"""
exampleRead = \
"""  e.g. scd read 0x5000
      print contents of scd address 5000 to console
"""
noScdError = "Error: No matching scd device found on this machine"
listFormat = '{:<15}{:<6}{:<15}'

def formatPciAddr( pciAddress ):
   if pciAddress is None:
      return None
   return str( Pci.Address( pciAddress ) )

def binary( value, numBytes=4 ):
   """String rep of an integer in binary, similar to 'hex'.
   Turns 0xc7 into 1100 0111.
   The optional bytes parameter indicates how many bytes wide the value is."""
   return "".join( [((value>>n &1) and "1" or "0") + ((n % 4 == 0) and " " or "")
                    for n in xrange(numBytes*8-1,-1,-1)] )

def errorAndExit( message ):
   sys.stderr.write( message + '\n' )
   sys.exit( 1 )

def parseArgs( args ):
   """Helper for parsing command line arguments"""
   # Preparser parses optional scdId and moduleType args
   # so the main parser only needs to worry about pciAddress
   parser = argparse.ArgumentParser( add_help=False )
   parser.add_argument( '-i', type=lambda x: None if x is None else int(x) )
   parser.add_argument( '-t', type=lambda x: None if x is None else int(x) )
   namespace, args = parser.parse_known_args( args )
   #add placeholder pciaddress if one wasn't passed
   if args[ 0 ] in [ "list", "write", "read" ]:
      args.insert( 0, None )

   # Start main parser
   parser = argparse.ArgumentParser( prog='scd',
      formatter_class=argparse.RawDescriptionHelpFormatter,
      epilog="Note: scdId or moduleType cannot be passed in with pciAddress." )

   parser.add_argument( 'pciAddress' , metavar='<pciaddress>', 
         nargs='?', type=formatPciAddr, help=pciAddressHelp)
   # Add these optional args again, even though they are already parsed by
   # the preparser, for helptext.
   parser.add_argument( '-i', metavar='<scdId>', help=scdIdHelp )
   parser.add_argument( '-t', metavar='<moduleType>', help=moduleTypeHelp )

   subparsers = parser.add_subparsers( dest="method", 
         help="scd subcommands.\n'scd <subcommand> -h' for details.")
   
   listHelp = "list information on scds in system"
   list_parser = subparsers.add_parser( 'list', epilog=exampleList, 
         formatter_class=argparse.RawDescriptionHelpFormatter,
         help=listHelp, description=listHelp)
   list_parser.set_defaults( run=scdList )
   
   readHelp = "read a scd register"
   read_parser = subparsers.add_parser( 'read', epilog=exampleRead, 
         formatter_class=argparse.RawDescriptionHelpFormatter,
         help=readHelp, description=readHelp)
   read_parser.add_argument( 'address', metavar='<addr>', 
         type=lambda x: int( x, 0 ), help=addrHelp ) 
   read_parser.set_defaults( run=scdIO )

   writeHelp = "write a value to a scd register"
   write_parser = subparsers.add_parser( 'write', epilog = exampleWrite, 
         formatter_class=argparse.RawDescriptionHelpFormatter,
         help=writeHelp, description=writeHelp)
   write_parser.add_argument( 'address', metavar='<addr>', 
         type=lambda x: int( x, 0 ), help=addrHelp )
   write_parser.add_argument( 'data', metavar='<data>',
         type=lambda x: int( x, 0 ), help=dataHelp )
   write_parser.set_defaults( run=scdIO )
   parsed = parser.parse_args(args, namespace=namespace )
   if parsed.pciAddress and ( parsed.i is not None or parsed.t is not None):
      parser.error( "scdId or moduleType cannot be passed in with pciAddress." )
   return parsed

def scdList( args ):
   """Helper for list command"""
   moduleType = { 
         0: 'Fixed-System', 
         1: 'Supervisor', 
         2: 'Fabric', 
         3: 'Linecard', 
         4: 'Switchcard' }
   devices = [ dev for dev in
         map( ScdRegisters.deviceInfo, Pci.allDevicesById( PciId ) )
         if dev.device is not None and
            ScdRegisters.scdAttributeMatch( dev,
                                            pciAddress=args.pciAddress,
                                            moduleType=args.t,
                                            scdId=args.i ) ]
   if not devices:
      errorAndExit( noScdError )

   result = [ listFormat.format( "pciAddress", "scdId", "moduleType" ) ]
   for dev in devices:
      result.append( listFormat.format(
         dev.pciAddress, dev.scdId,
         moduleType.get( dev.modType, "Unknown" ) ) )
   print( '\n'.join( result ) )

def scdIO( args ):
   """Helper for read/write command"""
   scd = ScdRegisters.scdPciResourceFile( pciAddress=args.pciAddress,
                                          scdId=args.i, moduleType=args.t )
   if not scd:
      errorAndExit( noScdError )
   offset = ScdRegisters.scdPciOffset()
   if args.method == 'write':
      val = Pci.MmapResource( scd ).write32( offset + args.address, args.data )
   elif args.method == 'read':
      val = Pci.MmapResource( scd, readOnly=True ).read32( offset + args.address )
      print( hex( args.address ), hex( val ), "==", binary( val, 4 ) )
   else:
      errorAndExit( "Error: unrecognized argument: " + args.method )

def main( args ):
   """Parse arguments and call appropriate function"""
   #remove first argument 'scd' before parsing
   parsed = parseArgs( args [ 1: ] )
   parsed.run( parsed )

if __name__ == "__main__":
   main( sys.argv )
