# 
import os
import re
import pin
import assembly as ASM
from controller import get_hex_str, save_hex_str

# 允许操作的寄存器
ALLOW_REGISTER_MAP = {
    "a": pin.A,
    "b": pin.B,
    "c": pin.C,
    "d": pin.D,
}

OP2_STR_MAP = {
    "mov": ASM.MOV,
    "add": ASM.ADD,
    "sub": ASM.SUB,

}

OP1_STR_MAP = {
    "inc": ASM.INC,
    "dec": ASM.DEC,
}

OP0_STR_MAP = {
    "nop": ASM.NOP,
    "hlt": ASM.HLT,
}

OP2_VALUE_SET = set(OP2_STR_MAP.values())
OP1_VALUE_SET = set(OP1_STR_MAP.values())
OP0_VALUE_SET = set(OP0_STR_MAP.values())


class Code:
    def __init__(self, number, source):
        self.number: int = number  # 行号
        self.source: str = source  # 源代码
        self._op_str = self._src = self._dst = None

        self.prepare_source()

    def __repr__(self):
        return f"[{self.number}] - op: {self._op_str}, dst: {self._dst}, src: {self._src} - '{self.source}'"

    def prepare_source(self):
        tup = self.source.split(",")
        if len(tup) > 2:
            raise Exception("指令格式错误")

        if len(tup) == 2:
            self._src = tup[1].strip()

        tup = re.split(r" +", tup[0].strip())

        if len(tup) > 2:
            raise Exception("指令格式错误")

        if len(tup) == 2:
            self._dst = tup[1].strip()

        self._op_str = tup[0]

    @property
    def op(self):
        # 根据 mov, add 等指令得到对应的操作二进制码
        if self._op_str in OP2_STR_MAP:
            op = OP2_STR_MAP[self._op_str]
        elif self._op_str in OP1_STR_MAP:
            op = OP1_STR_MAP[self._op_str]
        else:
            op = OP0_STR_MAP[self._op_str]
        return op

    def get_am(self, addr: str):
        # 得到目的操作数和源操作数, 返回(寻址方式, 地址)
        if addr is None:
            return None, None
        if addr in ALLOW_REGISTER_MAP:  # 如果是寄存器寻址, 且必须是允许的寄存器
            return pin.AM_REG, ALLOW_REGISTER_MAP[addr]
        if re.match(r"^[0-9]+$", addr):  # 如果是立即数
            return pin.AM_INS, int(addr)
        if re.match(r"^0[Xx][0-9A-Za-z]+$", addr):  # 如果是16进制立即数
            return pin.AM_INS, int(addr, 16)
        if re.match(r"^\[[0-9]+\]$", addr):  # 如果是立即数
            return pin.AM_DIR, int(addr[1:-1])
        if re.match(r"^\[0[Xx][0-9A-Za-z]+\]$", addr):  # 如果是16进制立即数
            return pin.AM_DIR, int(addr[1:-1], 16)
        if re.match(r"^\[.+\]$", addr):  # 如果是 [a]
            reg = addr[1:-1]
            if reg in ALLOW_REGISTER_MAP:
                return pin.AM_RAM, ALLOW_REGISTER_MAP[reg]

        raise Exception(f"操作数 地址错误: {self}")

    def compile_code(self):
        amd, dst = self.get_am(self._dst)  # 目的操作数
        ams, src = self.get_am(self._src)  # 源操作数
        # print(amd, dst, ams, src)

        if src and (amd, ams) not in ASM.INSTRUCTIONS[2][self.op]:
            raise Exception(f"操作数 地址错误: {self}")
        elif src is None and dst and amd not in ASM.INSTRUCTIONS[1][self.op]:
            raise Exception(f"操作数 地址错误: {self}")
        elif src is None and dst is None and self.op not in ASM.INSTRUCTIONS[0]:
            raise Exception(f"操作数 地址错误: {self}")

        amd = amd or 0
        dst = dst or 0
        ams = ams or 0
        src = src or 0

        # 根据op, 构建 ir寄存器
        if self.op in OP2_VALUE_SET:
            ir = self.op | amd << 2 | ams
        elif self.op in OP1_VALUE_SET:
            ir = self.op | amd
        else:
            ir = self.op
        return ir, dst, src


def compile_program(input_file: str):
    codes = []
    with open(input_file, "r", encoding="utf8") as f:
        lines = f.readlines()
        for idx, line in enumerate(lines):
            source = line.strip().split(";")[0]  # 去掉空格和注释
            if len(source) == 0:
                continue
            code = Code(idx + 1, source)
            codes.append(code)
    return codes


if __name__ == '__main__':
    data_width = 8
    addr_width = 16
    input_file = "./test_mov.asm"
    output_file = "test_mov.bin"

    program_bin_data = []
    for code in compile_program(input_file):
        values = code.compile_code()
        for value in values:
            program_bin_data.append(value)

    print(get_hex_str(program_bin_data, data_width / 4))
    save_hex_str(get_hex_str(program_bin_data, data_width / 4), 8, "program.bin")