import re
import sys

REGISTERS=["%rip","%rbp","%rsp","","%rax","%rdi","%r8","%r9"]
REG_STATUS_TO_COLOR = {
	"insert": "green",
	"change": "yellow",
	"remove": "red",
	"none":   "black",
	"jump":   "orange"
}

class TikzPainter:

	def __init__ (self, registers=None, max_stack_size=8):
		self.registers = registers if registers else REGISTERS
		self.text = []
		self.pos_x = 0
		self.max_stack_size = max_stack_size

	def addText (self, str, *args):
		self.text.append(str.format(*args))

	def getRegColor (self, reg_state):
		return REG_STATUS_TO_COLOR[reg_state]

	def drawState (self, emu, line_nr=None):
		x = (self.pos_x+1)*2.5-0.5
		self.pos_x += 1

		# Draw register cells
		pos = 0.5
		for reg_name in self.registers:
			pos -= 0.5
			if reg_name =="":
				continue
			reg_state = emu.regState(reg_name)
			reg_dect = "register_node, fill="+self.getRegColor(reg_state)+"!10"
			self.addText("\t\\node[{}]() at ({}, {}){{{}}};\n",
					reg_dect, x, pos, emu.getVal(reg_name))

		# Draw stack
		pos -= 1.5
		base_sp = max(emu.last_sp, emu.getVal('%rsp'))
		if emu.last_sp - 1 == emu.getVal('%rsp'):
			pos += 0.5
		for index in range(base_sp, max(-1, base_sp-self.max_stack_size),-1):
			reg_state = emu.regState("m"+str(index))
			reg_dect = "register_node, fill="+self.getRegColor(reg_state)+"!10"
			self.addText("\t\\node[{}]() at ({}, {}){{{}}};\n",
					reg_dect, x, pos, emu.stack[index])
			pos -= 0.5

		# Draw line signature
		if line_nr == None:
			return
		pos = 2
		signature = emu.getLineSignature(line_nr)
		for token in signature:
			self.addText("\t\\node[text_node]() at ({}, {}){{{}}};\n", x-1.25, pos, token)
			pos -= 0.5

	def drawNames (self, emu):
		x = 0
		pos = 0
		for reg_name in self.registers:
			self.addText("\t\\node[text_node]() at ({}, {}){{{}}};\n", x, pos, reg_name[1:])
			pos -= 0.5

		self.addText("\t\\node[text_node]() at ({}, {}){{TOS+1}};\n", x, pos-0.5)
		self.addText("\t\\node[text_node]() at ({}, {}){{TOS}};\n", x, pos-1.0)
		pos -= 1
		for index in range(1,self.max_stack_size):
			pos -= 0.5
			self.addText("\t\\node[text_node]() at ({}, {}){{TOS-{}}};\n", x, pos, index)

	def __str__ (self):
		return "".join(self.text)


class Emulator:

	def __init__ (self, source_text):
		self.registers = {}
		for reg_name in REGISTERS:
			self.registers[reg_name] = "Junk"
		self.stack = {}
		self.code = []
		self.labels = {}
		self.changes = {}
		self.status = {}
		self.last_sp = 0
		self.max_stack_size = 0
		index = 0
		for line in iter(source_text.splitlines()):
			index = self.processSourceLine(line, index)

	def setStack (self, *stack_list, **kwargs):
		"""
		Sets various stack elements, starting from 0 and going up.
		Automatically sets rsp. This can be disabled by passing set_rsp=False.
		"""
		i = -1
		for element in stack_list:
			i += 1
			self.stack[i] = element
		if (not 'set_rsp' in kwargs) or kwargs['set_rsp']:
			self.setRegs(rsp=i)

	def setRegs (self, **reg_dict):
		for reg_name, reg_val in reg_dict.iteritems():
			self.registers["%"+reg_name] = reg_val

	def getVal (self, val_text):
		if val_text[0] == "$":
			return int(val_text[1:])
		elif val_text[0] == "%":
			return self.registers[val_text]
		else:
			raise "NOT A VALUE! UGHGUIHGUE!"

	def compareVal (self, valT1, valT2):
		val1 = self.getVal(valT2)
		val2 = self.getVal(valT1)
		self.status = {"g":val1>val2}

	def changedRegisters (self, *args):
		for reg, val in args:
			self.changes[reg] = val

	def processSourceLine(self, line_text, index):
		tokens = re.findall(r'[\w%:$#]+', line_text)
		if len(tokens) <= 0:
			return index
		if tokens[0][-1] == ':':
			self.labels[tokens[0][:-1]] = index
			tokens = tokens[1:]
		for i in range(len(tokens)):
			if tokens[i][0] == "#":
				tokens = tokens[:i]
				break
		self.code.append(tokens)
		return index + 1

	def regState (self, reg_name):
		if reg_name in self.changes:
			return self.changes[reg_name]
		else:
			return "none"

	def getLineSignature (self, line_nr):
		return [ "\\"+token if token[0]=="$" or token[0]=="%" else token for token in self.code[line_nr] ]

	def pushToStack (self, new_element):
		self.registers['%rsp'] = self.registers['%rsp'] + 1
		self.stack[self.registers['%rsp']] = new_element
		self.changedRegisters(('%rsp',"change"), ("m"+str(self.registers['%rsp']),"insert"))
		self.max_stack_size = max(self.registers['%rsp'], self.max_stack_size)

	def popFromStack (self):
		temp = self.stack[self.registers['%rsp']]
		self.registers['%rsp'] -= 1
		self.changedRegisters(('%rsp',"change"),("m"+str(self.registers['%rsp']+1),"remove"))
		return temp

	def iterate (self):
		old_rip = self.registers['%rip']
		self.last_sp = self.registers['%rsp']
		if self.registers['%rip'] >= len(self.code):
			return None
		instruct = self.code[self.registers['%rip']]
		opcode = instruct[0]
		self.changes = {}
		if   opcode[:4] == "push":
			self.pushToStack(self.getVal(instruct[1]))
		elif opcode[:3] == "pop":
			self.registers[instruct[1]] = self.popFromStack()
			self.changedRegisters((instruct[1],"change"))

		elif opcode[:3] == "mov":
			self.changedRegisters((instruct[2],"change"))
			self.registers[instruct[2]] = self.getVal(instruct[1])
		elif opcode[:3] == "add":
			self.changedRegisters((instruct[2],"change"))
			self.registers[instruct[2]] = self.getVal(instruct[2]) + self.getVal(instruct[1])
		elif opcode[:3] == "sub":
			self.changedRegisters((instruct[2],"change"))
			self.registers[instruct[2]] = self.getVal(instruct[2]) - self.getVal(instruct[1])

		elif opcode[:3] == "cmp":
			self.compareVal(instruct[1],instruct[2])
		elif opcode == "jg":
			if self.status["g"]:
				self.registers['%rip'] = self.labels[instruct[1]]
		elif opcode == "jmp":
			self.registers['%rip'] = self.labels[instruct[1]]

		elif opcode == "call":
			self.pushToStack(self.registers['%rip']+1)
			self.registers['%rip'] = self.labels[instruct[1]]
		elif opcode == "leave":
			self.registers["%rsp"] = self.registers["%rbp"]
			self.registers["%rbp"] = self.popFromStack()
		elif opcode == "ret":
			self.registers['%rip'] = self.popFromStack()

		if self.registers['%rip'] == old_rip:
			self.registers['%rip'] += 1
		else:
			self.changedRegisters(('%rip',"jump"))

		return old_rip

	def __iter__(self):
		return self

	def next(self):
		output = self.iterate()
		if output is None:
			raise StopIteration()
		else:
			return output