import numpy as np

class make_inst(object):

    def parse_instruction(self, instruction, labels, pc):
        # 去掉所有逗号
        instruction = instruction.replace(',', ' ')
        parts = instruction.split()
        opcode = parts[0].upper().strip()

        if opcode.endswith(':'):
            # 处理标签
            label_name = opcode[:-1]
            labels[label_name] = pc
            return None

        operands = [op.strip() for op in parts[1:]]

        def parse_immediate(imm_str):
            try:
                if imm_str.startswith('0x') or imm_str.startswith('0X'):
                    return int(imm_str, 16)
                elif imm_str.startswith('0b') or imm_str.startswith('0B'):
                    return int(imm_str, 2)
                elif imm_str.startswith('-0x') or imm_str.startswith('-0X'):
                    return -int(imm_str[1:], 16)
                elif imm_str.startswith('-0b') or imm_str.startswith('-0B'):
                    return -int(imm_str[1:], 2)
                elif imm_str.startswith('-'):
                    return -int(imm_str[1:], 10)
                else:
                    return int(imm_str, 10)
            except ValueError:
                raise ValueError(f"Invalid immediate value: {imm_str}")

        if opcode == 'LUI':
            rd, imm = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            return format((self.opcode_map[opcode]) | (imm & 0xFFFFF) << 12 | (rd << 7), '032b')

        elif opcode == 'AUIPC':
            rd, imm = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            return format((self.opcode_map[opcode]) | (imm & 0xFFFFF) << 12 | (rd << 7), '032b')

        elif opcode == 'JAL':
            rd, label_or_imm = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            if label_or_imm.lstrip('-').isdigit() or label_or_imm.startswith(('0x', '0X', '0b', '0B', '-0x', '-0X', '-0b', '-0B')):
                imm = parse_immediate(label_or_imm)
            else:
                imm = labels.get(label_or_imm.upper().strip(), 0) - pc
            imm_bits = (((imm >> 20) & 0x1) << 19) | (((imm >> 1) & 0x3FF) << 9) | (((imm >> 11) & 0x1) << 8) | ((imm >> 12) & 0xFF)
            return format((self.opcode_map[opcode]) | (imm_bits) << 12 | (rd << 7), '032b')

        elif opcode == 'JALR':
            rd = operands[0]
            operands[1] = operands[1].rstrip(')')
            imm, rs1 = operands[1].split('(')
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | (imm & 0xFFF) << 20 | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')

        elif opcode in ['BEQ', 'BNE', 'BLT', 'BGE', 'BLTU', 'BGEU']:
            rs1, rs2, label_or_imm = operands
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            rs2 = int(rs2[1:])  # 去掉寄存器名称前的 'x'
            if label_or_imm.lstrip('-').isdigit() or label_or_imm.startswith(('0x', '0X', '0b', '0B', '-0x', '-0X', '-0b', '-0B')):
                imm = parse_immediate(label_or_imm)
            else:
                imm = labels.get(label_or_imm.upper().strip(), 0) - pc
            imm_high_bits = (((imm >> 12) & 0x1) << 6) | (((imm >> 5 ) & 0x3F)) 
            imm_low_bits  = (((imm >> 1 ) & 0xF) << 1) | (((imm >> 11) & 0x1 ))
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | (imm_high_bits << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | (imm_low_bits << 7), '032b')

        elif opcode in ['LB', 'LH', 'LW', 'LBU', 'LHU']:
            rd = operands[0]
            operands[1] = operands[1].rstrip(')')
            imm, rs1 = operands[1].split('(')
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | ((imm & 0xFFF) << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')

        elif opcode in ['SB', 'SH', 'SW']:
            rs2 = operands[0]
            operands[1] = operands[1].rstrip(')')
            imm, rs1 = operands[1].split('(')
            rs2 = int(rs2[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | (((imm >> 5) & 0x7F) << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | ((imm & 0x1F) << 7), '032b')

        elif opcode in ['ADDI', 'SLTI', 'SLTIU', 'XORI', 'ORI', 'ANDI']:
            rd, rs1, imm = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | ((imm & 0xFFF) << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')

        elif opcode in ['SLLI', 'SRLI', 'SRAI']:
            rd, rs1, shamt = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            shamt = parse_immediate(shamt)
            funct3 = self.opcode_funct3_map[opcode]
            funct7 = self.opcode_funct7_map[opcode]
            return format((self.opcode_map[opcode]) | (funct7 << 25) | ((shamt & 0x1F) << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')

        elif opcode in ['ADD', 'SUB', 'SLL', 'SLT', 'SLTU', 'XOR', 'SRL', 'SRA', 'OR', 'AND']:
            rd, rs1, rs2 = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            rs2 = int(rs2[1:])  # 去掉寄存器名称前的 'x'
            funct3 = self.opcode_funct3_map[opcode]
            funct7 = self.opcode_funct7_map[opcode]
            return format((self.opcode_map[opcode]) | (funct7 << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')
        
        elif opcode in ['WAIT', 'SEND', 'SENDC']:
            rd, rs1, imm = operands
            rd = int(rd[1:])  # 去掉寄存器名称前的 'x'
            rs1 = int(rs1[1:])  # 去掉寄存器名称前的 'x'
            imm = parse_immediate(imm)
            funct3 = self.opcode_funct3_map[opcode]
            return format((self.opcode_map[opcode]) | ((imm & 0xFFF) << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7), '032b')
        
        elif opcode in ['EXIT']:
            return format((self.opcode_map[opcode]), '032b')
        
        elif opcode in ['EXIT_IR']:
            return '00000000000000000001000000101011'

        else:
            raise ValueError(f"Unsupported opcode: {opcode}")

    def write(self, instructions, file_name, exaddr = 1, chip_id = 0, channel_id = 0, pc_start = 0, ard_flag = 0, show = False):
        if instructions == "":
            pass
        else:
            labels = {}
            binary_instructions = []

            # 将整段汇编代码拆成多条指令组成的字符串数组
            inst_list = instructions.split('\n')
            instructions = []
            for this_inst in inst_list:
                this_inst = this_inst.split('#')
                if this_inst[0].strip() != '':
                    instructions.append(this_inst[0].strip())

            # 第一遍扫描：记录标签位置
            pc = pc_start
            for instr in instructions:
                binary = self.parse_instruction(instr, labels, pc)
                if binary is not None:
                    binary_instructions.append(binary)
                    pc += 4
                else:
                    # 如果是标签，不增加pc
                    pass

            # 第二遍扫描：生成最终的二进制代码
            pc = pc_start
            final_binary_instructions = []
            for instr in instructions:
                binary = self.parse_instruction(instr, labels, pc)
                if binary is not None:
                    final_binary_instructions.append(binary)
                    if show:
                        print(f"{instr}: {binary}")
                    pc += 4
                else:
                    # 如果是标签，不增加pc
                    pass

            with open(file_name, "a") as f:
                base_addr = 0x010_0000 + pc_start + channel_id * 0x060_0000
                length = np.size(final_binary_instructions)
                f.write(f"{((ard_flag << 30) | (chip_id << 25) | (base_addr)):08x}\n")
                f.write(f"{((exaddr << 20) | (length<<2)):08x}\n")
                for binary_instr in final_binary_instructions:
                    f.write(f"{int(binary_instr,2):08x}\n")
                f.write("\n") 

            return final_binary_instructions

    opcode_map = {
        'LUI': 0x37,
        'AUIPC': 0x17,
        'JAL': 0x6F,
        'JALR': 0x67,
        'BEQ': 0x63,
        'BNE': 0x63,
        'BLT': 0x63,
        'BGE': 0x63,
        'BLTU': 0x63,
        'BGEU': 0x63,
        'LB': 0x03,
        'LH': 0x03,
        'LW': 0x03,
        'LBU': 0x03,
        'LHU': 0x03,
        'SB': 0x23,
        'SH': 0x23,
        'SW': 0x23,
        'ADDI': 0x13,
        'SLTI': 0x13,
        'SLTIU': 0x13,
        'XORI': 0x13,
        'ORI': 0x13,
        'ANDI': 0x13,
        'SLLI': 0x13,
        'SRLI': 0x13,
        'SRAI': 0x13,
        'ADD': 0x33,
        'SUB': 0x33,
        'SLL': 0x33,
        'SLT': 0x33,
        'SLTU': 0x33,
        'XOR': 0x33,
        'SRL': 0x33,
        'SRA': 0x33,
        'OR': 0x33,
        'AND': 0x33,
        'WAIT': 0x0B,
        'SEND': 0x0B,
        'SENDC': 0x0B,
        'EXIT': 0x2B,
    }

    opcode_funct3_map = {
        'JALR': 0x0,
        'BEQ': 0x0,
        'BNE': 0x1,
        'BLT': 0x4,
        'BGE': 0x5,
        'BLTU': 0x6,
        'BGEU': 0x7,
        'LB': 0x0,
        'LH': 0x1,
        'LW': 0x2,
        'LBU': 0x4,
        'LHU': 0x5,
        'SB': 0x0,
        'SH': 0x1,
        'SW': 0x2,
        'ADDI': 0x0,
        'SLTI': 0x2,
        'SLTIU': 0x3,
        'XORI': 0x4,
        'ORI': 0x6,
        'ANDI': 0x7,
        'SLLI': 0x1,
        'SRLI': 0x5,
        'SRAI': 0x5,
        'ADD': 0x0,
        'SUB': 0x0,
        'SLL': 0x1,
        'SLT': 0x2,
        'SLTU': 0x3,
        'XOR': 0x4,
        'SRL': 0x5,
        'SRA': 0x5,
        'OR': 0x6,
        'AND': 0x7,
        'WAIT': 0x0,
        'SEND': 0x2,
        'SENDC': 0x3,
        'EXIT': 0x0,
    }

    opcode_funct7_map = {
        'JALR': 0x00,
        'BEQ': 0x00,
        'BNE': 0x00,
        'BLT': 0x00,
        'BGE': 0x00,
        'BLTU': 0x00,
        'BGEU': 0x00,
        'LB': 0x00,
        'LH': 0x00,
        'LW': 0x00,
        'LBU': 0x00,
        'LHU': 0x00,
        'SB': 0x00,
        'SH': 0x00,
        'SW': 0x00,
        'ADDI': 0x00,
        'SLTI': 0x00,
        'SLTIU': 0x00,
        'XORI': 0x00,
        'ORI': 0x00,
        'ANDI': 0x00,
        'SLLI': 0x00,
        'SRLI': 0x00,
        'SRAI': 0x20,
        'ADD': 0x00,
        'SUB': 0x20,
        'SLL': 0x00,
        'SLT': 0x00,
        'SLTU': 0x00,
        'XOR': 0x00,
        'SRL': 0x00,
        'SRA': 0x20,
        'OR': 0x00,
        'AND': 0x00,
        'WAIT': 0x00,
        'SEND': 0x00,
        'SENDC': 0x00,
        'EXIT': 0x00,
    }