#!/usr/bin/python3.7

# Script Analyze WoT Blitz replays

import sys, argparse, json, os, inspect, aiohttp, asyncio, aiofiles, aioconsole, re, logging, time, xmltodict, collections
from beautifultable import BeautifulTable
import blitzutils as bu
from blitzutils import WG


logging.getLogger("asyncio").setLevel(logging.DEBUG)

WG_appID  = '81381d3f45fa4aa75b78a7198eb216ad'
wg = None
REPLAY_N = 0
STAT_TANK_BATTLE_MIN = 100

replay_summary_flds = [
	'battle_result',
	'battle_type',
	'map_name',
	'battle_duration',
	'protagonist'
]

replay_details_flds = [
	'damage_made',
	'damage_assisted_track',
	'damage_assisted',
	'damage_received',
	'damage_blocked',
	'enemies_damaged',
	'enemies_destroyed',
	'hitpoints_left',
	'time_alive',
	'distance_travelled',
	'base_capture_points',
	'base_defend_points',
	'wp_points_earned',
	'wp_points_stolen',
	'hits_received',
	'hits_bounced',
	'hits_splash',
	'hits_pen',
	'shots_made',
	'shots_hit',
	'shots_pen',
	'shots_splash',
	'enemies_spotted'
	]

result_ratios = {
	'KDR'				: [ 'enemies_destroyed', 'destroyed' ],
	'DR'				: [ 'damage_made', 'damage_received' ],
	'hit_rate'			: [ 'shots_hit', 'shots_made' ],
	'pen_rate'			: [ 'shots_pen', 'shots_hit' ],
	'dmg_block_rate' 	: [ 'damage_blocked', 'damage_received' ]
}

result_categories = {
	'total'				: [ 'TOTAL', 'Total' ],
	'battle_result'		: [ 'Result', ['Win', 'Loss', 'Draw']],
	'battle_type'		: [ 'Battle Type',['Encounter', 'Supremacy']],
	'battle_tier'		: [ 'Battle Tier', 'value' ],
	'top_tier'			: [ 'Tier', ['Bottom tier', 'Top tier']],
	'map_name'			: [ 'Map', 'value']		
	}

result_cat_header_frmt = '{:_<15s}'
result_cat_frmt = '{:>15s}'

result_fields = {
	'battles'			: [ 'Battles', 'Number of battles', 8, '{:^8.0f}' ],
	'win'				: [ 'WR', 'Win rate', 			6, '{:6.1%}' ],
	'damage_made'		: [ 'DPB', 'Average Damage', 	5, '{:5.0f}' ],
	'DR'				: [ 'DR', 'Damage Ratio', 		5, '{:5.1f}' ],
	'KDR'				: [ 'KDR', 'Kills / Death', 	4, '{:4.1f}' ],
	'enemies_spotted'	: [ 'Spot', 'Enemies spotted per battle', 		4, '{:4.1f}' ],
	'hit_rate'			: [ 'Hit rate', 'Shots hit / all shots made', 	8, '{:8.1%}' ],
	'pen_rate'			: [ 'Pen rate', 'Shots pen / shots hit', 		8, '{:8.1%}' ],
	'survived'			: [ 'Surv%', 'Survival rate', 						6, '{:6.1%}' ],
	'time_alive%'		: [ 'T alive%', 'Percentage of time being alive in a battle', 8, '{:8.0%}' ], 
	'top_tier'			: [ 'Top tier', 'Share of games as top tier', 	8, '{:8.0%}' ],
	'allies_win_rate'	: [ 'Allies WR', 'Average WR of allies: WR in the tank played counted if more than ' + str(STAT_TANK_BATTLE_MIN) + ' in the tank', 9, '{:9.2%}' ],
	'enemies_win_rate'	: [ 'Enemies WR', 'Average WR of enemies: WR in the tank played counted if more than ' + str(STAT_TANK_BATTLE_MIN) + ' in the tank', 10, '{:10.2%}' ],
	'allies_battles'	: [ 'Allies Btls', 'Average number battles of the allies', 		11, '{:11.0f}' ],
	'enemies_battles'	: [ 'Enemies Btls', 'Average number battles of the enemies', 	12, '{:12.0f}' ]
	
}

battle_stats_fields = {
	'WR'			: 'Win rate',
	'Battles'		: 'Battles'
	}

def defaultvalueZero():
    return 0

def defaultvalueBtlRec():
	return BattleRecord()

class BattleRecord():
	def __init__(self):
		try:
			self.battles = 0
			self.ratios_ready = False
			self.results_ready = False
			self.results = collections.defaultdict(defaultvalueZero)
			self.avg_fields = set(result_fields.keys())
			self.avg_fields.remove('battles')
			self.avg_fields.difference_update(result_ratios.keys())
			self.ratio_fields = set()
			for ratio in result_ratios.keys():
				self.ratio_fields.add(result_ratios[ratio][0])
				self.ratio_fields.add(result_ratios[ratio][1])

		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 

	def recordResult(self, result : dict) -> bool:
		try:
			for field in self.avg_fields | self.ratio_fields:
				if result[field] != None:
					self.results[field] += result[field]
			self.battles += 1
			return True
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		return False

	def calcRatios(self) -> bool:
		if self.ratios_ready:
			return True
		try:
			ratio_result_flds = set(result_ratios.keys()) & set(result_fields.keys())
			
			for field in ratio_result_flds:
				if self.results[result_ratios[field][1]] != 0:
					self.results[field] = self.results[result_ratios[field][0]] / self.results[result_ratios[field][1]]
				else:
					self.results[field] = float('Inf')
			self.ratios_ready = True
			return True
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found')
		return False

	def calcResults(self):
		if self.results_ready == True: 
			return True
		if not self.ratios_ready:
			self.calcRatios()
		try:
			for field in self.avg_fields:
				self.results[field] = self.results[field] / self.battles
			self.results['battles'] = self.battles
			self.results_ready = True
			return True
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		except Exception as err:
			bu.error(str(err)) 
		return False

	def getResults(self):
		if not self.results_ready:
			self.calcResults()
		results = []
		try:
			for field in result_fields.keys():
				results.append( result_fields[field][3].format(self.results[field]) )
			return results
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		except Exception as err:
			bu.error(str(err)) 
		return None
	
	def printResults(self):
		print(' : '.join(self.getResults()))


class BattleRecordCategory():
	def __init__(self, cat_name : str):
		self.category_name = cat_name
		self.category = collections.defaultdict(defaultvalueBtlRec)
	
	def getSubCategories(self):
		return self.category.keys()

	def recordResult(self, result: dict):
		try:
			if self.category_name == 'total':
				cat = 'Total'
			elif result_categories[self.category_name][1] == 'value':
				cat = str(result[self.category_name])
			else:
				cat = result_categories[self.category_name][1][result[self.category_name]]
			self.category[cat].recordResult(result)
			return True
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		except Exception as err:
			bu.error(str(err)) 
		return False
	
	def getHeaders(self):
		try:
			headers = [  result_cat_header_frmt.format(result_categories[self.category_name][0]) ]
			for field in result_fields.keys():
				print_format = '{:^' + str(result_fields[field][2]) + 's}'
				headers.append(print_format.format(result_fields[field][0]))
			return headers
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		except Exception as err:
			bu.error(str(err)) 
		return None

	def printResults(self):
		try:
			print('   '.join(self.getHeaders()))
			for row in self.getResults():
				print(' : '.join(row))
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		except Exception as err:
			bu.error(str(err)) 
		return None

	def getResults(self):
		try:
			results = []
			# results.append(self.getHeaders())			
			for cat in self.category.keys():
				row = [ result_cat_frmt.format(cat)  ]
				row.extend(self.category[cat].getResults())
				results.append(row)
			return results
		except KeyError as err:
			bu.error('Key ' + str(err) + ' not found') 
		return None

## main() -------------------------------------------------------------

async def main(argv):
	global wg
	
	TASK_N = 7
	WG_ID = 2008897346
	parser = argparse.ArgumentParser(description='ANalyze Blitz replay JSONs from WoTinspector.com')
	parser.add_argument('--output', default='plain', choices=['json', 'plain', 'db'], help='Select output mode: JSON, plain text or database')
	parser.add_argument('-id', dest='accountID', type=int, default=WG_ID, help='WG account_id to analyze')
	parser.add_argument('-a', '--account', dest='account', type=str, default=None, help='WG account nameto analyze. Format: ACCOUNT_NAME@SERVER')
	parser.add_argument('-u', '--url', dest= 'url', action='store_true', default=False, help='Print replay URLs')
	parser.add_argument('--tankfile', type=str, default='tanks.json', help='JSON file to read Tankopedia from')
	parser.add_argument('--mapfile', type=str, default='maps.json', help='JSON file to read Blitz map names from')
	parser.add_argument('-o','--outfile', type=str, default='-', metavar="OUTPUT", help='File to write results. Default STDOUT')
	parser.add_argument('-d', '--debug', action='store_true', default=False, help='Debug mode')
	parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Verbose mode')
	parser.add_argument('files', metavar='FILE1 [FILE2 ...]', type=str, nargs='+', help='Files to read. Use \'-\' for STDIN"')
	
	args = parser.parse_args()
	bu.setVerbose(args.verbose)
	bu.setDebug(args.debug)
	wg = WG(WG_appID, args.tankfile, args.mapfile)

	if (args.accountID == None) and (args.account != None):
		args.accountID = await wg.getAccountID(args.account)
		bu.debug('WG  account_id: ' + str(args.accountID))

	try:
		replayQ  = asyncio.Queue()	

		#resultsQ = asyncio.Queue()
		
		reader_tasks = []
		# Make replay Queue
		scanner_task = asyncio.create_task(mkReplayQ(replayQ, args.files))

		# Start tasks to process the Queue
		for i in range(TASK_N):
			reader_tasks.append(asyncio.create_task(replayReader(replayQ, i, args)))
			bu.debug('Task ' + str(i) + ' started')

		bu.debug('Waiting for the replay scanner to finish')
		await asyncio.wait([scanner_task])
		bu.debug('Scanner finished. Waiting for replay readers to finish the queue')
		await replayQ.join()
		bu.debug('Replays read. Cancelling Readers and analyzing results')
		for task in reader_tasks:
			task.cancel()	
		results = []
		playerstanks = set()
		for res in await asyncio.gather(*reader_tasks):
			results.extend(res[0])
			playerstanks.update(res[1])

		player_stats = await processTankStats(playerstanks, TASK_N)
		bu.verbose('')
		results = calcTeamStats(results, player_stats)
		processStats(results, args)	

	finally:
		## Need to close the aiohttp.session since Python destructors do not support async methods... 
		await wg.session.close()
	return None

def processStats(results: dict, args : argparse.Namespace):

	categories = {}
	for cat in result_categories.keys():
		categories[cat] = BattleRecordCategory(cat)

	for result in results:
		for cat in result_categories.keys():
			categories[cat].recordResult(result)

	for cat in result_categories.keys():
		for sub_cat in categories[cat].getSubCategories():
			categories[cat].category[sub_cat].calcResults()
		print('')
		categories[cat].printResults()

	return None

async def processTankStats(playerstanks, N_workers: int) -> dict:
	## Create queue of player-tank pairs to find stats for
	statsQ = asyncio.Queue()
	bu.debug('Create player/tank stats queue')
	for playertank in playerstanks:
		await statsQ.put(playertank)
	# Process player WR / battle stats
	stats_tasks = []
	bu.debug('Starting player/tank stats queue workers')
	for i in range(N_workers):
		stats_tasks.append(asyncio.create_task(statWorker(statsQ, i)))
	## let the workers to process	
	await statsQ.join()
	for task in stats_tasks: 
		task.cancel()	
	player_stats = {}
	for stats in await asyncio.gather(*stats_tasks):
		player_stats = {**player_stats, **stats}
	
	return player_stats

def calcTeamStats(result_list: list, player_stats  : dict) -> list:
	return_list = []
	stat_types = player_stats[list(player_stats.keys())[0]].keys()
	for result in result_list:
		try:
			n_allies = len(result['allies'])
			allies_stats = collections.defaultdict(defaultvalueZero)
			for ally in result['allies']:
				for stat in stat_types:
					if player_stats[ally][stat] == None:
						n_allies -= 1
					else:
						allies_stats[stat] += player_stats[ally][stat]
			for stat in stat_types:
				result['allies_' + str(stat)] = allies_stats[stat] / n_allies

			n_enemies = len(result['enemies'])
			enemies_stats = collections.defaultdict(defaultvalueZero)
			for enemy in result['enemies']:
				for stat in stat_types:
					if player_stats[enemy][stat] == None:
						n_enemies -= 1
					else:
						enemies_stats[stat] += player_stats[enemy][stat]
			for stat in stat_types:
				result['enemies_' + str(stat)] = enemies_stats[stat] / n_enemies


			return_list.append(result)
		except KeyError as err:
			bu.error('Key :' + str(err) + ' not found')
		except Exception as err:
			bu.error(err)
	
	return return_list

async def statWorker(queue : asyncio.Queue, workerID: int) -> list:
	"""Worker thread to find stats for player / tank pairs"""
	# global wg
	stats = {}

	try:
		while True:
			# item = await queue.get_nowait()
			item = await queue.get()
			try:
				bu.printWaiter()
				acc, tank = item.split(':')			
				account_id = int(acc)
				tank_id = int(tank)
				# bu.debug('[' +str(workerID) + '] AccountID: ' + acc + ' TankID: '  + tank)
				playerTankStat = await wg.getPlayerTankStats(account_id, tank_id, ['all.battles', 'all.wins'])
				# bu.debug('[' +str(workerID) + '] ' + str(playerTankStat))
				playerStat = await wg.getPlayerStats(account_id,  ['statistics.all.battles', 'statistics.all.wins'])
				# bu.debug('[' +str(workerID) + '] ' + str(playerStat))
			
				stats[item] = {}
				if (playerTankStat == None) or (playerStat == None):
					stats[item]['win_rate'] = None
					stats[item]['battles'] 	= None
				else:
					playerTankStat = playerTankStat['all']
					playerStat = playerStat['statistics']['all']

					battles = playerStat['battles']
					stats[item]['battles'] = battles

					battles_in_tank = playerTankStat['battles']
					if battles_in_tank >= STAT_TANK_BATTLE_MIN:
						stats[item]['win_rate'] =  min(playerTankStat['wins'] / battles_in_tank,1)  # To cope with broken stats in WG DB
					else:
						stats[item]['win_rate'] = min(playerStat['wins'] / battles, 1) # To cope with broken stats in WG DB
					bu.debug('[' +str(workerID) + '] Player[' + str(account_id) + '], Tank[' + str(tank_id) + '] : WR : ' + str(stats[item]['win_rate']) + ' Battles: ' + str(battles))
			except KeyError as err:
				bu.error('[' +str(workerID) + '] Key :' + str(err) + ' not found')
			except Exception as err:
				bu.error('[' +str(workerID) + '] ' + str(err))
			queue.task_done()

	except asyncio.CancelledError:
		bu.debug('Stats queue[' + str(workerID) + '] is empty')
	except Exception as err:
		bu.error(str(err))
	return stats
	
async def mkReplayQ(queue : asyncio.Queue, files : list):
	"""Create queue of replays to post"""
	p_replayfile = re.compile('.*\\.wotbreplay\\.json$')
	if files[0] == '-':
		bu.debug('reading replay file list from STDIN')
		stdin, _ = await aioconsole.get_standard_streams()
		while True:
			line = (await stdin.readline()).decode('utf-8').rstrip()
			if not line: 
				break
			else:
				if (p_replayfile.match(line) != None):
					await queue.put(await mkReaderQitem(line))
	else:
		for fn in files:
			if os.path.isfile(fn) and (p_replayfile.match(fn) != None):
				await queue.put(await mkReaderQitem(fn))
			elif os.path.isdir(fn):
				with os.scandir(fn) as dirEntry:
					for entry in dirEntry:
						if entry.is_file() and (p_replayfile.match(entry.name) != None): 
							bu.debug(entry.name)
							await queue.put(await mkReaderQitem(entry.path))
			bu.debug('File added to queue: ' + fn)
	bu.debug('Finished')
	return None

async def mkReaderQitem(filename : str) -> list:
	"""Make an item to replay queue"""
	global REPLAY_N
	REPLAY_N +=1
	return [filename, REPLAY_N]

async def replayReader(queue: asyncio.Queue, readerID: int, args : argparse.Namespace):
	"""Async Worker to process the replay queue"""
	#global SKIPPED_N
	account_id = args.accountID
	results = []
	playerstanks = set()
	try:
		while True:
			item = await queue.get()
			filename = item[0]
			replayID = item[1]

			try:
				if os.path.exists(filename) and os.path.isfile(filename):
					async with aiofiles.open(filename) as fp:
						replay_json = json.loads(await fp.read())
						if replay_json['status'] != 'ok':
							bu.verbose('Replay[' + str(replayID) + ']: ' + filename + ' is invalid. Skipping.' )
							#SKIPPED_N += 1
							queue.task_done()
							continue
						
						## Read the replay JSON
						bu.debug('[' + str(readerID) + '] reading replay: ' + filename)
						result = await readReplayJSON(replay_json, args)
						if result == None:
							bu.error('[' + str(readerID) + '] Error reading replay: ' + filename)
							queue.task_done()
							continue
						
						if (account_id != None):						
							playerstanks.update(set(result['allies']))
							playerstanks.update(set(result['enemies']))	
						
						results.append(result)
						bu.debug('Marking task ' + str(replayID) + ' done')
						
			except Exception as err:
				bu.error(err)
			queue.task_done()
	except asyncio.CancelledError:		
		return results, playerstanks
	return None

async def readReplayJSON(replay_json: dict, args : argparse.Namespace) -> dict:
	"""Parse replay JSON dict"""
	account_id = args.accountID
	url = args.url
	result = {}
	try:
		if account_id == None:
			account_id = replay_json['data']['summary']['protagonist'] 
		elif replay_json['data']['summary']['protagonist'] != account_id:
			if account_id in replay_json['data']['summary']['enemies']:
				# switch the teams...
				if replay_json['data']['summary']['battle_result'] != 2:
					if replay_json['data']['summary']['battle_result'] == 0:
						replay_json['data']['summary']['battle_result'] = 1
					else: 
						replay_json['data']['summary']['battle_result'] = 0
				tmp = replay_json['data']['summary']['enemies']
				replay_json['data']['summary']['enemies'] = replay_json['data']['summary']['allies']
				replay_json['data']['summary']['allies'] = tmp
			elif account_id not in replay_json['data']['summary']['allies']:
				# account_id looked for but not found in teams
				bu.debug('Replay ' + replay_json['data']['summary']['title'] + ' does not have account_id ' + str(account_id) + '. Skipping.')
				return None

		if url: 
			result['url'] = replay_json['data']['view_url']
		else:
			result['url'] = None

		for key in replay_summary_flds:
			result[key] = replay_json['data']['summary'][key]

		result['allies'] = set()
		result['enemies'] = set()

		btl_duration = 0
		btl_tier = 0
		for player in replay_json['data']['summary']['details']:
			btl_duration = max(btl_duration, player['time_alive'])
			player_tank_tier = wg.getTankData(player['vehicle_descr'], 'tier')
			btl_tier = max(btl_tier, player_tank_tier)
			if player['dbid'] == account_id:
				tmp = {}
				for key in replay_details_flds:
					tmp[key] = player[key]
				if player['hitpoints_left'] == 0:
					tmp['survived'] = 0
					tmp['destroyed'] = 1
				else:
					tmp['survived'] 	= 1
					tmp['destroyed'] 	= 0
				tmp['account_id'] 	= account_id
				tmp['tank_id'] 		= player['vehicle_descr']
				tmp['tank_tier'] 	= player_tank_tier
				
				for key in tmp.keys():
					result[key] = tmp[key]				

			elif player['dbid'] in replay_json['data']['summary']['allies']: 
				result['allies'].add(':'.join([ str(player['dbid']), str(player['vehicle_descr'])]))
			else:
				result['enemies'].add(':'.join([ str(player['dbid']), str(player['vehicle_descr'])]))
		result['time_alive%'] = result['time_alive'] / btl_duration  
		result['battle_tier'] = btl_tier
		result['top_tier'] = 1 if (result['tank_tier'] == btl_tier) else 0
		result['win'] = 1 if result['battle_result'] == 1 else 0
				
		bu.debug(str(result))
		return result
	except KeyError as err:
		bu.error('Key :' + str(err) + ' not found')
	except Exception as err:
		bu.error(err)
	return None

### main()
if __name__ == "__main__":
   asyncio.run(main(sys.argv[1:]), debug=True)
   #asyncio.run(main(sys.argv[1:]))