from ground_ast import RootNode
from typing import Any
from ground_ast import *


class SymbolTable:
	def __init__(self, parent=None):
		self.table = {}  # variable name -> info (e.g., stack offset, type)
		self.parent = parent

	def define(self, name, value):
		self.table[name] = value

	def lookup(self, name):
		scope = self
		while scope:
			if name in scope.table:
				return scope.table[name]
			scope = scope.parent
		return None
		#raise KeyError(f"Undefined variable: {name}")


class Generator:
	def __init__(self, ast: RootNode, code: str, output_path: str):
		self.ast = ast
		self.lines: list[str] = []
		self.code = code
		self.output_path = output_path
		self.global_scope = SymbolTable()
		self.current_var_scope = self.global_scope
		self.constants = {}
		self.structs = {}
		self.functions: dict[str, dict] = {}
		self.labels = []
		self.arg_list = []
		self.constants_reverse = {}
		self.constant_counter = 0

	def ground_type_to_node(self, ground_type: str):
		if ground_type == "string":
			return StringNode
		elif ground_type == "int":
			return IntNode
		elif ground_type == "float":
			return FloatNode
		elif ground_type == "bool":
			return BoolNode
		else:
			return ground_type

	def add_constant(self, value: Any, no_string: bool = False):
		existing_constant_name = self.constants_reverse.get(value, None)
		if existing_constant_name != None: return f"[.{existing_constant_name}]"
		self.constants["LC" + str(self.constant_counter)] = {"value": value, "no_string": no_string}
		self.constants_reverse[value] = "LC" + str(self.constant_counter)
		self.constant_counter += 1
		return "[.LC" + str(self.constant_counter-1) + "]"
	
	def add_function(self, node: FunctionNode):
		self.functions[node.name] = {
			"func": node,
			"scope": SymbolTable(self.current_var_scope)
		}
		return self.functions[node.name]
	
	def clamp_instruction_args(self, instruction: InstructionNode, min_args: int, max_args: int):
		if len(instruction.arguments) < min_args:
			traceback(self.code, "TypeError", f"{instruction.instruction} expects at least {min_args} arguments.")
		elif len(instruction.arguments) > max_args:
			traceback(self.code, "TypeError", f"{instruction.instruction} expects at most {max_args} arguments.")

	def init(self):
		pass

	def generate_node(self, node):
		#self.lines.append(f"; {node}\n\t")
		node_type = str(type(node))[19:-2]
		if not hasattr(self, f"generate_{node_type}"):
			raise NotImplementedError(f"Generator has no generate method for {node_type}.")
		getattr(self, f"generate_{node_type}")(node, self.lines)

	def generate(self):
		for statement in self.ast.statements:
			self.generate_node(statement)

	def write(self):
		with open(self.output_path + ".asm", "w") as f:
			f.write(self.lines)