#!/usr/bin/env python3

###
# LLM generated
###

"""
Assembler for nandgame, matching the custom disassembler above.

Syntax (as produced by print_decoded):

    mnemonic[.jump]  DEST, OP1[, OP2]

Examples:

    mov        A, #123
    add.jgt    D, D, A
    sub        _, D, M
    not        D, D
    inc.jeq    D, D
    and        _, D, M
    xor        M, D, M

DEST:
    A, D, M, any combination like AD, AM, DM, ADM, or "_" for no destination.

OP1 / OP2:
    D, A, M, #0 (for OP1 only), or #<number> for mov A,#imm (A-instruction).

Jumps:
    jlt, jle, jeq, jne, jgt, jge, jmp, or none.
"""

import sys

ZERO = "#0"
DEST_NONE = "_"
JUMP_NONE = ""
ENDIANNESS = "little"

# mapping from mnemonic to (opcode, two_op)
MNEMONICS = {
    "and": (0b000, True),
    "or":  (0b001, True),
    "xor": (0b010, True),
    "not": (0b011, False),
    "add": (0b100, True),
    "inc": (0b101, False),
    "sub": (0b110, True),
    "dec": (0b111, False),
}

# jump mnemonic -> bits 0..2
JUMP_ENCODE = {
    "":    0b000,
    "jgt": 0b001,
    "jeq": 0b010,
    "jge": 0b011,
    "jlt": 0b100,
    "jne": 0b101,
    "jle": 0b110,
    "jmp": 0b111,
}


def encode_dest(dest: str) -> int:
    """
    dest is something like "A", "D", "M", "AD", "ADM", or "_" for none.
    Returns bits for A,D,M in positions 5,4,3.
    """
    dest = dest.strip()
    if dest == DEST_NONE:
        return 0

    bits = 0
    if "A" in dest:
        bits |= (1 << 5)
    if "D" in dest:
        bits |= (1 << 4)
    if "M" in dest:
        bits |= (1 << 3)
    return bits


def encode_jump(jump: str) -> int:
    jump = jump.strip()
    if jump not in JUMP_ENCODE:
        raise ValueError(f"Unknown jump condition: {jump}")
    return JUMP_ENCODE[jump]


def encode_args_two_op(op1: str, op2: str) -> int:
    """
    For two-operand instructions, find zx, sw, use_mem bits that reproduce
    the given op1/op2 under decode_arg1/decode_arg2.

    op1 in {D, A, M, #0}
    op2 in {D, A, M}
    """
    op1 = op1.strip()
    op2 = op2.strip()

    # brute-force all combinations of zx, sw, use_mem and pick the one that matches
    for zx in (0, 1):
        for sw in (0, 1):
            for use_mem in (0, 1):
                # simulate decode_arg1/2
                if zx:
                    dec_op1 = ZERO
                else:
                    if not sw:
                        dec_op1 = "D"
                    else:
                        dec_op1 = "M" if use_mem else "A"

                if sw:
                    dec_op2 = "D"
                else:
                    dec_op2 = "M" if use_mem else "A"

                if dec_op1 == op1 and dec_op2 == op2:
                    bits = 0
                    if use_mem:
                        bits |= (1 << 12)
                    if zx:
                        bits |= (1 << 7)
                    if sw:
                        bits |= (1 << 6)
                    return bits

    raise ValueError(f"Unsupported operand combination for two-op: {op1}, {op2}")


def encode_args_one_op(op1: str) -> int:
    """
    For one-operand instructions, only decode_arg1 matters.
    We choose canonical encodings:

        D  -> zx=0, sw=0
        A  -> zx=0, sw=1, use_mem=0
        M  -> zx=0, sw=1, use_mem=1
        #0 -> zx=1
    """
    op1 = op1.strip()
    bits = 0

    if op1 == ZERO:
        bits |= (1 << 7)  # zx
        # sw/use_mem don't matter for arg1 when zx=1, but keep them 0
        return bits

    if op1 == "D":
        # zx=0, sw=0, use_mem=0
        return bits

    if op1 == "A":
        bits |= (1 << 6)  # sw=1
        # use_mem=0
        return bits

    if op1 == "M":
        bits |= (1 << 6)  # sw=1
        bits |= (1 << 12)  # use_mem=1
        return bits

    raise ValueError(f"Unsupported operand for one-op: {op1}")


def encode_instruction(mnemonic: str, dest: str, op1: str, op2: str, jump: str) -> int:
    """
    Encode a single instruction into a 16-bit integer.
    """
    mnemonic = mnemonic.strip()

    if mnemonic == "hlt":
        return (0xFFFF & ~0x4000)

    dest = dest.strip()
    op1 = op1.strip()
    op2 = op2.strip()
    jump = jump.strip()

    # A-instruction: mov A, #imm
    if mnemonic == "mov":
        if dest == "A" and op1.startswith("#") and not op2 and not jump:
            imm_str = op1[1:]
            if imm_str.startswith("0x") or imm_str.startswith("0X"):
                value = int(imm_str, 16)
            else:
                value = int(imm_str, 10)
            if not (0 <= value < 0x8000):
                raise ValueError(f"Immediate out of range (0..32767): {value}")
            return value & 0x7FFF
        else:
            raise ValueError(f"Invalid args to mov.")

    # C-instruction
    if mnemonic not in MNEMONICS:
        raise ValueError(f"Unknown mnemonic: {mnemonic}")

    opcode, two_op = MNEMONICS[mnemonic]

    # bit 14, 15 = 1
    ins = 0xC000

    # opcode bits: low 2 bits in 8..9, high bit in 10 (ar_n_log)
    low2 = opcode & 0b11
    high1 = (opcode >> 2) & 0b1
    ins |= (low2 << 8)
    if high1:
        ins |= (1 << 10)

    # dest bits
    ins |= encode_dest(dest)

    # jump bits
    ins |= encode_jump(jump)

    # arg bits
    if two_op:
        if not op2:
            raise ValueError(f"Two-op instruction {mnemonic} requires two operands")
        ins |= encode_args_two_op(op1, op2)
    else:
        if not op1:
            raise ValueError(f"One-op instruction {mnemonic} requires one operand")
        ins |= encode_args_one_op(op1)

    return ins


def parse_line(line: str):
    """
    Parse a single assembly line into (mnemonic, dest, op1, op2, jump).
    Returns None if the line is empty or comment.
    """
    # strip comments starting with ';' or '#'
    for sep in (";"):
        idx = line.find(sep)
        if idx != -1:
            line = line[:idx]
    line = line.strip()
    if not line:
        return None

    # first token: mnemonic[.jump]
    parts = line.split(None, 1)
    if not parts:
        return None
    opcode_part = parts[0]
    rest = parts[1] if len(parts) > 1 else ""

    if "." in opcode_part:
        mnemonic, jump = opcode_part.split(".", 1)
    else:
        mnemonic, jump = opcode_part, ""

    # operands: dest, op1[, op2]
    dest = ""
    op1 = ""
    op2 = ""

    if rest:
        ops = [o.strip() for o in rest.split(",")]
        ops = [o for o in ops if o]  # remove empty
        if len(ops) >= 1:
            dest = ops[0]
        if len(ops) >= 2:
            op1 = ops[1]
        if len(ops) >= 3:
            op2 = ops[2]
        if len(ops) > 3:
            raise ValueError(f"Too many operands: {rest}")

    # normalize no-dest
    if dest == "":
        dest = DEST_NONE

    return mnemonic, dest, op1, op2, jump


def assemble_file(in_filename: str, out_filename: str):
    with open(in_filename, "r") as fin, open(out_filename, "wb") as fout:
        lineno = 0
        for line in fin:
            lineno += 1
            try:
                parsed = parse_line(line)
                if parsed is None:
                    continue
                mnemonic, dest, op1, op2, jump = parsed
                ins = encode_instruction(mnemonic, dest, op1, op2, jump)
                fout.write(ins.to_bytes(2, byteorder=ENDIANNESS))
            except Exception as e:
                raise SystemExit(f"{in_filename}:{lineno}: {e}")


def main():
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} input.asm output.bin")
        sys.exit(1)

    assemble_file(sys.argv[1], sys.argv[2])


if __name__ == "__main__":
    main()