#!/usr/bin/python2

import sys

import re
import argparse

from Emulator import Emulator, CodeParseException, REGISTERS
from TikzPainter import TikzPainter
from AsciiPainter import AsciiPainter
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 = 'input-file')
	parser.add_argument('-o', '--output-file', default='-', help = '.S file to use as output', dest = 'output-file')
	for register in REGISTERS:
		parser.add_argument('--'+register[1:], nargs = 1, type = int, default = None, dest = register, help = 'the initial value of the register')
	args = vars(parser.parse_args())

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

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

	# Output file:
	output_file = None
	if args['output-file'] == '-':
		output_file = sys.stdout
	else:
		output_file = open(args['output-file'], 'w')

	##

	return (program, registers_init, output_file)

def main ():

	(program, registers_init, output_file) = 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_class = TikzPainter # AsciiPainter

	painter = painter_class(registers_to_draw)

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

	# Display result
	output_file.write(painter.to_string(emu))
	output_file.close()

if __name__ == "__main__":
	main()