import sys
import argparse
import re

entryformat = '{:>10} {:>10}    {}'

class Costs:
	def __init__(self, msize, vsize):
		self.msize = msize
		self.vsize = vsize
		
class CostsDict(dict):
	pass
	
class Bucket:
	def __init__(self, callstacks, messages):
		self.callstacks = callstacks
		self.messages = messages
		
class BucketDict(dict):
	pass
		
class ResourceMemory:
	def __init__(self, totals, scenecosts, loaded, externalallocs, buckets):
		self.totals = totals
		self.scenecosts = scenecosts
		self.loaded = loaded
		self.externalallocs = externalallocs
		self.buckets = buckets
		
def read_costs_helper(stream, guard = None):
	costs = CostsDict()

	for line in stream:
		if guard is not None and guard in line:
			break

		match = re.match( '(?P<name>.*), (?P<msize>\d*)kb, (?P<vsize>\d*)kb', line )
		if match:
			name = match.group( 'name' )
			msize = int( match.group( 'msize' ) )
			vsize = int( match.group( 'vsize' ) )
			costs[name] = Costs( msize, vsize )

	return costs

def read_moduleinfo(filename):
	costs = CostsDict()

	with open( filename, 'r' ) as stream:
		for line in stream:
			try:
				parts = [ part.strip() for part in line.split() ]
				name, msize, vsize = parts[1], int( parts[2].replace('K','') ), int( parts[3].replace('K','') )
				costs[name] = Costs( msize,vsize )
			except Exception as ex:		# nasty but quick way to ignore the lines that don't match the parsing above
				pass

	return costs

def read_scenecosts(filename):
	with open( filename, 'r' ) as stream:
		return read_costs_helper( stream )

def read_resourcememory(filename):
	with open( filename, 'r' ) as stream:
		totals = read_costs_helper( stream, '** END LOG: Totals **' )
		scenecosts = read_costs_helper( stream, '** END LOG: Scene Cost **' )
		loaded = read_costs_helper( stream, '** END LOG: Loaded Required Files Cost **' )
		externalallocs = read_costs_helper( stream, 'Detailed Breakdown (CALLSTACK)' )

		buckets = BucketDict()
		for bucket in externalallocs:
			callstacks = read_costs_helper( stream, 'End Detailed Breakdown (CALLSTACK)' )
			messages = read_costs_helper( stream, 'End Detailed Breakdown (CONTEXT MESSAGE)' )
			buckets[ bucket ] = Bucket( callstacks, messages )
			
		return ResourceMemory( totals, scenecosts, loaded, externalallocs, buckets )

def compare_costs(costs1, costs2, name1, name2, includezero):

	def print_entries(entries, title):
		print( 80 * '-' )
		print( title )
		print('')
		print( entryformat.format( 'MAIN', 'VRAM', 'OBJECT NAME' ) )

		entries.sort( key = lambda x: x[2], reverse = True )
		mtotal = 0
		vtotal = 0
		for key, msize, vsize in entries:
			if (not includezero) and (msize == 0 and vsize == 0):
				continue

			mtotal += msize
			vtotal += vsize
			print( entryformat.format( msize, vsize, key ) )

		print('')
		print( entryformat.format( mtotal, vtotal, '*** TOTAL ***' ) )
		print('')
		
	entries = [ (name, costs1[name].msize, costs1[name].vsize) for name in costs1.viewkeys() - costs2.viewkeys() ]
	print_entries( entries, 'Contained only in {}'.format( name1 ) )

	entries = [ (name, costs2[name].msize, costs2[name].vsize) for name in costs2.viewkeys() - costs1.viewkeys() ]
	print_entries( entries, 'Contained only in {}'.format( name2 ) )

	entries = [ (name, costs2[name].msize - costs1[name].msize, costs2[name].vsize - costs1[name].vsize) for name in costs1.viewkeys() & costs2.viewkeys() ]
	print_entries( entries, 'Contained in both files (size differences)' )
	
def compare_costs_concise(costs1, costs2, title, includezero, callstacks = False):
	entries = [ ('{} (REMOVED)'.format(name), -costs1[name].msize, -costs1[name].vsize) for name in costs1.viewkeys() - costs2.viewkeys() ]
	entries += [ ('{} (NEW)'.format(name), costs2[name].msize, costs2[name].vsize) for name in costs2.viewkeys() - costs1.viewkeys() ]
	entries += [ (name, costs2[name].msize - costs1[name].msize, costs2[name].vsize - costs1[name].vsize) for name in costs1.viewkeys() & costs2.viewkeys() ]
	
	mtotal = sum( msize for name, msize, vsize in entries )
	vtotal = sum( vsize for name, msize, vsize in entries )
	
	if mtotal is 0 and vtotal is 0:
		return
		
	entries.sort( key = lambda x: x[2], reverse = True )
	
	print( 80 * '-' )
	print( title )
	print( '' )
	print( entryformat.format( 'MAIN', 'VRAM', 'NAME' if not callstacks else 'CALLSTACK' ) )
	
	for name, msize, vsize in entries:
		if msize is not 0 or vsize is not 0 or includezero:
			if callstacks:
				frames = [ frame.strip() for frame in name.split('|') if len( frame.strip() ) > 0 ]
				print( entryformat.format( msize, vsize, frames[0] ) )
				for frame in frames[1:]:
					print( entryformat.format( '', '', frame ) )
				print( '' )
			else:
				print( entryformat.format( msize, vsize, name ) )
	
	print( entryformat.format( mtotal, vtotal, '*** TOTAL ***' ) )
	print( '' )
	
def compare_scenecosts(args):
	costs2 = read_scenecosts( args.file2 )
	costs1 = read_scenecosts( args.file1 )
	compare_costs( costs1, costs2, args.file1, args.file2, args.includezero )
	
def compare_moduleinfo(args):
	costs1 = read_moduleinfo( args.file1 )
	costs2 = read_moduleinfo( args.file2 )
	compare_costs( costs1, costs2, args.file1, args.file2, args.includezero )
	
def compare_resourcememory(args):
	resmem1 = read_resourcememory( args.file1 )
	resmem2 = read_resourcememory( args.file2 )
	
	compare_costs_concise( resmem1.totals, resmem2.totals, 'Totals', args.includezero )
	compare_costs_concise( resmem1.scenecosts, resmem2.scenecosts, 'Scene Costs', args.includezero )
	compare_costs_concise( resmem1.loaded, resmem2.loaded, 'Loaded Required Files', args.includezero )
	compare_costs_concise( resmem1.externalallocs, resmem2.externalallocs, 'External Allocations', args.includezero )
	
	for bucket in resmem1.buckets.viewkeys() | resmem2.buckets.viewkeys():
		callstacks1 = resmem1.buckets[ bucket ].callstacks if bucket in resmem1.buckets else CostsDict()
		callstacks2 = resmem2.buckets[ bucket ].callstacks if bucket in resmem2.buckets else CostsDict()
		messages1 = resmem1.buckets[ bucket ].messages if bucket in resmem1.buckets else CostsDict()
		messages2 = resmem2.buckets[ bucket ].messages if bucket in resmem2.buckets else CostsDict()
		
		compare_costs_concise( callstacks1, callstacks2, 'Bucket {} (CALLSTACK)'.format(bucket), args.includezero, callstacks = True )
		compare_costs_concise( messages1, messages2, 'Bucket {} (CONTEXT MESSAGE)'.format(bucket), args.includezero )
	
def parse_args(argv):
	parser = argparse.ArgumentParser( description = 'Compares two lists of streaming objects as returned by Scene Costs tracker or Dump Module info.' )
	parser.add_argument( '-s', '--scenecosts', action = 'store_true', help = 'Compares Scene Cost lists' )
	parser.add_argument( '-m', '--moduleinfo', action = 'store_true', help = 'Compares Dump Module Info lists' )
	parser.add_argument( '-r', '--resourcememory', action = 'store_true', help = 'Compares Resource Memory logs' )
	parser.add_argument( '-z', '--includezero', action = 'store_true', help = 'Include zero-sized costs in the output' )
	parser.add_argument( 'file1', help = 'The first file to compare' )
	parser.add_argument( 'file2', help = 'The second file to compare' )
	args = parser.parse_args( argv )

	num_commands = sum( 1 if x is True else 0 for x in (args.scenecosts, args.moduleinfo, args.resourcememory) )
	if num_commands is not 1:
		print( 'Error: you have to specify exactly one between --scenecosts, --moduleinfo and --resourcememory.\nUse -h/--help for usage help.' )
		sys.exit(1)
	
	return args

if __name__ == '__main__':
	args = parse_args( sys.argv[1:] )
	
	if args.scenecosts:
		compare_scenecosts( args )
	elif args.moduleinfo:
		compare_moduleinfo( args )
	elif args.resourcememory:
		compare_resourcememory( args )