from generators.generator import Generator
from ground_ast import *
from error import traceback

class X86_64Generator(Generator):
	def __init__(self, ast, code, output_path):
		super().__init__(ast, code, output_path)
		self.stack_size = 0

	def init(self):
		self.lines.append("global _start\n\n")
		self.lines.append("_start:\n\t")

		# generate code
		self.generate()
		self.write()

	def push(self, reg: str):
		self.lines.append("push " + reg + "\n\t")
		self.stack_size += 1

	def pop(self, reg: str):
		self.lines.append("pop " + reg + "\n\t")
		self.stack_size -= 1

	def get_variable(self, var_name: str, reg: str):
		try:
			self.push(
				f"QWORD [rsp + {(self.stack_size - self.variables.get(var_name, None)['stack_loc'] - 1) * 8}]"
			)
		except TypeError: # variable doesnt exist
			traceback(self.code, "NameError", f"\"{var_name}\" is not defined.")
		self.pop(reg)

	def get_var_pos(self, var_name: str):
		return (self.stack_size - self.variables.get(var_name)['stack_loc'] - 1) * 8

	def create_variable(self, var_name: str, starting_value):
		new_var = {"stack_loc": (self.stack_size), "type": type(starting_value)}
		if type(starting_value) == NumberNode:
			self.lines.append(f"mov rax, {starting_value.value}\n\t")
			self.push("rax")
		elif type(starting_value) == VarRefNode:
			self.get_variable(starting_value.var_name, "rax")
			#self.lines.append(f"mov QWORD [rsp + {var_pos}], rax\n\t")
			self.push("rax")
		self.variables[var_name] = new_var

	def change_variable(self, var_name: str, new_value):
		var_pos = self.get_var_pos(var_name)

		if type(new_value) == NumberNode: # we're changing a variable to a number
			#self.lines.append(f"mov rax, {new_value.value}\n\t")
			self.lines.append(f"mov QWORD [rsp + {var_pos}], {new_value.value}\n\t")

		elif type(new_value) == VarRefNode: # we're changing a variable to the value of another variable
			self.get_variable(new_value.var_name, "rax")
			self.lines.append(f"mov QWORD [rsp + {var_pos}], rax\n\t")

	def generate_InstructionNode(self, node: InstructionNode):
		if node.instruction == "end":
			if len(node.arguments) == 0: # example: "end"
				traceback(self.code, "TypeError", "end expects atleast 1 argument.")
			elif len(node.arguments) > 1: # example: "end 1 1"
				traceback(self.code, "TypeError", "end expects only 1 argument.")
			if not type(node.arguments[0]) in [NumberNode, VarRefNode]: # example: "end true"
				traceback(self.code, "TypeError", f"end expects an integer, not {type(node.arguments[0])}")

			self.lines.append("mov rax, 60\n\t")
			if isinstance(node.arguments[0], NumberNode):
				self.lines.append("mov rdi, " + str(node.arguments[0].value) + "\n\t")
			elif isinstance(node.arguments[0], VarRefNode):
				self.get_variable(node.arguments[0].var_name, "rdi")
				#self.lines.append("mov rdi, " + str(self.get_variable(node.arguments[0].var_name)) + "\n\t")
			self.lines.append("syscall\n\t")

		elif node.instruction == "set":
			if len(node.arguments) < 2: # example: "set" or "set &hi"
				traceback(self.code, "TypeError", "set expects atleast 2 arguments.")
			elif len(node.arguments) > 2: # example: "set &hi 0 123"
				traceback(self.code, "TypeError", "set expects only 2 arguments.")
			if not isinstance(node.arguments[0], VarPointerNode):
				traceback(self.code, "TypeError", f"the first argument of set should be a variable pointer, not \"{type(node.arguments[0])}\"")
			if type(node.arguments[1]) not in [NumberNode, VarRefNode]:
				traceback(self.code, "TypeError", f"variables can't be of type \"{type(node.arguments[1])}\"")

			variable_exists = self.variables.get(node.arguments[0].var_name, None) != None

			if not variable_exists: # create a new variable
				self.create_variable(node.arguments[0].var_name, node.arguments[1])
			else: # modify the existing one
				self.change_variable(node.arguments[0].var_name, node.arguments[1])

				
		else:
			self.lines.append("; FUCK\n\t")
			#raise NotImplementedError(f"A generate method hasn't been made for the \"{node.instruction}\" instruction.")