#!/usr/bin/env python3

"""
Disassembler for nandgame.

Using my own flavor of assembly language.
I don't like the "C-style" one nandgame introduces.
"""

import sys

ZERO = "#0"
DEST_NONE = "_"
JUMP_NONE = ""
JUMPS_IF_NZERO = ["jgt", "jlt", "jne", "jmp"]
JUMPS_IF_ZERO = ["jge", "jle", "jeq", "jmp"]

def decode_jump(ins: int) -> str:
    if (ins & 0x7) == 0:
        return JUMP_NONE

    if (ins & 0x7) == 0x7:
        return "jmp"

    jl = (ins & (1 << 2)) != 0
    je = (ins & (1 << 1)) != 0
    jg = (ins & (1 << 0)) != 0

    # implied: and not jg
    if jl and je:
        return "jle"
    # implied: and not je
    if jl and jg:
        return "jne"
    # implied: and not je
    if je and jg:
        return "jge"

    # implied: only one flag is 1
    if jl:
        return "jlt"
    if je:
        return "jeq"
    if jg:
        return "jgt"

    return "<unknown>"


# return op, and whether it's a one-op or two-op
def decode_ins(ins: int) -> (str, bool):
    opcode = (ins >> 8) & 0x03
    ar_n_log = (ins & (1 << 10)) != 0
    opcode |= ar_n_log << 2

    if opcode == 0b000:
        return "and", True
    if opcode == 0b001:
        return "or", True
    if opcode == 0b010:
        return "xor", True
    if opcode == 0b011:
        return "not", False
    if opcode == 0b100:
        return "add", True
    if opcode == 0b101:
        return "inc", False
    if opcode == 0b110:
        return "sub", True
    if opcode == 0b111:
        return "dec", False

    return "<?>"


# normally, X = arg1 = D
def decode_arg1(ins: int) -> str:
    use_mem = (ins & (1 << 12)) != 0
    zx = (ins & (1 << 7)) != 0
    sw = (ins & (1 << 6)) != 0

    if zx:
        return ZERO
    if not sw:
        return "D"
    return "M" if use_mem else "A"


# normally, Y = arg2 = A
def decode_arg2(ins: int) -> str:
    use_mem = (ins & (1 << 12)) != 0

    # don't care, only X is zeroed
    # zx = (ins & (1 << 7)) != 0
    sw = (ins & (1 << 6)) != 0

    if sw:
        return "D"
    return "M" if use_mem else "A"


def decode_dest(ins: int) -> str:
    dA = (ins & (1 << 5)) != 0
    dD = (ins & (1 << 4)) != 0
    dM = (ins & (1 << 3)) != 0
    dest = ""
    if dA:
        dest += "A"
    if dD:
        dest += "D"
    if dM:
        dest += "M"

    return dest if dest else DEST_NONE


def decode_instruction(ins: int) -> list[str]:
    """
    Will return a 5 element list/tuple/whatever
    mnemonic, destination, X, Y, jumpdest
    """
    if ins & 0x8000 == 0:
        # mov? ldr? ldi? aaaaaaaaaaa....
        return ["mov", "A", f"#{ins}", "", ""]
    else:
        mnemonic, two_op = decode_ins(ins)
        dest = decode_dest(ins)
        op1 = decode_arg1(ins)
        op2 = decode_arg2(ins) if two_op else ""
        jumpdest = decode_jump(ins)
        return [mnemonic, dest, op1, op2, jumpdest]


def fixup_ins(ins: int) -> list[str]:
    (mnemonic, dest, op1, op2, jumpdest) = decode_instruction(ins)
    # fixups
    if op1 == ZERO:
        # subtract something from #0 - subtraction
        if  mnemonic == "sub":
            return ["neg", dest, op2, "", jumpdest]

        # 0 AND something = 0
        if mnemonic == "and":
            # if no dest, only jump matters
            if dest == DEST_NONE:
                # jump always or jump-if-zero --> always jump
                if jumpdest in JUMPS_IF_ZERO:
                    return ["jmp", "", "", "", ""]
                # all other jumps? <, >, <>, nojmp
                else:
                    return ["nop", "", "", "", ""]
            else:
                if jumpdest in JUMPS_IF_ZERO:
                    newjmp = "jmp"
                else:
                    newjmp = ""
                return ["mov", dest, ZERO, "", newjmp]

        # 0 +/|/^ something = something
        if mnemonic in ["add", "or", "xor"]:
            if dest == DEST_NONE and jumpdest == JUMP_NONE:
                return ["nop", "", "", "", ""]
            else:
                return ["mov", dest, op2, "", jumpdest]

        # basically, not 0 == 0xFFFF....
        # opposite of what AND is doing?
        if mnemonic == "not":
            # if no dest, only jump matters
            if dest == DEST_NONE:
                # 0xFFFF is not jgt, since highest bit is always signed.
                if jumpdest in ["jeq", "jgt", "jge", JUMP_NONE]:
                    return ["nop", "", "", "", ""]
                else:
                    return ["jmp", "", "", "", ""]

    return [mnemonic, dest, op1, op2, jumpdest]


def print_decoded(ins: int, simplify: bool) -> str:
    if simplify:
        (mnemonic, dest, op1, op2, jumpdest) = fixup_ins(ins)
    else:
        (mnemonic, dest, op1, op2, jumpdest) = decode_instruction(ins)

    jumpdest_str = f".{jumpdest}" if jumpdest else ""
    opcode_str = f"{mnemonic}{jumpdest_str}"
    dest_str = f"{dest}, " if dest else 7 * " "
    op1_str = f"{op1}{", " if op2 else ""}"
    return f"{opcode_str:<9}{dest_str:<6}{op1_str:<4}{op2}"


def main():
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} [filename]")
        sys.exit(1)

    try:
        filename = sys.argv[1]
        with open(filename, "rb") as f:
            while True:
                insb = f.read(2)
                if not insb:
                    break
                ins = int.from_bytes(insb)
                raw_ins = f"{insb[0]:02x} {insb[1]:02x}"
                decoded_ins = print_decoded(ins, False)
                decoded_ins2 = print_decoded(ins, True)
                if decoded_ins == decoded_ins2:
                    line = f"\t{raw_ins}\t{decoded_ins}"
                else:
                    line = f"\t{raw_ins}\t{decoded_ins2:<25}; {decoded_ins}"

                print(line)

    except FileNotFoundError:
        print(f"File {filename} not found.")
        sys.exit(1)
    # head, tail...
    except BrokenPipeError:
        sys.exit(0)


if __name__ == "__main__":
    main()