#!/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

PAINTER_ID_TO_CLASS = {
    'ascii': AsciiPainter,
    'tikz': TikzPainter
}

def parse_args ():
	parser = argparse.ArgumentParser(description="For fun x86-64 emulator and stack visualizer.")
	parser.add_argument('-i', '--input-file',  default='-',     help = 'File to use as input',  dest = 'input-file')
	parser.add_argument('-o', '--output-file', default='-',     help = 'File to use as output', dest = 'output-file')
	parser.add_argument('-p', '--painter',     default='ascii', help = 'Drawer to use',         dest = 'painter-name', choices=PAINTER_ID_TO_CLASS.keys())
	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')


        # Determine painter class
        painter_class = PAINTER_ID_TO_CLASS[args['painter-name']]

	## Return
	return (program, registers_init, output_file, painter_class)

def main ():

	(program, registers_init, output_file, painter_class) = 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 = 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()