class Things:
    opcodes = {
        'nop': 0x00,
        'hlt': 0x01,
        'ld': 0x04,
        'push': 0x05,
        'pop': 0x06,
        'jmp': 0x07,
        'jz': 0x08,
        'cmp': 0x09,
        'inc': 0x0A,
        'dec': 0x0B,
        'addi': 0x0C,
        'subi': 0x0D,
        'add': 0x0E,
        'sub': 0x0F,
        'muli': 0x10,
        'divi': 0x11,
        'mul': 0x12,
        'div': 0x13,
        'test': 0x16,
        'prt': 0x17,

        'je': 0x08,
    }

    registers = {
        'ip': 0x00,
        'sp': 0x01,
        'flags': 0x03,
        'a': 0x04,
        'b': 0x05,
        'c': 0x06,
        'd': 0x07,
        'r8': 0x08,
        'r9': 0x09,
        'r10': 0x0A,
        'r11': 0x0B,
        'r12': 0x0C,
        'r13': 0x0D,
        'r14': 0x0E,
        'r15': 0x0F,
        'r16': 0x10,
    }

class ZASM:
    def __init__(self):
        self._out = []
        self._pc = 0
        self._org = 0
        self._labels = {}
        self._to_be_resolved = {}

    def emit8(self, value):
        self._out.append(value & 0xFF)
        self._pc += 1

    def emit16(self, value):
        self._out.append((value & 0xFF00) >> 8)
        self._out.append(value & 0xFF)
        self._pc += 2

    def set8(self, addr, value):
        self._out[addr - self._org] = bytes([value & 0xFF])

    def set16(self, addr, value):
        print(self._out)
        self._out[addr - self._org] = (value & 0xFF00) >> 8
        self._out[addr - self._org + 1] = value & 0xFF

    def assemble(self, path):
        with open(path, 'r') as f:
            lines = f.readlines()

        line_no = 0
        for line in lines:
            # Strip away bullshit
            line = line.replace(',', '')
            line = line.replace('\n', '')

            # Increment line number
            line_no += 1

            # Ignore blank lines and comments
            if line == '':
                continue
            if line.startswith('#'):
                continue

            # Is it a label?
            if not line.startswith('    '):
                if line.startswith('org'):
                    self._pc = int(line.split(' ')[1], 16)
                    self._org = self._pc
                else:
                    self._labels[line.split(':')[0]] = self._pc
            # It is code
            else:
                line = line[4:]
                chunks = line.split(' ')
                if chunks[0] in Things.opcodes:
                    self.emit8(Things.opcodes[chunks[0]])
                else:
                    print(f'Unknown opcode: {chunks[0]} at {line_no}: {line}')
                    break

                # SUUUUPER dirty hack to do indirect addressing
                #if chunks[1].startswith('[') and chunks[1].endswith(']'):
                #    self.set8(self._pc - 1, Things.opcodes[chunks[0]] - 2)
                #    chunk[1] = chunk[1][1:-1]
                if len(chunks) == 1:
                    self.emit16(0)
                elif chunks[1].startswith('#'):
                    self.emit16(int(chunks[1][1:]))
                elif chunks[1] in Things.registers:
                    self.emit16(Things.registers[chunks[1]])
                else:
                    self._to_be_resolved[self._pc] = (chunks[1], line_no, line)
                    self.emit16(0xDADA)
                    # print(f'Unknown register or invalid immediate: {chunks[1]!r}')
                    # break

                # SUUUUPER dirty hack to do indirect addressing
                #if chunks[2].startswith('[') and chunks[2].endswith(']'):
                #    self.set8(self._pc - 2, Things.opcodes[chunks[0]] - 2)
                #    chunk[2] = chunk[2][1:-1]
                if len(chunks) <= 2:
                    self.emit8(0)
                elif chunks[2] in Things.registers:
                    self.emit8(Things.registers[chunks[2]])
                elif chunks[2].startswith('#'):
                    self.emit8(int(chunks[2][1:]))
                else:
                    print(f'Invalid immediate: {chunks[2]} at {line_no}: {line}')
                    break

        for loc, (label, line_no, line) in self._to_be_resolved.items():
            print(loc, label, line_no, line)
            print(self._labels)
            if label in self._labels.keys():
                self.set16(loc, self._labels[label])
            else:
                print(f'Failed to resolve {label} at {line_no}: {line}')
                break

    def write(self, path):
        with open(path, 'wb') as f:
            f.write(bytes(self._out))
        self._out = b''

if __name__ == '__main__':
    asm = ZASM()

    asm.assemble('test.zasm')
    asm.write('test.bin')