#!/usr/bin/env python3

from collections import namedtuple
from itertools import *
import argparse
import os
import re
import subprocess
import toml
import xml.etree.ElementTree as ET

reset     = '\033[0m'
bold      = '\033[1m'
dim       = '\033[2m'
underline = '\033[4m'
blink     = '\033[5m'
inverted  = '\033[7m'
hidden    = '\033[8m'

black   = '\033[30m'
red     = '\033[31m'
green   = '\033[32m'
yellow  = '\033[33m'
blue    = '\033[34m'
magenta = '\033[35m'
cyan    = '\033[36m'
white   = '\033[37m'

bright_black   = '\033[90m'
bright_red     = '\033[91m'
bright_green   = '\033[92m'
bright_yellow  = '\033[93m'
bright_blue    = '\033[94m'
bright_magenta = '\033[95m'
bright_cyan    = '\033[96m'
bright_white   = '\033[97m'

fancy_print_enabled = True

def fancy_print(text, color, *styles, newline=True, indent=0):
	builder = []
	if fancy_print_enabled:
		builder.append(color)
		builder.extend(styles)
		builder.append(text)
		builder.append(white)
		builder.append(reset)
	else:
		builder.append(text)
	lines = map(lambda x: (' ' * indent) + x + '\n', ''.join(builder).split('\n'))
	if newline:
		builder.append('\n')
	print(''.join(lines), end='')

def validate_directory(directory):
	# Ensure that directory exists
	if not os.path.isdir(directory):
		print('\'{0}\' is not a directory'.format(directory))
		exit(1)
	# Ensure that manifest.toml exists
	if 'manifest.toml' not in os.listdir(directory):
		print('No manifest.toml found in \'{0}\''.format(directory))
		exit(1)

def load_manifest(path):
	with open(path, 'r') as f:
		text = f.read()
		data = toml.loads(text)
	return data

Header = namedtuple('Header', 'file out_type')

def parse_header(src):
	src = src.strip()
	regex = re.compile(r"\"([^\"\n]*)\"[\ ]*(out|error)")
	match = regex.match(src)
	if not match:
		print('Invalid header \'{0}\''.format(src))
		exit(1)
	return Header(match.group(1), match.group(2))

def read_outfile(path, sigil):
	with open(path, 'r') as f:
		text = f.read()
	sections = text.split(sigil)
	sections = sections[1:] # Ditch everything before first sigil
	def split_section(section):
		part = section.partition('\n')
		return part[0], part[2]
	sections = [(parse_header(h), b) for h, b in map(split_section, sections)]
	return sections

def test_section(header, body, directory, options):
	indent = 2
	path = os.path.join(directory, header.file)
	shortpath = os.path.basename(path)
	if not os.path.isfile(path):
		print('Source file \'{0}\' does not exist!'.format(path))
	command = options['command'].replace('%%', path)
	result = subprocess.run(command, shell=True, capture_output=True)
	succeeded = False
	def success():
		nonlocal succeeded
		succeeded = True
		fancy_print('{0} was a success!'.format(shortpath), green, bold, indent=indent)
	def print_context():
		fancy_print('stdout:', green, indent=indent)
		if options['output_bin']:
			fancy_print(str(result.stdout), dim, indent=indent*2)
		else:
			fancy_print(result.stdout.decode(), dim, indent=indent*2)
		fancy_print('stderr:', red, indent=indent)
		if options['output_bin']:
			fancy_print(str(result.stderr), dim, indent=indent*2)
		else:
			fancy_print(result.stderr.decode(), dim, indent=indent*2)
	if header.out_type == 'out':
		if result.returncode != 0:
			fancy_print('{0} recieved an error.'.format(shortpath), red, bold, indent=indent)
			print_context()
		else:
			if result.stdout.decode() == body:
				success()
			else:
				fancy_print('{0} got the wrong output.'.format(shortpath), red, bold, indent=indent)
				print_context()
				fancy_print('expected stdout:', blue, bold, indent=indent)
				if options['output_bin']:
					fancy_print(str(body.encode()), dim, indent=indent*2)
				else:
					fancy_print(body, dim, indent=indent*2)
	elif header.out_type == 'error':
		if result.returncode != 0:
			success()
		else:
			fancy_print('{0} expected an error, but didn\'t get one.'.format(shortpath), red, bold, indent=indent)
			print_context()
	else:
		assert(False) # Regular expression should prevent this from ever happening
	return succeeded

def test_section_for_leaks(header, body, directory, options):
	if header.out_type == 'error':
		return True
	path = os.path.join(directory, header.file)
	shortpath = os.path.basename(path)
	if not os.path.isfile(path):
		print('Source file \'{0}\' does not exist!'.format(path))
	command = options['command'].replace('%%', path)
	command = 'valgrind --leak-check=full --show-leak-kinds=all --xml=yes --xml-fd=1 {0}'.format(command)
	result = subprocess.run(command, shell=True, capture_output=True)
	# Parse XML
	tree = ET.fromstring(result.stdout)
	if any(filter(lambda x: x.tag == 'error', tree)) > 0:
		fancy_print('   Leaked!', red, bold)
		return False
	return True

def test_from_directory(path, options):
	if 'out' not in os.listdir(path):
		print('\'{0}\' contains no out file'.format(path))
		return True
	sections = read_outfile(os.path.join(path, 'out'), options['sigil'])
	fancy_print('{0}'.format(path), reset)
	any_failed = False
	for section in sections:
		if not test_section(*section, path, options):
			any_failed = True
		elif options['leak_check']:
			if not test_section_for_leaks(*section, path, options):
				any_failed = True
	return any_failed

def main():
	parser = argparse.ArgumentParser(
		description='Testing suite for language development')
	parser.add_argument('directory', metavar='directory',
						help='Central testing directory')
	parser.add_argument('--output-bin', action='store_true',
						help='Write output as binary string')
	parser.add_argument('--leak-check', action='store_true',
						help='Check for leaks')
	args = parser.parse_args()
	
	validate_directory(args.directory)
	options = load_manifest(os.path.join(args.directory, 'manifest.toml'))
	options['output_bin'] = args.output_bin
	options['leak_check'] = args.leak_check

	any_failures = False
	for directory in filter(lambda p: os.path.isdir(p),
							map(lambda p: os.path.join(args.directory, p),
								os.listdir(args.directory))):
		if test_from_directory(directory, options):
			any_failures = True
	if any_failures:
		fancy_print('Some tests failed!', red, bold)
	else:
		fancy_print('All tests passed!', green, bold)

if __name__ == '__main__':
	main()
