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

# pylint: disable=consider-using-f-string

# pylint: disable-next=deprecated-module
import tempfile, Tac, optparse, sys, os, struct, Pci, time

op = optparse.OptionParser( prog='flashFpga', usage='''
   %prog -- Wrapper Utility around flashrom to perform read
            write action on FPGA flash.

   %prog [-v] [-n] [--reset] [--strip] -f <FPGA> <-r|w> <FILENAME>
   * not all sections are supported on all platdforms.
   ''' )

op.add_option( "-r", "--read", action='store_true', default=None,
      help="Read from Flash" )

op.add_option( "-w", "--write", action='store_true', default=None,
      help="Write to Flash" )

op.add_option( "-f", "--fpga", action='store', default=None,
      help="Specify FPGA to Flash" )

op.add_option( "", "--reset", action='store_true', default=None,
      help="Reset the FPGA with the new firmware after upgrade" )

op.add_option( "", "--strip", action='store_true', default=None,
      help="File contains Arista header and needs to be stripped" )

op.add_option( "-v", "--verbose", dest='verbose',
         action='store_true', default=None, help="verbose output" )

opts, args = op.parse_args()

# For debugging purposes, to test different flashroms
flashrom = "flashrom"

def run( cmd, verboseOut=False ):
   if verboseOut:
      return Tac.run( cmd )

   return Tac.run( cmd,
            stdout=Tac.CAPTURE,
            stderr=Tac.DISCARD )

def exitWithPrint( code, err=None, msg=None ):
   if msg:
      sys.stdout.write( msg + '\n' )
   if err:
      sys.stderr.write( err + '\n' )
   sys.exit( code )

class CaveCreekGpio( ):

   def rdReg( self, addr ):
      os.lseek( self.fd, self.offset + addr, os.SEEK_SET )
      value = os.read( self.fd, 4 )
      return struct.unpack( '=L', value )[0]

   def wrReg( self, addr, data ):
      os.lseek( self.fd, self.offset + addr, os.SEEK_SET )
      retVal = os.write( self.fd, struct.pack( '=L', data ) )
      return retVal

   def __init__( self ):
      acpi = Pci.Device( "0000:00:1f.0" )
      self.config = acpi.config( readOnly=True )
      # Read GPIO Base address from 0x48
      self.offset = self.config.read32( 0x48 )
      # Mask out bits 0-6 and 16-31
      self.offset &= 0xFF80
      # GPIO should already be enabled and set as output but set it anyways.
      # WP is GPIO 26.  USE_SEL is reg 0x0, IO_SEL is reg 0x4, GP_LVL is reg 0xc.
      self.fd = os.open( "/dev/port", os.O_RDWR )

      # Set the GPIO setting for SpiMuxGpio in case it wasn't set correctly
      use_sel = self.rdReg( 0x30 )
      use_sel |= ( 1 << ( 57 % 32 ) )
      self.wrReg( 0x30, use_sel )
      io_sel = self.rdReg( 0x34 )
      io_sel &= ~( 1 << ( 57 % 32 ) )
      self.wrReg( 0x34, io_sel )

   def __del__( self ):
      os.close( self.fd )

   def setSpiMuxGpio( self, value ):
      gp_lvl = self.rdReg( 0x38 )
      if value:
         gp_lvl |= ( 1 << ( 57 % 32 ) )
      else:
         gp_lvl &= ~( 1 << (57 % 32 ) )
      self.wrReg( 0x38, gp_lvl )

class FlashFpga:

   def __init__( self ):
      self.length = 0x0800000
      self.start = 0
      self.blockSize = 1024
      self.flashrom = "flashrom"

   def setMux( self, enableFpga ):
      '''Set mux to flash Fpga'''
      pass # pylint: disable=unnecessary-pass

   def stripHeader( self, outputFile, inputFilename ):
      pass

   def read( self, filename, options ):
      '''Read FPGA image from flash memory'''
      run( [ self.flashrom, '--start=%d' % self.start, '--numBytes=%d' % self.length,
             '-r', filename ], options.verbose )

   def write( self, filename,  options ):
      '''Write FPGA image into flash memory'''
      # pylint: disable-next=consider-using-with
      tmpOutputFile = tempfile.NamedTemporaryFile()

      if options.strip:
         # pylint: disable-next=consider-using-with
         tmpInputFile = tempfile.NamedTemporaryFile()
         self.stripHeader( tmpInputFile, filename )
         tmpFilename = tmpInputFile.name
      else:
         tmpFilename = filename

      # need to create padded image
      run( [ 'dd', 'if=/dev/zero', 'of=' + tmpOutputFile.name,
             'bs=%d' % self.blockSize, 'count=%d' % \
                ( self.length // self.blockSize ) ],
           options.verbose )
      run( [ 'dd', 'if=' + tmpFilename, 'of=' + tmpOutputFile.name,
             'bs=%d' % self.blockSize, 'seek=%d' % ( self.start // self.blockSize ),
             'conv=notrunc' ], options.verbose )
      run( [ self.flashrom, '-w', tmpOutputFile.name ], options.verbose )

   def reset( self ):
      '''Reset FPGA to load new image'''
      pass # pylint: disable=unnecessary-pass

class Jalapeno( FlashFpga ):
   '''
   Schooner FPGA (contains both jalapeno and onion)
   Also used by TimberCove FPGA (serrano)
   '''
   def __init__( self ):
      FlashFpga.__init__( self )
      self.gpio = CaveCreekGpio()

   def setMux( self, enableFpga ):
      if ( enableFpga ): # pylint: disable=superfluous-parens
         self.gpio.setSpiMuxGpio( 1 )
      else:
         self.gpio.setSpiMuxGpio( 0 )

   def stripHeader( self, outputFile, inputFilename ):
      """Read through the Arista header, which is terminated by null character"""
      headerSize = 0
      fileSize = os.path.getsize( inputFilename )
      arbfFile = open( inputFilename, 'rb' ) # pylint: disable=consider-using-with
      # First read through the Arista header, which is terminated by null character
      while( arbfFile.read( 1 ) != '\x00' ): # pylint: disable=superfluous-parens
         headerSize += 1
         continue
      headerSize += 1
      size = fileSize - headerSize
      outputFile.write( arbfFile.read( size ) )
      outputFile.flush()

   def reset( self ):
      '''Set the SCD_PRGM bit on CPLD'''
      run( [ "i2cset", "-y", "1", "0x23", "0x10", "0x2" ] )
      time.sleep( 1.0 )
      run( [ "i2cset", "-y", "1", "0x23", "0x10", "0x0" ] )

read = opts.read
write = opts.write

# Check right num args
if len( args ) != 1:
   op.print_usage()
   exitWithPrint( 1, err="Invalid argument count." )

# Get the FPGA name
if opts.fpga in ( "jalapeno", "serrano" ):
   fpga = Jalapeno()
else:
   op.print_usage()
   exitWithPrint( 1, err="FPGA not specified." )

# Check that either read or write is specified
if read and not write:
   target = read
elif write and not read:
   target = write
else:
   op.print_usage()
   exitWithPrint( 1, err="Invalid option." )

fileArg = args[ 0 ]

# Set the mux, read/write image, and unset the mux
try:
   # Set any muxes that are required.
   fpga.setMux( enableFpga=True )
   if read:
      # Operation is a read.
      fpga.read( fileArg, opts )
   else:
      # Operation is a write.
      fpga.write( fileArg, opts )
finally:
   # Set any muxes that are required.
   fpga.setMux( enableFpga=False )

if opts.reset:
   fpga.reset()


