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

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

import Tac, Tracing, SmbusClient
import random, time, ctypes, sys

from PyClient import PyClient
from PowerDiagLib import OptionParser, hwPowerSupplySlotConfig, envPowerStatusDir
from PowerDiagLib import exitWithError, configsAndStatuses
from datetime import datetime, timedelta

t0 = Tracing.trace0
failed = 0

def agentCfgs( sysdbRoot ):
   agentConfigDir = sysdbRoot.entity[ "sys/config/agentConfigDir" ]
   return agentConfigDir

def getPmbus11Config( sysdbRoot ):
   pmbus11Config = sysdbRoot.entity[ "hardware/powerSupply/pmbus11" ]
   return pmbus11Config

def _deviceLoc( sysdbRoot, psNum ):
   pmbus11Config = getPmbus11Config( sysdbRoot )

   pmbusConfig = pmbus11Config.powerSupply.get( int( psNum ) )
   if not pmbusConfig:
      exitWithError( 'Could not find PowerSupply%s' % psNum )
   pmbusDeviceOffset = pmbusConfig.psmiDeviceOffset

   ahamAddress = SmbusClient.decodeAhamAddress( pmbusDeviceOffset )
   accelId = ahamAddress.accelId
   busId = ahamAddress.busId
   deviceId = ahamAddress.deviceId

   t0( "Device Location: AccelId=%s, BusId=%s, DeviceId=%s" %
       ( accelId, busId, deviceId ) )

   return ( accelId, busId, deviceId )

def dumpReg( title, regVal, numBytes, aham, pmbusOff, loud=True ):
   pmbusDeviceBuf = SmbusClient.newBuf( numBytes )
   req = aham.newRequest
   req.nextOp = Tac.Value( "Hardware::AhamRequest::OpDesc",
                           op="read", addr=pmbusOff+regVal,
                           data=pmbusDeviceBuf.address,
                           count=numBytes, elementSize=1, accessSize=numBytes )
   req.state = 'started'

   Tac.waitFor( lambda: req.status != 'inProgress',
                description='dump of register 0x%02x to complete' % regVal )

   if req.status != 'succeeded':
      exitWithError( "Failed to read the PMBus device of the power supply. "
                     "Is it properly inserted? Does the device actually "
                     "support PMBus?" )

   toPrint = "0x"
   for byte in range(numBytes-1, -1, -1):
      assert pmbusDeviceBuf.str[byte], "missing byte to print. software bug"
      toPrint += "%02x" % ord( pmbusDeviceBuf.str[byte] )
   if loud:
      print( f"{title} (0x{regVal:x}):", toPrint )
   req.state = 'cancelled'
   return toPrint

# --------------------------------------------------------------------------------
# Class for various kinds of requests, that may be read or write
# --------------------------------------------------------------------------------

class OpAndBuf:
   def __init__( self, op, buf, pmbusOffset ):
      self.op_ = op
      self.buf_ = buf
      self.pmbusOffset = pmbusOffset
      assert self.op_.op in [ 'write', 'read' ]

   def __str__( self ):
      if self.op_.op == 'write':
         return "[W 0x%x]" % ( self.regVal() )
      elif self.op_.op == 'read':
         return f"[R 0x{self.regVal():x}: 0x{self.resultFromRead():x}]"
      else:
         assert 0

   def regVal( self ):
      return self.op_.addr - self.pmbusOffset

   def resultFromRead( self ):
      if self.op_.op == 'write':
         return None
      elif self.op_.op == 'read':
         result = 0
         buf = self.buf_
         for ch in buf.str[ ::-1 ]:
            result = result << 8
            result = result | ord( ch )
         return result
      else:
         assert 0

class GenericRequest:
   def __init__( self ):
      pass

   def randomRequest( self, pmbusOffset ):
      raise NotImplementedError

class PmbusReadRequest( GenericRequest ):
   def __init__( self, regVal, numBytes ):
      GenericRequest.__init__( self )
      self.regVal = regVal
      self.numBytes = numBytes

   def randomRequest( self, pmbusOffset ):
      buf = SmbusClient.newBuf( self.numBytes )

      op = Tac.Value( "Hardware::AhamRequest::OpDesc",
                      op="read", addr=pmbusOffset+self.regVal,
                      data=buf.address, count=self.numBytes,
                      elementSize=1, accessSize=self.numBytes )

      return OpAndBuf( op, buf, pmbusOffset )

class PmbusWriteRequest( GenericRequest ):
   def __init__( self, regVal, numBytes, possibleValues ):
      GenericRequest.__init__( self )
      self.regVal = regVal
      assert 0 <= numBytes <= 2
      self.numBytes = numBytes
      self.possibleValues = possibleValues
      # assert that possibleValues can be represented by numBytes
      assert all( val >> self.numBytes * 8 == 0 for val in possibleValues )

   def _randomPossibleValue( self ):
      return random.choice( self.possibleValues )

   def randomRequest( self, pmbusOffset ):
      value = self._randomPossibleValue()

      # Allocate space on a 0-byte write even though it isn't used
      if self.numBytes <= 1:
         buf = ctypes.c_uint8( value )
      elif self.numBytes == 2:
         buf = ctypes.c_uint16( value )

      op = Tac.Value( "Hardware::AhamRequest::OpDesc",
                      op="write", addr=pmbusOffset+self.regVal,
                      data=ctypes.addressof( buf ), count=self.numBytes,
                      elementSize=1, accessSize=self.numBytes )

      return OpAndBuf( op, buf, pmbusOffset )

# --------------------------------------------------------------------------------
# Class that generates requests
# --------------------------------------------------------------------------------

class RequestGenerator():
   def __init__( self, pmbusOffset, supplyState, debug=False ):
      self.opsAndBufs = []
      self.pmbusOffset = pmbusOffset
      self.debug = debug

      assert supplyState in [ 'ok', 'powerLoss' ]
      self.supplyState = supplyState

   def generateRequests( self, numRequests ):
      raise NotImplementedError

   def validate( self ):
      if self.debug:
         # Print register actions
         print( ' '.join( [ x.__str__() for x in self.opsAndBufs ] ) )

      # Gather results
      results = [ ( x.regVal(), x.resultFromRead() ) for x in self.opsAndBufs ]

      # Check that all status word reads agree with supply state from
      # the start (if off, stay off; if on, stay on)
      statusWordReads = [ resultFromRead for regVal, resultFromRead in results if
                          ( regVal == 0x79 and resultFromRead is not None ) ]
      statusFans12Reads = [ resultFromRead for regVal, resultFromRead in results if
                            ( regVal == 0x81 and resultFromRead is not None ) ]

      if self.supplyState == 'ok':
         expectVal = 0x0
      elif self.supplyState == 'powerLoss':
         expectVal = 0x2841
      else:
         assert 0

      # Allow bit 2 to be set in case of OVERTEMP such as in EDVT high
      # temp corner
      # Allow fan fault since fan speed override may count as a fault
      # Allow CML faults since accessing unsupported registers may count as a fault
      statusResults = [ not ( x & ~( expectVal | 0x0406 ) )
                        for x in statusWordReads ]
      fanResults = [ not ( x & ~( 0x08 ) ) for x in statusFans12Reads ]
      if not all( statusResults + fanResults ):
         global failed
         failed = failed + 1
      if not all( statusResults ):
         print()
         failureString = "Not all the statusWordReads were the same. " \
             "statusResults=%s" % statusResults
         print( failureString )
         print( "failureString:", [ "0x%x" % y for y in statusWordReads ] )
      if not all( fanResults ):
         print()
         failureString = "Not all the statusWordReads were the same. " \
             "fanResults=%s" % fanResults
         print( failureString )
         print( "failureString:", [ "0x%x" % y for y in statusFans12Reads ] )

   def clear( self ):
      self.opsAndBufs = []

class RandomPmbusRequestGenerator( RequestGenerator ):
   "Pass different kinds of requests and choose any of those at any given time"
   def __init__( self, pmbusOffset, supplyState, debug=False ):
      RequestGenerator.__init__( self, pmbusOffset, supplyState, debug )
      requestTypes = []
      byteReads = [ 0x78, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x80, 0x81, 0x82 ]
      for read in byteReads:
         requestTypes.append( PmbusReadRequest( read, 1 ) )

      twoByteReads = [ 0x79, 0x8d ]
      for read in twoByteReads:
         requestTypes.append( PmbusReadRequest( read, 2 ) )

      self.requestTypes = requestTypes

   def generateRequests( self, numRequests ):
      for _ in range( numRequests ):
         requestType = random.choice( self.requestTypes )
         opAndBuf = requestType.randomRequest( self.pmbusOffset )
         self.opsAndBufs.append( opAndBuf )
         yield opAndBuf

class OrderedPmbusRequestGenerator( RequestGenerator ):
   """ Rotate through a given order """
   def __init__( self, pmbusOffset, supplyState, debug=False ):
      RequestGenerator.__init__( self, pmbusOffset, supplyState, debug )
      self.orderedRequests = []

   def generateRequests( self, numRequests ):
      for i in range( numRequests ):
         length = len( self.orderedRequests )
         opAndBuf = self.orderedRequests[ i %
                                          length ].randomRequest( self.pmbusOffset )
         self.opsAndBufs.append( opAndBuf )
         yield opAndBuf

class PmbusAgentOrderRequestGenerator( OrderedPmbusRequestGenerator ):
   def __init__( self, pmbusOffset, supplyState, debug=False ):
      OrderedPmbusRequestGenerator.__init__( self, pmbusOffset, supplyState, debug )
      # Misc regular reads
      for register, numBytes in [ ( 0x8d, 2), ( 0x20, 1 ), ( 0x88, 2 ), ( 0x89, 2 ),
                                  ( 0x8b, 2 ), ( 0x8c, 2 ), ( 0x96, 2 ), ( 0x97, 2 ),
                                  ( 0x79, 2 ), ( 0x7a, 1 ), ( 0x7b, 1 ), ( 0x7c, 1 ),
                                  ( 0x7d, 1 ), ( 0x7e, 1 ), ( 0x7f, 1 ), ( 0x80, 1 ),
                                  ( 0x81, 1 ), ( 0x82, 1 ) ]:
         self.orderedRequests.append( PmbusReadRequest( register, numBytes ) )

      # Add a write to CLEAR_FAULTS
      self.orderedRequests.append( PmbusWriteRequest( 0x3, 1, [ 1 ] ) )

      # Read READ_FAN_SPEED_1
      self.orderedRequests.append( PmbusReadRequest( 0x90, 2 ) )

      # Add a write to WRITE_PROTECT (only do 0)
      self.orderedRequests.append( PmbusWriteRequest( 0x10, 1, [ 0 ] ) )

      # Add a write to FAN_COMMAND_1
      self.orderedRequests.append( PmbusWriteRequest( 0x3b,
            2, list( range( 101 ) ) ) )

class WriteFanAndReadStatusOrderRequestGenerator( OrderedPmbusRequestGenerator ):
   def __init__( self, pmbusOffset, supplyState, debug=False ):
      OrderedPmbusRequestGenerator.__init__( self, pmbusOffset, supplyState, debug )
      # Read STATUS_WORD
      self.orderedRequests.append( PmbusReadRequest( 0x79, 2 ) )

      # Set CLEAR_FAULTS
      self.orderedRequests.append( PmbusWriteRequest( 0x3, 1, [ 1 ] ) )

      # Write something to fan
      self.orderedRequests.append( PmbusWriteRequest( 0x3b,
            2, list( range( 101 ) ) ) )

def pummelRegisters( generator, aham, queueLength=100 ):

   req = aham.newRequest
   for opAndBuf in generator.generateRequests( queueLength ):
      req.nextOp = opAndBuf.op_

   req.state = 'started'

   Tac.waitFor( lambda: req.status != 'inProgress',
                description='pummel to complete' )

   if req.status != 'succeeded':
      global failed
      failed = failed + 1
      print( "Warning, there was a failed request. I will keep going but I will "
             "fail this test." )

   req.state = 'cancelled'

   generator.validate()
   generator.clear()
   if failed:
      print( "Failed", failed, "times so far" )

def main():
   parser = OptionParser()
   parser.add_option( "", "--powerSupply", action="store",
                      help="The power supply to test " )
   parser.add_option('--printRegisters', action='store_true',
                     help='Print interesting registers' )
   parser.add_option('--debug', action='store_true',
                     help='Debug mode, print out the order of events',
                     default=False )
   parser.add_option('--pummelTest', action='store',
                     help='Pummel power supply with many requests and check ' \
                        'statusWord stays consistent (default: %default)',
                     default='pmbusAgentOrder' )
   parser.add_option('--duration', action='store', type=int,
                     help='minimum test duration (default=%default)',
                     default=5 )
   parser.add_option('--minQueueLength', action='store', type=int,
                     help='The minimum length of a request during this test ' \
                        '(default=%default)',
                     default=50 )
   parser.add_option('--maxQueueLength', action='store', type=int,
                     help='The maximum length of a request during this test ' \
                        '(default=%default)',
                     default=200 )
   parser.add_option('--keepAgentAlive', action='store_true',
                     help='run the diag at the same time as the agent (for example '\
                        ' before 4.9.2 where agent shutdown is unsupported)',
                     default=False )
   ( options, args ) = parser.parse_args()


   seed = int( time.time() )
   print( "Random seed is", seed )
   random.seed( seed )

   if args:
      parser.error( "unexpected arguments" )
   if not options.powerSupply:
      parser.error( "you must specify a power supply to test" )

   sysname = options.sysname
   sysdbRoot = PyClient( sysname, "Sysdb" ).agentRoot()

   powerSupplyConfig = hwPowerSupplySlotConfig( sysdbRoot )
   envPowStatusDir = envPowerStatusDir( sysdbRoot )

   ( envTemperatureStatus, thermostatConfig,
     thermostatStatus ) = configsAndStatuses( sysdbRoot )

   agentConfigDir = agentCfgs( sysdbRoot )

   def agentShutdown( agent, shutdown=True ):
      if agent in agentConfigDir.agent:
         agentCfg = agentConfigDir.agent[ agent ]
         agentCfg.shutdown = shutdown
      Tac.flushEntityLog()

   if options.keepAgentAlive:
      pass
   else:
      print( "Killing agent" )
      agentShutdown( 'Pmbus', shutdown=True )
      agentShutdown( 'PowerSupplyDetector', shutdown=True )

   SmbusClient.setup( sysname, "Pmbus11DeviceRegTest" )

   supplyConfig = powerSupplyConfig.slotConfig.get( options.powerSupply )
   powerStatus = envPowStatusDir.get( "PowerSupply%s" % options.powerSupply )

   assert powerStatus.state in [ 'ok', 'powerLoss' ]

   # Actually we poll more than 1 temp sensor on the power supply, but
   # this is good enough for now
   temperatureStatus = envTemperatureStatus.tempSensor.get(
      "TempSensorP%s/1" % options.powerSupply )

   if not supplyConfig:
      exitWithError( "Could not find supply %s" % options.powerSupply )

   print( "pmbus11DeviceRegTest: checking the PMBus device contents" )

   aham = supplyConfig.powerSupplyBaseAhamDesc.aham
   print( "Bus Timeout:", supplyConfig.powerSupplyBaseAhamDesc.busTimeout )
   print( "Read Delay:", supplyConfig.powerSupplyBaseAhamDesc.readDelay )
   print( "Write Delay:", supplyConfig.powerSupplyBaseAhamDesc.writeDelay )

   (accel, bus, device) = _deviceLoc( sysdbRoot, options.powerSupply )
   pmbusOff = SmbusClient.ahamAddress( accelId=accel, busId=bus,
                                       addrSize='addrSize1', deviceId=device,
                                       offset=0 )

   if options.printRegisters:
      # Print out some of the interesting registers.  The PMBus Spec
      # defines the various STATUS bits.
      dumpReg( "STATUS_BYTE", 0x78, 1, aham, pmbusOff )
      dumpReg( "STATUS_WORD", 0x79, 2, aham, pmbusOff )
      dumpReg( "STATUS_VOUT", 0x7A, 1, aham, pmbusOff )
      dumpReg( "STATUS_IOUT", 0x7B, 1, aham, pmbusOff )
      dumpReg( "STATUS_INPUT", 0x7C, 1, aham, pmbusOff )
      dumpReg( "STATUS_TEMP", 0x7D, 1, aham, pmbusOff )
      dumpReg( "STATUS_CML", 0x7E, 1, aham, pmbusOff )
      dumpReg( "STATUS_OTHER", 0x7F, 1, aham, pmbusOff )
      dumpReg( "STATUS_MFR", 0x80, 1, aham, pmbusOff )
      dumpReg( "STATUS_FANS_1_2", 0x81, 1, aham, pmbusOff )
      dumpReg( "STATUS_FANS_3_4", 0x82, 1, aham, pmbusOff )
      # Expected to work for PWR-2900
      dumpReg( "PRI_CODE_REV", 0xee, 1, aham, pmbusOff )
      dumpReg( "SEC_CODE_REV", 0xef, 1, aham, pmbusOff )

      print( "Output Voltage (0x8b): ",
             powerStatus.outputVoltageSensor.voltage, "V" )
      print( "Output Current (0x8c): ",
             powerStatus.outputCurrentSensor.current, "A" )
      print( "Output Power (0x96): ", powerStatus.outputPower, "W" )
      print( "Input Voltage: (0x88)", powerStatus.inputVoltageSensor.voltage, "V" )
      print( "Input Current (0x89): ", powerStatus.inputCurrentSensor.current, "A" )
      print( "Input Power (0x97): ", powerStatus.inputPower, "W" )
      print( "Temperature 1 (0x8d): ", temperatureStatus.temperature, "C" )

   pummelTestPossibleOptions = [ 'randomRead', 'pmbusAgentOrder',
                                 'writeFanAndReadStatusOrder' ]
   if options.pummelTest and options.pummelTest not in pummelTestPossibleOptions:
      print( "Please enter a valid pummelTest option (%s)" % ", ".join(
         pummelTestPossibleOptions ) )
      sys.exit( 1 )

   if not options.pummelTest or options.pummelTest == 'pmbusAgentOrder':
      print( "Running register pummel test with similar request order "
             "as PMBus agent" )
      gen = PmbusAgentOrderRequestGenerator( pmbusOff, powerStatus.state,
                                             options.debug )
   elif options.pummelTest == 'writeFanAndReadStatusOrder':
      print( "Running register pummel test by doing lots of fan speed writes" )
      gen = WriteFanAndReadStatusOrderRequestGenerator( pmbusOff, powerStatus.state,
                                                        options.debug )
   elif options.pummelTest == 'randomRead':
      print( "Running register pummel test by doing plenty of reads" )
      gen = RandomPmbusRequestGenerator( pmbusOff, powerStatus.state, options.debug )

   savedThermostatMode = thermostatConfig.mode
   try:
      # Set the thermostat agent into debug mode so that it does not attempt to
      # configure the fan speeds
      thermostatConfig.mode = 'debug'
      # Wait for the Thermostat agent to acknowledge the mode change
      Tac.waitFor( lambda: thermostatStatus.mode == 'debug',
                   timeout=10,
                   description="Thermostat agent to go into debug mode" )

      testBeginTime = datetime.now()
      while datetime.now() - testBeginTime < \
             timedelta( 0, options.duration, 0 ):
         runStartTime = datetime.now() # print start time
         numToQueue = random.randint( options.minQueueLength,
                                      options.maxQueueLength )
         pummelRegisters( gen, aham, queueLength=numToQueue )
         runEndTime = datetime.now() # print end time
         print( "[queueLen=%d] Elapsed: %s" % ( numToQueue,
                                               str( runEndTime - runStartTime ) ) )
   finally:
      # Undo all configuration
      thermostatConfig.mode = savedThermostatMode
      # Wait for the Thermostat agent to acknowledge the mode change
      Tac.waitFor( lambda: thermostatStatus.mode == savedThermostatMode,
                   timeout=10,
                   description="Thermostat agent to go back to it's original mode" )
      # Make sure all changes get to sysdb before we exit. Although at this point
      # there shouldn't be any...
      Tac.flushEntityLog()

   if options.keepAgentAlive:
      pass
   else:
      print( "Re-enabling agent" )
      agentShutdown( 'Pmbus', shutdown=False )
      agentShutdown( 'PowerSupplyDetector', shutdown=False )



   if failed:
      print( "Failed", failed, "total during this test" )
      exitWithError( "I Failed to read the PMBus device of the power supply during "
                     "the test or the status_word flapped. Is it properly inserted? "
                     "Does the device actually support PMBus?" )


   print( "pmbus11DeviceRegTest: power supply", options.powerSupply, "passed!" )


if __name__ == "__main__":
   main()
