# 
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,  ALU经常使用这俩寄存器, 还是禁止使用算了
    # "b": pin.B,
    "c": pin.C,
    "d": pin.D,
    "ss": pin.SS,
    "sp": pin.SP,
}

OP2_STR_MAP = {
    "mov": ASM.MOV,
    "add": ASM.ADD,
    "sub": ASM.SUB,
    "cmp": ASM.CMP,
    "and": ASM.AND,
    "or": ASM.OR,
    "xor": ASM.XOR,
}

OP1_STR_MAP = {
    "inc": ASM.INC,
    "dec": ASM.DEC,
    "not": ASM.NOT,
    "jmp": ASM.JMP,
    "jo": ASM.JO,
    "jno": ASM.JNO,
    "jz": ASM.JZ,
    "jnz": ASM.JNZ,
    "jp": ASM.JP,
    "jnp": ASM.JNP,
    "push": ASM.PUSH,
    "pop": ASM.POP,
    "call": ASM.CALL,
    "int": ASM.INT,
}

OP0_STR_MAP = {
    "nop": ASM.NOP,
    "hlt": ASM.HLT,
    "ret": ASM.RET,
    "iret": ASM.IRET,
    "sti": ASM.STI,
    "cli": ASM.CLI,
}

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:
    TYPE_CODE = 1
    TYPE_LABLE = 2
    def __init__(self, number, source):
        self.number: int = number  # 源文件中行号
        self.source: str = source  # 源代码
        self._op_str = self._src = self._dst = None

        self.code_type = ...
        self.name = ...
        self.code_from_bin_addr = ...  # 对应bin文件中的位置

        self.prepare_source()

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

    def prepare_source(self):
        if self.source.endswith(":"):
            self.code_type = self.TYPE_LABLE
            self.name = self.source.strip(":")
            return

        self.code_type = self.TYPE_CODE

        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):
        # 得到目的操作数和源操作数, 返回(寻址方式, 地址)
        global marks
        if addr in marks:
            return pin.AM_INS, marks[addr].code_from_bin_addr * 3  # 我们所有的指令都是占了3个字节: ir dst src
        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)

    # 保底的
    codes.append(Code(idx + 1 + 1, "hlt"))

    ret = []
    marks = {}
    current_code = ...
    for code in reversed(codes):  # 从后往前
        if code.code_type == Code.TYPE_CODE:
            current_code = code
            ret.insert(0, code)
            continue
        if code.code_type == Code.TYPE_LABLE:
            marks[code.name] = current_code
            continue

        raise Exception("语法之外的 错误")

    # 更新代码在bin中的索引
    for idx, code in enumerate(ret):
        code.code_from_bin_addr = idx

    return ret, marks


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

    program_bin_data = []
    ret_codes, marks = compile_program(input_file)
    for code in ret_codes:
        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")