CALL_ASM = """@{return_address}
D=A
@SP
A=M
M=D
@LCL
D=M
@SP
AM=M+1
M=D
@ARG
D=M
@SP
AM=M+1
M=D
@THIS
D=M
@SP
AM=M+1
M=D
@THAT
D=M
@SP
AM=M+1
M=D
@SP
MD=M+1
@LCL
M=D
@{arg_lcl_offset}
D=D-A
@ARG
M=D
@{function_label}
0;JMP
({return_address})
"""

FUNCTION_ASM = "({function_label})\n"

# repeated (n-1) times for function with n arguments
INIT_ONE_VAR_ASM = """@SP
AM=M+1
M=0
"""

# concatenated only when function has arguments
INIT_VARS_ASM = """@SP
A=M
M=0
{init_vars}@SP
M=M+1
"""

RETURN_ASM = """@LCL
D=M
@5
A=D-A
D=M
@{tmp}
M=D
@SP
A=M-1
D=M
@ARG
A=M
M=D
@ARG
D=M+1
@SP
M=D
@LCL
AM=M-1
D=M
@THAT
M=D
@LCL
AM=M-1
D=M
@THIS
M=D
@LCL
AM=M-1
D=M
@ARG
M=D
@LCL
AM=M-1
D=M
@LCL
M=D
@{tmp}
A=M
0;JMP
"""

ret_idx = {}


def translate_function(action, function, n, prog, verbose=False):
    global ret_idx
    try:
        n = int(n)
    except ValueError:
        exit_on_error(f"Syntax error: not an integer: {n}", EXIT_CODE_SYNTAX_ERROR)

    asm = f"// {action} {function} {n}\n" if verbose else ""
    if action == "call":
        if prog not in ret_idx:
            ret_idx[prog] = 0
        else:
            ret_idx[prog] += 1
        asm += CALL_ASM.format(
            return_address=f"{prog}$ret.{ret_idx[prog]}",
            arg_lcl_offset=(5 + n),
            function_label=f"{function}",
        )
    else:
        asm += FUNCTION_ASM.format(function_label=f"{function}")
        if n > 0:
            # function has arguments
            asm += INIT_VARS_ASM.format(init_vars=INIT_ONE_VAR_ASM * (n - 1))

    return asm


def translate_return(verbose=False):
    asm = "// return\n" if verbose else ""
    asm += RETURN_ASM.format(tmp="R13")
    return asm