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

import functools
import sys, re

def cPlusPlusFilt( orig ):
   """A very restricted version of c++filt."""
   s = orig
   components = []
   try:
      assert s[ 0 ] == 'N'
      s = s[ 1: ]
      while s:
         if s[ 0 ] == 'E': # pylint: disable=no-else-break
            break
         else:
            m = re.match( r'(\d+)(.*)', s )
            assert m
            n = int( m.group( 1 ) )
            s = m.group( 2 )
            assert len( s ) >= n
            components.append( s[ :n ] )
            s = s[ n: ]
      assert s[ 0 ] == 'E'
      s = s[ 1: ]
      assert not s
      return '::'.join( components )
   except ( AssertionError, ValueError, IndexError ):
      return orig

def makeTypeRe( types ):
   # pylint: disable-next=consider-using-f-string
   return re.compile( r"(%s)\b" % ( '|'.join( types ) ) )

introspectionTypes = [ "Tac::TypeMap" ]
introspectionRe = makeTypeRe( introspectionTypes )

dirTypes = [ "Tac::Dir" ]
dirRe = makeTypeRe( dirTypes )

mountTypes = [
   "Tac::AgentEnvConfig",
   "Tac::AgentEnvConfigDir",
   "Tac::AgentEnvDirSm",
   "Tac::AgentEnvDirSmWrapper",
   "Tac::AgentEnvSm",
   "Tac::AgentEnvStatus",
   "Tac::AgentEnvStatusDir",
   "Tac::NboAttrLog::In",
   "Tac::NboAttrLog::Out",
   "Tac::ConnectionDirWrapper",
   "Tac::DemuxerManagerDir",
   "Tac::MountConfig",
   "Tac::MountConfigDir",
   "Tac::MountDirAgentConfig",
   "Tac::MountDirSm",
   "Tac::MountSm",
   "Tac::MountStatus",
   "Tac::MountStatusDir",
   "Tac::MountTable",
   "Tac::MountTableProxy",
   "Tac::MsgChannel",
   "Tac::NboAttrLog",
   "Sysdb::AgentEnvStatusSm",
   "Sysdb::ConnectionConfig",
   "Sysdb::ConnectionConfigDir",
   "Sysdb::ConnectionManagerSm",
   "Sysdb::ConnectionManagerSmWrapper",
   "Sysdb::ConnectionSm",
   "Sysdb::ConnectionStatus",
   "Sysdb::ConnectionStatusDir",
   "Sysdb::GroupSpecificMountReference",
   "Sysdb::MountGroupDir",
   "Sysdb::MountReference",
]
mountRe = makeTypeRe( mountTypes )

activityTypes = [
   "Tac::ActivityConfig",
   "Tac::ActivityManagerConfig",
   "Tac::ActivityManagerSm",
   "Tac::ActivityManagerStatus",
   "Tac::ActivitySm",
   "Tac::ActivityStatus",
]
activityRe = makeTypeRe( activityTypes )

loggingTypes = [
   "Tac::LogConfig",
   "Tac::LogFacility",
   "Tac::LogFacilityConfig",
   "Tac::LogManager",
   "Tac::LogMsg",
   "Tac::LogStatus",
]
loggingRe = makeTypeRe( loggingTypes )

tracingTypes = [
   "Tac::TraceFacility",
   "Tac::TraceFacilityMan",
]
tracingRe = makeTypeRe( tracingTypes )

freeKinds = [ "<free>", "<top>" ]
specialKinds = freeKinds + [
   "<unknown>", "<c-string>", "<c++-string>", "<tac-string>", "<python-arena>",
   "<python-aux>" ]

KIND_UNKNOWN = 0
KIND_FREE = 1
KIND_CPLUSPLUS = 2
KIND_PYTHON = 3

def main():
   d = {}

   for line in sys.stdin:
      line = line.rstrip()
      m = re.match( r"  (\S+) (\d+) (\d+) (\d+) (\d+) (\d+)$", line )
      if not m:
         print( line )
         continue

      typename = cPlusPlusFilt( m.group( 1 ) )
      kindNum = int( m.group( 2 ) )
      objectCount = int( m.group( 3 ) )
      objectSize = int( m.group( 4 ) )
      uniqueObjectSize = int( m.group( 5 ) )
      proportionalObjectSize = int( m.group( 6 ) )

      if typename.endswith( ( "::GenericIf_", "::GenericIf_Adapter_" ) ):
         kind = "GenericIf"
      elif typename.endswith( ( "::TacNboAttrLogSa", "::TacNboAttrLogSaAdapter_" ) ):
         kind = "TacNboAttrLogSa"
      elif typename.startswith( "N3Tac13HashMapVTable" ):
         kind = "HashMapVTable"
      elif typename in specialKinds:
         kind = typename
      elif introspectionRe.match( typename ):
         kind = "<introspection>"
      elif dirRe.match( typename ):
         kind = "<directories>"
      elif mountRe.match( typename ):
         kind = "<mounts>"
      elif activityRe.match( typename ):
         kind = "<activities>"
      elif loggingRe.match( typename ):
         kind = "<logging>"
      elif tracingRe.match( typename ):
         kind = "<tracing>"
      elif kindNum == KIND_CPLUSPLUS:
         if '::' in typename:
            namespace = typename.split( "::", 1 )[ 0 ]
            kind = '::' + namespace
         else:
            kind = "Other C++"
      else:
         assert kindNum == KIND_PYTHON
         kind = "Other Python"

      vals = d.setdefault( kind, [ 0, 0, 0, 0 ] )
      vals[ 0 ] += objectCount
      vals[ 1 ] += objectSize
      vals[ 2 ] += uniqueObjectSize
      vals[ 3 ] += proportionalObjectSize

   formatStr = "%-25s %14s %14s %14s %14s"
   divider = "-" * 85

   if not d:
      # Invalid input data.
      return

   def sums( tuples ):
      return tuple( functools.reduce(
         lambda a, b: list( map( sum, zip( a, b ) ) ), tuples ) )

   print()
   print( formatStr % ( "Category", "objects", "bytes", "uss", "pss" ) )

   print( divider )
   for k in sorted( d ):
      if k in freeKinds:
         continue
      print( formatStr % ( ( k, ) + tuple( d[ k ] ) ) )
   print( formatStr % (
      # pylint: disable-next=consider-using-dict-items
      ( "Total Used", ) + sums( d[ k ] for k in d if k not in freeKinds ) ) )

   print( divider )
   for k in sorted( d ):
      if k not in freeKinds:
         continue
      print( formatStr % ( ( k, ) + tuple( d[ k ] ) ) )
   print( formatStr % (
      # pylint: disable-next=consider-using-dict-items
      ( "Total Free", ) + sums( d[ k ] for k in d if k in freeKinds ) ) )

   print( divider )
   # pylint: disable-next=consider-using-dict-items
   print( formatStr % ( ( "TOTAL", ) + sums( d[ k ] for k in d ) ) )

if __name__ == "__main__":
   main()

