#!/usr/bin/env python3

"""
Generates the ARM64 Assembly for an executable that performs the given 
calculation and exits with that number as the exit code. It currently
only allows for multiplication and addition. You can check the output
modulo 256 using `echo $?`. The only defined variables are the letters
a-z, where they represent their 1-based index into the alphabet.
"""

import string
import sys, os; sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) # just becuase this is in a subdirectory
from ppp_lexer import *
from ppp_parser import *

_program, file_path = sys.argv

lexer = Lexer.from_file(file_path)
expr = parse_expression(lexer)
lexer.assert_token(EofToken())

def get_index(char: str) -> int:
	"""
	Assuming char is [a-z]
	"""
	return ord(char)-96

def get_stack_address(index: int) -> int:
	return 112-index*4

def to_asm(expr: Expression, *, taken: int=0):
	"""
	Generates the assembly using registers w0, w1, w2, ... as a stack.

	`taken` is how many registers are already taken by previous elements 'on the stack'
	"""
	if isinstance(expr, Variable):
		assert expr.name in string.ascii_lowercase
		address = get_stack_address(get_index(expr.name))
		print(f"\tldr w{taken}, [sp, {address}] // {expr.name}")
	else:
		match expr:
			case Addition():       word = 'add'
			case Multiplication(): word = 'mul'
			case _: assert False, f"Unimplemented: {expr}"
		to_asm(expr.lhs, taken=taken)
		to_asm(expr.rhs, taken=taken+1)

		print(f"\t{word} w{taken}, w{taken}, w{taken+1} // {expr.represent()}")

print(".global main")
print(".align 2")
print()
print("main:")
print("\tsub sp, sp, #112") # to fit 26*4=104 bytes. It seems that the sp should only be offset by 16 bytes at a time.
for i in range(1,27):
	print(f"\tmov w0, {i}")
	print(f"\tstr w0, [sp, {get_stack_address(i)}]")
to_asm(expr)
print("\tadd sp, sp, 112")
print("\tret")