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

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

import email
import optparse # pylint: disable=deprecated-module
import os
import re
import sys

import Tac
import Tracing
import ScdRegisters
import Swag
import PlutoIdentifyLib

__defaultTraceHandle__ = Tracing.Handle( "IdentifyCell" )
t0 = Tracing.trace0
t1 = Tracing.trace1
t2 = Tracing.trace2

cmdLineFile = os.environ.get( "CMDLINE", '/proc/cmdline' )
prefdlPath = os.environ.get( 'PREFDL_PATH', '/mnt/flash/.prefdl' )

def writeIdFiles( outputDir, slotId, cellType, cellId, electionMgrSupported,
                  cpuBmcMode, cardSlotsSupported ):
   def path( filename ):
      return outputDir + "/" + filename
   open( path( 'slotid' ) , 'w' ).write( '%d' % slotId )
   open( path( 'celltype' ), 'w' ).write( cellType )
   open( path( 'cellid' ), 'w' ).write( '%d' % cellId )
   open( path( 'electionMgrSupported' ), 'w' ).write( 'TRUE' if electionMgrSupported
                                                      else 'FALSE' )
   open( path( 'cpuBmcMode' ), 'w' ).write( 'TRUE' if cpuBmcMode else 'FALSE' )
   open( path( 'cardSlotsSupported' ), 'w' ).write( 'TRUE' if cardSlotsSupported
                                                   else 'FALSE' )

def computeCellId( slotId, cellType ):
   """Compute cellId from slotId and cardType, following the rules of
   AID158: supervisors are 1-255, line cards are 257-511, fabric cards
   are 513-767."""
   if cellType in [ "supervisor", "fixed" ]:
      cellId = slotId
   else:
      cellId = slotId & 0xf
      if slotId & 0x10:
         cellId += 0x200
      else:
         cellId += 0x100
   return cellId

def identifyCell( scd ):
   """Return a tuple of slotId, cellType, fdl, fdlError.  Looks for
   SLOTID, /proc/cmdline, and then /sys/class/pci Ham in
   that order to decide.  A failure to find the slotId via any of the
   above mechanisms is reported in fdlError and the tuple of
   (0,'generic',None) is returned.  scd.slotId, scd.cellType, and
   scd.cellId are not modified.  The slotId is one-based."""

   errors = []

   # Check environment for CPUBMCMODE, mainly for btest
   cpuBmcMode = False
   cpuBmcModeVal = os.environ.get( "CPUBMCMODE" )
   if cpuBmcModeVal is not None:
      cpuBmcMode = bool( cpuBmcModeVal == "TRUE" )
   # Check environment for CARDSLOTSSUPPORTED, mainly for btest
   cardSlotsSupported = False
   cardSlotsSupportedVal = os.environ.get( "CARDSLOTSSUPPORTED" )
   if cardSlotsSupportedVal is not None:
      cardSlotsSupported = bool( cardSlotsSupportedVal == "TRUE" )

   # linux passes all unrecognized boot args to 'init' as environment variables
   # so SLOTID will always be found on
   # NETDEV environment variable to determine if we are a card or not.
   # Check the SLOTID environment var => cellType is 'supervisor'
   slotId = os.environ.get( "SLOTID" )
   if slotId is not None:
      print( 'found slotId', slotId, 'in env' )
      # Check environment for ELECTIONMGRSUPPORTED, mainly for btest
      electionMgrSupported = False
      electionMgrSupportedVal = os.environ.get( "ELECTIONMGRSUPPORTED" )
      if electionMgrSupportedVal is not None:
         electionMgrSupported = bool( electionMgrSupportedVal == "TRUE" )
      try:
         slotId = int( slotId )
      except ValueError:
         errors.append( "Illegal format for SLOTID: %s" % slotId )
      else:
         netdev = os.environ.get( "NETDEV" )
         if netdev:
            print( 'found netdev', netdev )

         if netdev and netdev.startswith( 'internal' ):
            cardType = 'card'
         else:
            cardType = 'supervisor'

         return ( slotId, cardType, None, electionMgrSupported, cpuBmcMode,
                  cardSlotsSupported )
   else:
      errors.append( "SLOTID not set in environment" )

   # Check for SLOTID=<N> in /proc/cmdline => cellType is 'card'
   t0( "opening", cmdLineFile )
   try:
      cmdline = open( cmdLineFile ).read()
   except OSError as e:
      errors.append( "failed to read kernel boot parameters: %s" % str( e ) )
   else:
      t0( "opened", cmdLineFile)
      m = re.search( r"SLOTID=(\d+)", cmdline )
      if m:
         slotId = int( m.group( 1 ) )
         t0( 'found slotId',slotId,'in cmdline')
         return slotId, 'card', None, False, cpuBmcMode, cardSlotsSupported
      else:
         errors.append( "SLOTID not present in kernel boot parameters" )

   # Try to read from the SCD PCI file => cellType is "supervisor" or "fixed"
   scdPath = ScdRegisters.scdPciResourceFile()
   if not scdPath:
      errors.append( "No PCI scd found" )
   elif not os.path.exists( scdPath ):
      errors.append( "Scd PCI file at " + scdPath + " does not exist" )
   else:
      scd.ham = ( "scdHam", "hamTypeMemMapped", None, 0, 0, scdPath )
      scd.doReadSlotId()
      if scd.slotIdError:
         errors.append( scd.slotIdError )
      else:
         return ( scd.slotId, scd.cellType, None, scd.electionMgrSupported,
                  scd.cpuBmcMode, scd.cardSlotsSupported )

   # The registers provided by the SCD for an LCPU are not the same as on a fixed
   # system, so it will not be identified above. We're just going to say that if
   # /mnt/flash/.prefdl exists then it's an LCPU and we want to treat it as fixed.
   if os.path.exists( prefdlPath ):
      return ( 1, "fixed", None, False, False, False )

   # None of the above so we are a generic linecard
   return 0, "generic", " and ".join( errors ), False, cpuBmcMode, cardSlotsSupported

def detectPlutoYaml( options ):
   ( slotId, cellType, identifyError ) = PlutoIdentifyLib.identifyCell()
   if identifyError:
      return False

   cellId = computeCellId( slotId, cellType )
   t0( "Pluto says:", "cellId", cellId, "slotId", slotId, "cellType", cellType )

   # Pluto/whitebox systems are only for fixed so disable ElectionMgr
   writeIdFiles( options.dir, slotId, cellType, cellId, False, False, False )

   return True

def detectScd( options ):
   # Create the scdDriver
   scd = Tac.root.newEntity( "IdentifyCell::Scd", "IdentifyCell::Scd" )

   ( slotId, cellType, identifyError, electionMgrSupported,
     cpuBmcMode, cardSlotsSupported ) = identifyCell( scd )
   if Swag.isMember():
      print( "Overriding cellType to 'supervisor' as a SWAG member" )
      cellType = "supervisor"
   if identifyError:
      print( "Error when identifying the cell via the SCD: %s" % identifyError,
             file=sys.stderr )
      return False

   if slotId <= 0 or slotId > 1023:
      print( "Error slotId out of range: %d" % slotId )
      return False

   cellId = computeCellId( slotId, cellType )

   t0( "cellId", cellId, "slotId", slotId, "cellType", cellType )
   writeIdFiles( options.dir, slotId, cellType, cellId, electionMgrSupported,
                 cpuBmcMode, cardSlotsSupported )

   return True

def detectPrefdl( options ):
   try:
      with open( options.prefdl ) as prefdlFile:
         prefdlRaw = prefdlFile.read()
   except OSError:
      return False

   prefdl = email.message_from_string( prefdlRaw )
   slotId = prefdl.get( 'slotId', None )
   cellType = prefdl.get( 'cellType', None )

   # The prefdl need to at least contain those 2 fields in order to be considered
   # like a valid identification.
   if slotId is None or cellType is None:
      return False

   slotId = int( slotId )
   cellId = computeCellId( slotId, cellType )
   electionMgrSupported = prefdl.get( 'electionMgrSupported', False )
   cpuBmcMode = prefdl.get( 'cpuBmcMode', False )
   cardSlotsSupported = prefdl.get( 'cardSlotsSupported', False )

   t0( "Prefdl says:", "cellId", cellId, "slotId", slotId, "cellType", cellType,
       "electionMgr", electionMgrSupported, "cpuBmcMode", cpuBmcMode,
       "cardSlotsSupported", cardSlotsSupported )

   writeIdFiles( options.dir, slotId, cellType, cellId, electionMgrSupported,
                 cpuBmcMode, cardSlotsSupported )

   return True

def genericCell( options ):
   slotId = 1
   cellType = "generic"
   cellId = 1
   electionMgrSupported = False
   cpuBmcMode = False
   cardSlotsSupported = False

   t0( "Generic cell:", "cellId", cellId, "slotId", slotId, "cellType", cellType )
   writeIdFiles( options.dir, slotId, cellType, cellId, electionMgrSupported,
                 cpuBmcMode, cardSlotsSupported )

def main():
   Tac.sysnameIs( 'ar' )

   parser=optparse.OptionParser()
   parser.add_option( "-t", "--test", action="store_true",
                      help="Run in diagnostic test mode" )
   parser.add_option( "-d", "--dir", action="store",
                      help="Where to write output files: (default: %default)",
                      default="/etc" )
   parser.add_option( "-p", "--prefdl", action="store",
                      help="Path to the prefdl (default: %default)",
                      default="/etc/prefdl" )

   ( options, args ) = parser.parse_args()
   if args:
      parser.error( "Unexpected arguments" )

   if detectPlutoYaml( options ):
      return 0

   if detectPrefdl( options ):
      return 0

   if detectScd( options ):
      return 0

   genericCell( options )
   return 1

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