import sys

import re
import argparse

from Emulator import Emulator, CodeParseException, REGISTERS
from TikzPainter import TikzPainter
import Junk

def parse_args ():
	parser = argparse.ArgumentParser(description="For fun x86-64 emulator and stack visualizer.")
	parser.add_argument('-i', '--input-file', default='-', help = '.S file to use as input', dest = 'filename')
	for register in REGISTERS:
		parser.add_argument('--'+register[1:], nargs = 1, type = int, default = [Junk.Junk()], dest = register)
	args = vars(parser.parse_args())

	# Determine args #
	registers_init = {}
	for register in REGISTERS:
		registers_init[register[1:]] = args[register][0]
	registers_init['rip'] = 0
	registers_init['rsp'] = 0
	registers_init['rbp'] = Junk.Junk('old bp')

	program = ""
	if args['filename'] == '-':
		program = sys.stdin.read()
	else:
		with open(args['filename']) as file:
			program = file.read()

	return (program, registers_init)

def main ():

	(program, registers_init) = parse_args()

	# Determine registers to display
	registers_to_draw = ["%rip","%rbp","%rsp", ""]
	for match in re.findall(r"%r[a-z0-9]{1,2}", program):
		if match not in registers_to_draw:
			registers_to_draw.append(match)

	# Setup emulator and drawer
	emu = Emulator(program)
	emu.setRegs(**registers_init)
	emu.setStack(Junk.Junk("junk..."), Junk.Junk("calling eip"))

	painter = TikzPainter(registers_to_draw)

	# Iteratively draw states.
	painter.drawState(emu)
	for line_nr in emu:
		painter.drawState(emu, line_nr)
	painter.drawNames(emu)

	# Display result
	print(painter.to_string())

if __name__ == "__main__":
	main()