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

from collections import namedtuple
import argparse
import sys
import re
import subprocess
from prettytable import PrettyTable

# Definition for output
ProfileEntry = namedtuple( 'ProfileEntry',
                           [ 'index', 'count', 'datetime',
                             'averageStackTime', 'totalStackTime',
                             'averageSelfTime', 'totalSelfTime',
                             'totalPercent',
                             'location', 'message', 'native' ] )

# Simple class to "formalize" align options from prettytable
class Justify:
   Left = 'l'
   Right = 'r'
   Center = 'c'

# Maps the column headers to the the ProfileEntry fields with formatting info
columnToAttrMap = {
   'index':( 'index', int, Justify.Right ),
   'count':( 'count', int, Justify.Right ),
   'lastTime':( 'datetime', str, Justify.Left ),
   'avgTime':( 'averageStackTime', float, Justify.Right ),
   'totalTime':( 'totalStackTime', float, Justify.Right ),
   'avgSelfTime':( 'averageSelfTime', float, Justify.Right ),
   'totalSelfTime':( 'totalSelfTime', float, Justify.Right ),
   'totalPercent':( 'totalPercent', float, Justify.Right ),
   'location':( 'location', str, Justify.Left ),
   'message':( 'message', str, Justify.Left ),
   'native': ( 'native', str, Justify.Left ),
}

# Regex's used for parsing the qtcat output
indexExp = r'\s+(?P<index>\d+)'
countExp = r'\s+(?P<count>\S+)'
dateTimeExp = r'\s+(?P<dateTime>\d+-\d+-\d+ \d+:\d+:\d+.\d+)'
timestampExp = r'\s+(\S+)'
locationExp = r'\s+(?P<location>\S+)'
messageExp = r'\s+(?P<message>.*)'

# the standard line is missing a couple columns present in the self-profile lines
stdLineRe = re.compile( indexExp + countExp + dateTimeExp +
                        timestampExp + timestampExp +
                        locationExp + messageExp )
selfLineRe = re.compile( indexExp + countExp + dateTimeExp +
                         timestampExp + timestampExp +
                         timestampExp + timestampExp +
                         locationExp + messageExp )
locationDetailRe = re.compile(
   r'/bld/(?P<package>\S+)/\S+/\S+/(?P<filename>\S+):(?P<fileline>\d+)' )

def timestamp( s ):
   """Remove 's' from the input value"""
   return s.replace( 's', '' )

# import from qtparse?
def hrToInt( s ):
   """
   Conversion from human readable integer to plain integer, so that
   spreadsheets can deal with the result.
   """
   if s.find( "Px" ) >= 0:
      tmp = s.replace( "Px", "000000000000000" )
   elif s.find( "Tx" ) >= 0:
      tmp = s.replace( "Tx", "000000000000" )
   elif s.find( "Gx" ) >= 0:
      tmp = s.replace( "Gx", "000000000" )
   elif s.find( "Mx" ) >= 0:
      tmp = s.replace( "Mx", "000000" )
   elif s.find( "Kx" ) >= 0:
      tmp = s.replace( "Kx", "000" )
   else:
      tmp = s.replace( "x", "" )
   return int( tmp )

def getEntryString( entry, column ):
   """
   Gets the string for the specified column from the entry.
   """
   return str( getattr( entry, columnToAttrMap[ column ][ 0 ] ) )

def getEntryValue( entry, column ):
   """
   Gets the value of the appropriate type for the column, with some special
   sauce to accomodate '-' that qtcat supplies when numeric values are 0.
   """
   attrType = columnToAttrMap[ column ][ 1 ]
   value = getEntryString( entry, column )
   if value == '-' and attrType in ( int, float ):
      value = '0'
   return attrType( value )

def getTextFromQtFile( qtfilename ):
   qtCatCmd = [ 'qtcat', '--profile', '--selfProfiling', '--parsable', qtfilename ]
   text = subprocess.check_output( qtCatCmd, text=True )
   return text.split( '\n' )

def getProfileEntriesFromText( fh ):
   """Parses the lines form the input and provides a list of ProfileEntry objects
   corresponding to the different lines."""
   entries = []
   for line in fh:
      m = selfLineRe.search( line )
      if m:
         averageSelfTime = timestamp( m.group( 6 ) )
         totalSelfTime = timestamp( m.group( 7 ) )
      else:
         averageSelfTime = 0.0
         totalSelfTime = 0.0
         m = stdLineRe.search( line )
         if not m:
            continue

      # this is the BIG location...
      location = m.group( 'location' )

      # attempt to reduce full paths to 'package/file:line' -- in some cases
      # the location does not contain the full path (not sure why)
      locDetail = locationDetailRe.search( location )
      if locDetail:
         location = '%s/%s:%s' % (
            locDetail.group( 'package' ),
            locDetail.group( 'filename' ),
            locDetail.group( 'fileline' ) )

      # convert 'count' to raw integers
      count = hrToInt( m.group( 'count' ) )

      averageStackTime = timestamp( m.group( 4 ) )
      totalStackTime = timestamp( m.group( 5 ) )

      # make a value-added calculation
      if ( totalStackTime == '-' or totalSelfTime == '-' or
           float( totalStackTime ) == 0.0 or float( totalSelfTime ) == 0.0 ):
         totalPercent = 0.0
      else:
         totalPercent = 100.0 * float( totalSelfTime ) / float( totalStackTime )
      totalPercent = "%0.2f" % totalPercent

      e = ProfileEntry( index=m.group( 'index' ),
                        count=count,
                        datetime=m.group( 'dateTime' ),
                        averageStackTime=averageStackTime,
                        totalStackTime=totalStackTime,
                        averageSelfTime=averageSelfTime,
                        totalSelfTime=totalSelfTime,
                        totalPercent=totalPercent,
                        location=location,
                        message=m.group( 'message' ),
                        native=line
      )
      entries.append( e )

   return entries

def renderProfileEntrySeparated( _entries, columns, separator ):
   output = [ separator.join( columns ) ]
   for entry in _entries:
      output.append( separator.join( [ getEntryString( entry, c )
                                       for c in columns ] ) )
   return "\n".join( output )

def createProfileEntryTable( _entries, columns ):
   """
   Provides common means of creating a PrettyTable for entries.
   """
   table = PrettyTable( field_names=columns, border=False )
   # pylint: disable-msg=protected-access
   # The alignment methods only allow for one alignment for the whole table, but
   # it is desirable/supported to have different alignments per column. So, this
   # sets it directly.
   table._align = { c: columnToAttrMap[ c ][ 2 ] for c in columns }
   # The desired padding combination does not appear to be settable to the desired
   # values during construction, so override them here.
   table._set_padding_width( 0 )
   table._set_right_padding_width( 1 )
   # pylint: enable-msg=protected-access

   for entry in _entries:
      table.add_row( [ getEntryString( entry, c ) for c in columns ] )

   return table

def renderProfileEntryTableText( _entries, columns ):
   table = createProfileEntryTable( _entries, columns )

   # remove trailing whitespace from the last column
   output = table.get_string().split( '\n' )
   output = [ line.rstrip() for line in output ]
   output.append( '' )
   return '\n'.join( output )

def renderProfileEntryTableHtml( _entries, columns ):
   table = createProfileEntryTable( _entries, columns )
   return table.get_html_string()

def printProfileEntryTable( _entries, renderFormat, columns ):
   output = "Unhandled render format: %s" % renderFormat
   if renderFormat == 'text':
      output = renderProfileEntryTableText( _entries, columns )
   elif renderFormat == 'html':
      output = renderProfileEntryTableHtml( _entries, columns )
   elif renderFormat == 'csv':
      output = renderProfileEntrySeparated( _entries, columns, ',' )
   elif renderFormat == 'tsv':
      output = renderProfileEntrySeparated( _entries, columns, '\t' )
   print( output )

def parseArgs( sysargs ):
   defaultColumns = [
      #'index',
      'count',
      #'lastTime',
      'avgTime',
      'totalTime',
      'avgSelfTime',
      'totalSelfTime',
      #'totalPercent',
      #'location',
      'message',
   ]
   outputOptions = [ 'text', 'csv', 'tsv', 'html', 'raw' ]
   columnOptions = list( columnToAttrMap )
   orderOptions = list( columnToAttrMap )
   description = ( "Reads profiling data from QuickTrace file, and provides means "
                   "of sorting, and other display adjustments." )
   parser = argparse.ArgumentParser( description=description )
   parser.add_argument( dest='inQtFile', metavar='FILENAME', nargs='?',
                        help="QuickTrace file from which to read profiling data" )
   parser.add_argument( '--text', dest='inText', metavar='FILENAME',
                        help=( "Text file to parse for captured output, or '-' to "
                               "read from stdin (instead of QuickTrace file)." ) )
   parser.add_argument( '--columns', dest='columns', nargs='+',
                        metavar='col',
                        choices=columnOptions,
                        help=( "Ordered list of columns to display. Options are " +
                               ", ".join( columnOptions ) + "." ),
                        default=defaultColumns )
   parser.add_argument( '--format', dest='outFmt', default=outputOptions[ 0 ],
                        choices=outputOptions,
                        metavar='<' + '|'.join( outputOptions ) + '>',
                        help="Output format type (default=%(default)s)" )
   parser.add_argument( '--order', dest='order',
                        choices=orderOptions,
                        metavar='<' + '|'.join( orderOptions ) + '>',
                        help="Order according to column" )
   parser.add_argument( '--reverse', action='store_true',
                        help='Reverse the sort order when specified' )
   parser.add_argument( '--head', dest='headCount', type=int,
                        help="Maximum profile entries at the head" )
   parser.add_argument( '--tail', dest='tailCount', type=int,
                        help="Maximum profile entries at the tail" )
   return parser.parse_args( sysargs )

def qtprofHandler( sysargs ):
   args = parseArgs( sysargs )

   textFile = args.inText
   if textFile:
      if textFile == '-':
         text = sys.stdin
      else:
         text = open( textFile ) # pylint: disable=consider-using-with
   elif args.inQtFile:
      text = getTextFromQtFile( args.inQtFile )
   else:
      print( "Either a QuickTrace filename, or --text argument must be provided." )
      return -1

   if args.outFmt == 'raw' :
      print( "\n".join( text ) )
      return 0

   entries = getProfileEntriesFromText( text )

   if args.order:
      entries = sorted( entries, reverse=args.reverse,
                        key=lambda entry: getEntryValue( entry, args.order ) )

   if args.headCount:
      entries = entries[ : args.headCount ]
   if args.tailCount:
      entries = entries[ -args.tailCount : ]

   printProfileEntryTable( entries, args.outFmt, args.columns )
   return 0

if __name__ == '__main__':
   qtprofHandler( sys.argv[ 1: ] )
