import sys

memory_block_start = """
    mov UP, LEFT

    mov UP, ACC
    mov UP, LEFT

    jnz is_write
    mov LEFT, DOWN
    mov 0, DOWN
    mov UP, DOWN
    mov LEFT, NIL
    jmp wrap

is_write:
    mov LEFT, DOWN
    mov 1, DOWN
    mov UP, ACC
    xor LEFT
    mov ACC, DOWN
    swp

wrap:
    mov DOWN, ACC
    xor LEFT
    mov LEFT, RIGHT
    mov ${{addr}}, RIGHT
    mov ${OP_ACK}, RIGHT
    mov ACC, RIGHT
    mov ${OP_NIL}, RIGHT
"""

memory_block_start_left = """
    mov RIGHT, ACC
    swp
    mov RIGHT, ACC
    mov ACC, RIGHT
    mov ACC, RIGHT
    mov ACC, RIGHT
    swp
    mov ACC, RIGHT
"""

memory_block_start_right = """
    mov LEFT, UP
"""

code_memory_split = """
    mov UP, ACC
    mov ACC, LEFT
    mov ACC, RIGHT

    mov UP, ACC
    mov ACC, LEFT
    mov ACC, RIGHT

    mov UP, ACC
    mov ACC, LEFT
    mov ACC, RIGHT

    mov ANY, ACC
    add ANY
    mov ACC, UP
"""

code_memory_split_rev = """
    mov UP, ACC
    mov ACC, RIGHT
    mov ACC, LEFT

    mov UP, ACC
    mov ACC, RIGHT
    mov ACC, LEFT

    mov UP, ACC
    mov ACC, RIGHT
    mov ACC, LEFT

    mov ANY, ACC
    add ANY
    mov ACC, UP
"""

code_memory_cell = """
    mov 0x{val:02x}, ACC
    sav
start:
    mov UP, ACC
    xor {idx}
    jnz skip
    mov UP, ACC
    jnz WRITE
    mov UP, NIL
    swp
    mov ACC, UP
    sav
    jmp start
write:
    mov UP, ACC
    swp
    mov ACC, UP
    jmp start
skip:
    mov UP, NIL
    mov UP, NIL
    mov 0, UP
    jmp start
"""

code_memory_right = """
mov LEFT, RIGHT
mov LEFT, RIGHT
mov LEFT, RIGHT
mov RIGHT, LEFT
"""

code_memory_right_corner = """
mov LEFT, DOWN
mov LEFT, DOWN
mov LEFT, DOWN
mov DOWN, LEFT
"""

code_memory_left = """
mov RIGHT, LEFT
mov RIGHT, LEFT
mov RIGHT, LEFT
mov LEFT, RIGHT
"""

code_memory_left_corner = """
mov RIGHT, DOWN
mov RIGHT, DOWN
mov RIGHT, DOWN
mov DOWN, RIGHT
"""

class BlockMap(dict):
    def __setattr__(self, key, value):
        if key in self.__dict__:
            raise ValueError("Already exists")
        super().__setattr__(key, value)

def gen_memory(blocks, start, idx_offset, idx, depth, contents):
    if depth <= 0:
        addr = idx_offset + idx
        blocks[start] = code_memory_cell.format(idx=addr,val=addr^contents[idx])
        return
    sx, sy = start
    blocks[start] = code_memory_split_rev if depth % 2 else code_memory_split
    sidelen = int(2**(depth-1))
    for i in range(1, sidelen):
        blocks[(sx + i, sy)] = code_memory_right
        blocks[(sx - i, sy)] = code_memory_left
    blocks[(sx + sidelen, sy)] = code_memory_right_corner
    blocks[(sx - sidelen, sy)] = code_memory_left_corner
    gen_memory(blocks, (start[0] + sidelen, start[1] + 1),
        idx_offset, idx + 2**(depth-1), depth - 1, contents)
    gen_memory(blocks, (start[0] - sidelen, start[1] + 1),
        idx_offset, idx, depth - 1, contents)

def write_blocks(blocks, filepath, addr_tmpl):
    with open(filepath, "w+") as f:
        for i,(x,y) in enumerate(blocks):
            f.write(f"tpu X{x} Y{y}\n")
            if i == 0:
                f.write("\t## ${" + addr_tmpl + "}\n")
            for l in blocks[(x,y)].strip("\n").split("\n"):
                f.write(f"{l}\n")
            f.write(f"end\n")
            f.write("\n")

if __name__ == "__main__":
    blocks = {}
    blocks[(0, 0)] = memory_block_start.replace("{addr}", sys.argv[5])
    blocks[(-1, 0)] = memory_block_start_left
    blocks[(1, 0)] = memory_block_start_right
    gen_memory(blocks, (0, 1), int(sys.argv[2]), 0, int(sys.argv[3]),
        open(sys.argv[4], "rb").read().ljust(256, b"\x00"))
    write_blocks(blocks, sys.argv[1], sys.argv[5])