#!/usr/bin/env python3
# Copyright (c) 2017 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.
#
# This program provides a peek/poke interaction with Cfp2 modules.

import Cell, PyClient
import argparse
import sys, re


# ---------------------------------------------------------------------------------#
# Helper functions
# ---------------------------------------------------------------------------------#
# Returns true if system is fixed, false otherwise
def isFixed():
   return Cell.cellType() == "fixed"

# Returns true if the passed in string value is a number, return None otherwise.
# Accepts decimal, hexidecimal, and binary numbers of these forms:
#     8, 0x8, 0X8, 0b1000, 0B1000
def strToInt( s ):
   try:
      if s.lower().startswith( "0x" ):
         return int( s, 16 )
      elif s.lower().startswith( "0b" ):
         return int( s[2:], 2 )
      else:
         return int( s )
   except ValueError:
      return None

# Calls sys.exit( 1 ) if an invalid port number is seen, otherwise returns a
# port string of the form "EthernetX" or "EthernetX/X". The returned intf string
# should be a valid index into a xcvrStatus collection.
def parsePort( portStr ):
   def invalid():
      print( "Invalid port name: " + portStr )
      sys.exit( 1 )

   portStr = portStr.lower()
   validMatch = re.match( r'^(et|ethernet)?\d+(?:/\d+)?(?:/\d+)?$', portStr )
   if validMatch is None:
      print( "Invalid port name: " + portStr )
      sys.exit( 1 )

   # if we reached this point, the port string is, in general, a real port name
   ethernetStr = validMatch.group( 1 )
   portStr = portStr.strip( ethernetStr )
   r = re.match( r'^(?:(\d+)/)?(\d+)', portStr )
   r1 = r.group( 1 )
   r2 = r.group( 2 )

   if isFixed():
      if r1 is not None:
         return "Ethernet%s" % r1 # pylint: disable=consider-using-f-string
      else:
         return "Ethernet%s" % r2 # pylint: disable=consider-using-f-string
   else:
      if( r1 is None or r2 is None ):
         invalid()
      return "Ethernet%s/%s" % ( r1, r2 ) # pylint: disable=consider-using-f-string


usage = """
cfp2Mdio read <intfs> address [options]
cfp2Mdio write <intfs> address data
   <intfs> is accepted in the following forms (case-insenitive):
      4/5
      et4/5
      ethernet4/5
   If <intfs> includes the sub-interface number, like 4/5/1 for example,
   it is ignored and only 4/5 is used.

   <intfs> accepts a comma-seperated list of interface names.
   For example, et3/1,et4/6,et4/2,et5/5,et6/4.

   <intfs> can be replaced with "all". This is a shortcut to apply the
   operation to all recognized cfp2 modules in the system.

   address
      Address must be given in hexidecimal.

   data
      Data can be given in hexidecimal, decimal, or binary. For example,
      0x8, 0X8, 8, 0B1000, or 0b1000.

   -f --force is an optional argument to force the action even if module is not
      recognized as cfp2. Can be used to decode corrupted eeprom

Examples:
   1) To perform a write to address A010 on the module installed in port 41 on a
      fixed system:
         cfp2Mdio write 41 0xA010 0x4
   2) To read the identifier address from all cfp2 modules installed in a system:
         cfp2Mdio read all 0x8000
   3) To read the identifier address from a selection of ports spanning different
      linecards in a system:
         cfp2Mdio read et4/7,et3/28,et5/1,et5/2,et5/6 0x8000

WARNING: It is recommended to put modules into DIAGNOSTIC mode to avoid
possible MDIO Flow Control conflicts.
"""


# ---------------------------------------------------------------------------------#
# Function to parse user arguments and build a list of ports to peek/poke.
# Returns a tuple of ( user arguments, list of ports )
# ---------------------------------------------------------------------------------#
def parse():
   parser = argparse.ArgumentParser( prog="cfp2Mdio", usage=usage )

   # Create subparser's for read and write operations
   operationParser = parser.add_subparsers( dest="op",
                                            help="specifies operation type" )
   readParser = operationParser.add_parser( "read", usage=usage )
   writeParser = operationParser.add_parser( "write", usage=usage )

   # This is kind of wonky. Loop over the read and write subparsers to add their
   # shared arguments to each parser. We could add these arguments to the parent
   # parser, but then we can't get the ordering we want.
   for op_parser in [readParser, writeParser]:
      op_parser.add_argument( "ports", action="store", type=str,
                              help="specifies the port(s) to receive operation" )
      op_parser.add_argument( "address", action="store", type=str,
                              help="specifies module address" )

   # write specific arg(s)
   writeParser.add_argument( "data", action="store", type=str,
                             help="specifies value to write" )

   # read specific arg(s)
   readParser.add_argument( "-f", "--force", action='store_true',
                              default=False,
                              help="force the action" )

   #readParser.add_argument(
   #                   "count", action="store", type=int,
   #                   nargs='?',
   #                   default=1,
   #                   help="number of addresses to read starting at 'address'." )

   args = parser.parse_args()

   argPorts = None
   if args.ports.lower() == "all":
      argPorts = "all"
   elif args.ports:
      argPorts = []
      portStrs = args.ports.split( "," )
      for p in portStrs:
         argPorts.append( parsePort( p ) )
   else:
      parser.error( "too few arguments" )

   # attempt to convert 'register' and 'data' arguments to numbers
   if not args.address.lower().startswith( "0x" ):
      parser.error( "Please provide the module register as a hexidecimal number." )

   args.address = strToInt( args.address )
   if args.address is None:
      parser.error( "Invalid register value, register must be a number." )

   if args.op == "write":
      args.data = strToInt( args.data )
      if args.data is None:
         parser.error( "Invalid data value, data must be a number." )

   # done processing user input

   # Amass a list of xcvrStatus directories
   pycl = PyClient.PyClient( "ar", "Sysdb" )
   xcvrStatusDirs = []
   if isFixed():
      xcvrStatusAll = pycl.agentRoot().entity[ 'hardware/archer/xcvr/status/all' ]
      xcvrStatusDirs.append( xcvrStatusAll )
   else:
      sliceDir = pycl.agentRoot().entity[ 'hardware/archer/xcvr/status/slice/' ]
      if argPorts == "all":
         for lc in sliceDir.values():
            xcvrStatusDirs.append( lc )
      else:
         # Grab any Linecard dirs that coorespond to a user-specified port
         tmpLCs = [] # temporarily store linecards we've already added
         for port in argPorts:
            m = re.match( r'Ethernet(\d+)/\d+', port )
            lc = "Linecard" + m.group( 1 )
            if lc not in tmpLCs:
               xcvrStatusDirs.append( sliceDir[ "Linecard" + m.group( 1 ) ] )
               tmpLCs.append( lc )

   # Create set of valid ports by iterating over xcvrStatuses
   # A port is valid if:
   #     1. xcvr type is cfp2
   #     2. xcvr is present
   validPorts = []
   for xcvrDir in xcvrStatusDirs:
      for port in xcvrDir.values():
         # if user specified some ports, filter out the non-specified ones
         if( port.name not in argPorts and
             not argPorts == "all" ):
            continue

         if( port.presence == "xcvrPresent" and
             port.xcvrType == 'cfp2' or args.force ):
            validPorts.append( port.name )
         elif not argPorts == "all":
            # print message if user specified an existing port thinking it was
            # installed as a cfp2 port, but actually was not
            print( port.name + " is not recognized as cfp2, skipping..." )

   return ( validPorts, args )


def sendMdioRequest( ports, args ):
   # Print output header
   fmtStr = "{0:<6} {1:<15} {2:<9} {3:<8}"
   print( fmtStr.format( "Op", "Port", "Address", "Result" ) )
   print( fmtStr.format( "-" * 6, "-" * 15, "-" * 9, "-" * 8 ) )

   # Add SCD offset
   scdOffset = 0x10000
   addr = scdOffset + args.address

   # function to read module given a ctrlSmDir and a port.
   def doRead( ctrlSmDir, port, reg, output=True ):
      evalStr = ".aham.hamImpl.data[ " + str( reg ) + " ]"
      try:
         result = ctrlSmDir[ port ].config.ahamDesc.eval( evalStr )
      except PyClient.RpcError:
         print( "Internal script error" )
         sys.exit( 1 )

      if output:
         regStr = hex( reg - scdOffset )
         print( fmtStr.format( "Read", port, regStr, hex( result ) ) )
      return result

   # function to write to module given a ctrlSmDir and port.
   def doWrite( ctrlSmDir, port, reg, data ):
      evalStr = ".aham.hamImpl.data.__setitem__( " \
                + str( reg ) + ", " + str( data ) + " )"
      try:
         ctrlSmDir[ port ].config.ahamDesc.eval( evalStr )
      except PyClient.RpcError:
         print( "Internal script error" )
         sys.exit( 1 )

      # perform a read to verify the write went through
      result = doRead( ctrlSmDir, port, reg, output=False )
      writeStatus = "Success" if result == data else "Failed"

      regStr = hex( reg - scdOffset )
      print( fmtStr.format( "Write", port, regStr, writeStatus ) )

   # perform read(s)/write(s) based on system type (fixed or modular)
   if isFixed():
      pycl = PyClient.PyClient( "ar", "XcvrAgent" )
      ctrlSmDir = pycl.agentRoot().entity[ 'xcvrCtrlAgent' \
                        ].xcvrControllerAgentSm.cfp2CtrlSm

      for port in ports:
         if args.op == "read":
            doRead( ctrlSmDir, port, addr )
         else:
            doWrite( ctrlSmDir, port, addr, args.data )

   else:  # Modular
      cfp2CtrlSmDirs = {}  # map to memoize dir lookup operation
      for port in ports:
         # grab the cfp2CtrlSm directory for this port's linecard,
         # and memoize it to save us from recomputing
         ctrlSmDir = None
         m = re.match( r'Ethernet(\d+)/\d+', port )
         lcNum = m.group( 1 )
         if lcNum in cfp2CtrlSmDirs:
            ctrlSmDir = cfp2CtrlSmDirs[ lcNum ]
         else:
            xcvrAgent = "XcvrAgent-" + "Linecard" + lcNum
            pycl = PyClient.PyClient( "ar", xcvrAgent )
            ctrlSmDir = pycl.agentRoot().entity[ 'xcvrCtrlAgent' \
                  ].xcvrControllerAgentSm.cfp2CtrlSm
            cfp2CtrlSmDirs[ lcNum ] = ctrlSmDir

         # perform mdio
         if args.op == "read":
            doRead( ctrlSmDir, port, addr )
         else:
            doWrite( ctrlSmDir, port, addr, args.data )


# ---------------------------------------------------------------------------------#
# Main
# ---------------------------------------------------------------------------------#
if __name__ == "__main__":
   ( userPorts, arguments ) = parse()
   sendMdioRequest( userPorts, arguments )
