from .utils import *

# common instructions all push commands end with
PUSH_COMMON_ASM = """@SP
A=M
M=D
@SP
M=M+1
"""

# push <segment> <index>
# when <segment> is one of local, argument, this and that
PUSH_ASM = (
    """@{segment}
D=M
@{index}
A=D+A
D=M
"""
    + PUSH_COMMON_ASM
)

# when <segment> is one of static, temp and pointer
PUSH_FIXED_ASM = (
    """@{addr}
D=M
"""
    + PUSH_COMMON_ASM
)

# push constant <constant>
PUSH_CONSTANT_ASM = (
    """@{constant}
D=A
"""
    + PUSH_COMMON_ASM
)


# pop <segment> <index>
# when <segment> is one of local, argument, this and that
POP_ASM = """@{index}
D=A
@{segment}
M=M+D
@SP
AM=M-1
D=M
@{segment}
A=M
M=D
@{index}
D=A
@{segment}
M=M-D
"""

# when <segment> is one of static, temp and pointer
POP_FIXED_ASM = """@SP
AM=M-1
D=M
@{addr}
M=D
"""

SEGMENT_PTR = {
    "local": "LCL",
    "argument": "ARG",
    "this": "THIS",
    "that": "THAT",
}

SEGMENT_SIZE = {
    "pointer": 2,  # 3..4, i.e. THIS and THAT
    "temp": 8,  # 5..12
    "static": 240,  # 16..255
}


def translate_memory(action, segment, index, prog, verbose=False):
    try:
        index = int(index)
    except ValueError:
        exit_on_error(
            f"Syntax error: invalid segment index {index}", EXIT_CODE_SYNTAX_ERROR
        )

    if index < 0:
        exit_on_error(
            f"Address error: negative segment index {index}", EXIT_CODE_ADDR_ERROR
        )

    asm = f"// {action} {segment} {index}\n" if verbose else ""
    if segment in ["local", "argument", "this", "that"]:
        # these segments are dynamic
        asm += PUSH_ASM if action == "push" else POP_ASM
        asm = asm.format(segment=SEGMENT_PTR[segment], index=index)
    elif segment in ["static", "temp", "pointer"]:
        # these segments have a fixed location and size on memory
        if index >= SEGMENT_SIZE[segment]:
            exit_on_error(
                f"Address error: segment index {index} exceeds size of {segment}",
                EXIT_CODE_ADDR_ERROR,
            )

        addr = ""
        if segment == "static":
            # translate "static i" in "prog.vm" into "@prog.i"
            addr = f"{prog}.{index}"
        elif segment == "temp":
            addr = str(5 + index)
        else:
            # "pointer 0" is THIS, "pointer 1" is THAT
            addr = "THIS" if index == 0 else "THAT"

        asm += PUSH_FIXED_ASM if action == "push" else POP_FIXED_ASM
        asm = asm.format(addr=addr)
    elif segment == "constant":
        if action == "pop":
            exit_on_error(
                f"Syntax error: popping to constant not allowed", EXIT_CODE_SYNTAX_ERROR
            )
        asm += PUSH_CONSTANT_ASM.format(constant=index)
    else:
        exit_on_error(
            f"Syntax error: invalid memory segment {segment}", EXIT_CODE_SYNTAX_ERROR
        )

    return asm