# from select import kevent
# from os import sync
# from select import kevent
from logging import config
import numpy as np
import re
from typing import List
# from chip_define.reg_define import reg_define['mcu_reg'], addr_base['DTCM0_BASE']
from reg_define import *
class AssemblyTemplateManager:
    """汇编指令模板管理器"""

    def __init__(self, mk_instance, mk_instr, **kwargs):
        self.mk = mk_instance
        self.mk_instr = mk_instr
        self.config_file = kwargs.get('config_file')
        self.templates = {
            'general_send_wait': self._generic_awg_control_template,
            'ramp_fixed': self._ramp_mcu_fixed_template,
            'ramp_step': self._ramp_mcu_template
        }

    def create_instructions(self, **kwargs):
        template_type = kwargs.get('instr_type', str)
        params = {**kwargs}

        return self.templates[template_type](**params)

    def _codeword_encode(self, **kwargs):

        sendc        = kwargs.pop('sendc'       , 0)
        wave_hold    = kwargs.pop('wave_hold'   , 0)
        ff_amp_index = kwargs.pop('ff_amp_index', 0)
        fm_amp_index = kwargs.pop('fm_amp_index', 0)
        bias_index   = kwargs.pop('bias_index'  , 0)
        fcw_index    = kwargs.pop('fcw_index'   , 0)
        pcw_index    = kwargs.pop('pcw_index'   , 0)
        code_clr     = kwargs.pop('code_clr'    , 0)
        env_index    = kwargs.pop('env_index'   , 0)

        codeword = 0
        codeword  |= sendc        << 31
        codeword  |= wave_hold    << 30
        codeword  |= ff_amp_index << 28
        codeword  |= fm_amp_index << 26
        codeword  |= bias_index   << 24
        codeword  |= fcw_index    << 22
        codeword  |= pcw_index    << 19
        codeword  |= code_clr     << 18
        codeword  |= env_index    << 12

        return codeword

    def _codeword_gen(self, **kwargs):
        codeword_configs = kwargs.pop('codeword_configs', [])
        codeword_list = []
        for config in codeword_configs:
            codeword = self._codeword_encode(**config)
            codeword_list.append(codeword)
        return codeword_list

    def write_register(self, address, value):
        self.mk.rw_once('w', address, value, self.config_file)

    def _generic_awg_control_template(self, **kwargs):
        """
        通用 AWG汇编控制模版
        支持扫参功能
        """
        # ==========================================
        # 1. 准备并打包 DTCM 数据 (Python 侧)
        # ==========================================
        codeword_list = self._codeword_gen(**kwargs)
        send_interval_list = kwargs.get('send_interval', [100])
        cycle_num = kwargs.get('cycle_num', 1)
        # 提取扫参相关的参数配置
        sweep_config = kwargs.get('sweep_config', {})   # 扫描参数包
        sweep_num = sweep_config.get('sweep_num', 1)    # 扫描次数，代表最外层大循环需要mcu参数重载的次数
        sweep_offsets = sweep_config.get('offsets', [])  # mcu_regfile中的偏移地址，代表需要重载的寄存器
        sweep_steps = sweep_config.get('steps', [])     # 增量步进
        # 自动获取需要扫描的寄存器个数
        sweep_reg_num = len(sweep_offsets)

        # 构造 DTCM Payload 列表
        # [宏观头部] -> [波形序列] -> [Offsets表] -> [Starts表] -> [Steps表]
        # 宏观头部：重载次数，波形播放次数，码字个数,需要重载的寄存器个数
        dtcm_payload = [sweep_num, cycle_num, len(codeword_list), sweep_reg_num]

        # 压入波形序列：[码字0, 间隔0, 码字1, 间隔1 ...]
        for cw, wait_clk in zip(codeword_list, send_interval_list):
            dtcm_payload.append(cw)
            dtcm_payload.append(wait_clk)

        # 如果有扫参任务，压入地址表和步长表
        if sweep_reg_num > 0:
            dtcm_payload.extend(sweep_offsets)
            dtcm_payload.extend([s & 0xFFFFFFFF for s in sweep_steps])

        # 统一将这批动态参数写入到 DTFR(0xD8) 的下一个地址，即 0xDC 开始的内存中
        target_addr = addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['DTFR'] + 4
        self.write_register(target_addr, dtcm_payload)

        # ==========================================
        # 2. 生成标准 RISC-V 汇编指令文本 (MCU 侧)
        # ==========================================
        return f"""
        start:
            # ---------------------------------------------------------
            # 基地址初始化
            # ---------------------------------------------------------
            lui  x1 , 0x100                # x1 = 数据内存 (DTCM) 基地址 (0x00100000)
            lui  x2 , 0x200                # x2 = 硬件控制寄存器基地址 (0x00200000)

            # ---------------------------------------------------------
            # 批量搬移固化的 NCO 外设参数 (0x40 ~ 0x94)
            # 包含 22 个寄存器：FCW0~3, CWPRR, GAPR0~7, LCPR, AMPR0~3, BIASR0~3
            # ---------------------------------------------------------
            addi x3 , x0 , 4               # x3 = 4 (地址递增步长为 4 字节)
            addi x4 , x0 , 22              # x4 = 22 (需要搬移的寄存器总数)
            addi x5 , x1 , 0x40            # x5 = 数据源首地址 (DTCM 的 0x40 偏移)
            addi x6 , x2 , 0x40            # x6 = 目标首地址 (外设的 0x40 偏移)

        load_nco_params_loop:
            addi x4 , x4 , -1              # 循环计数减 1
            lw   x31, 0x00(x5)             # 从 DTCM 读出 1 个参数
            sw   x31, 0x00(x6)             # 写入到控制寄存器
            add  x5 , x5 , x3              # 源地址 + 4
            add  x6 , x6 , x3              # 目标地址 + 4
            bne  x4 , x0 , load_nco_params_loop

            # =========================================================
            # 单独搬移 0xB4 地址的参数 (1000b4 -> 2000b4) FM_EN 。至此上述共23个参数搬移完成
            # =========================================================
            lw   x31, 0xb4(x1)             # 从 DTCM (0x1000b4) 读出参数到 x31
            sw   x31, 0xb4(x2)             # 将 x31 写入到外设控制寄存器 (0x2000b4)
            
            # ---------------------------------------------------------
            # 模块 A: 读取宏观参数 (0xDC)
            # ---------------------------------------------------------
            lw   x10, 0xdc(x1)             # x10 = sweep_num (扫描总步数)
            lw   x11, 0xe0(x1)             # x11 = cycle_num (系综平均次数)
            lw   x12, 0xe4(x1)             # x12 = wave_num  (波形数)
            lw   x17, 0xe8(x1)             # x17 = sweep_reg_num (修改个数)
            
            addi x13, x1 , 0xec            # x13 = 波形序列首地址 (0xEC)
            slli x14, x12, 3               
            add  x18, x13, x14             # x18 = [Offsets] 地址表首地址
            slli x14, x17, 2               
            add  x20, x18, x14             # x20 = [Steps] 步长表首地址

           # ==================== 主控执行入口 (Do-While 模式) ====================
        run_sweep_iteration:
            # ====== 1. 系综平均循环 ======
            addi x28, x11, 0               
        middle_ensemble_loop:
            addi x28, x28, -1              
            addi x26, x12, 0               
            addi x25, x13, 0       
                    
            # ====== 2. 波形发送内循环 ======
        inner_wave_send_loop:
            addi x26, x26, -1              
            lw   x31, 0x00(x25)            
            lw   x30, 0x04(x25)            
            addi x25, x25, 8               

            send x0 , x31, 0      
                     
            bne  x26, x0 , inner_wait_branch
            beq  x0 , x0 , outer_wait_branch
    
        inner_wait_branch:
            wait x0 , x30, -24             
            bne  x26, x0 , inner_wave_send_loop
    
        outer_wait_branch:
            wait x0 , x30, -36             
            bne  x28, x0 , middle_ensemble_loop
    
            # ====== 3. 扫参结束判断 ======
            # 如果 sweep_num == 0，说明当前序列已经打完，直接下班！
            beq  x10, x0 , mcu_exit        
            addi x10, x10, -1              # 否则步数 -1，准备更新参数
            wait x0 , x0 , 100             # 参数切换保护死时间
    
            # ====== 4. 硬件 ALU 更新寄存器 ======
            beq  x17, x0 , run_sweep_iteration # 若没配寄存器(防呆)，直接进入下一轮
            addi x4 , x17, 0               # 循环次数
            addi x21, x18, 0               # 游标 x21 -> Offsets
            addi x23, x20, 0               # 游标 x23 -> Steps
    
        update_param_loop:
            addi x4 , x4 , -1
            lw   x29, 0(x21)               # 读 相对偏移地址 (例如 0x78)
            lw   x24, 0(x23)               # 读 步长 Step
            
            add  x5 , x1 , x29             # x5 = DTCM中该参数的地址 (x1 + 0x78)
            lw   x31, 0(x5)                # 从 DTCM 读出当前真值！
            add  x31, x31, x24             # 当前值 = 当前值 + 步长
            sw   x31, 0(x5)                # 将新值存回 DTCM，作为下一次的基准！
    
            add  x6 , x2 , x29             # x6 = 硬件物理地址 (x2 + 0x78)
            sw   x31, 0(x6)                # 写入硬件生效
    
            addi x21, x21, 4               # 游标下移
            addi x23, x23, 4               
            bne  x4 , x0 , update_param_loop
    
            # 参数更新完毕，无条件跳回上面执行新一轮波形
            beq  x0 , x0 , run_sweep_iteration
    
        mcu_exit:
            exit x0 , x0 , 0               
        """


    def _ramp_mcu_template(self, **kwargs):
        ramp_mcu_registers = []
        ramp_mcu_registers.append(0 << 16)
        ramp_mcu_registers += [1 << 31]
        param_num = kwargs.pop('param_num')
        ensemble_num = kwargs.pop('ensemble_num')
        ramp_mcu_registers.append(param_num)
        ramp_mcu_registers.append(ensemble_num)
        height_list = kwargs.pop('height', 0)
        length_list = kwargs.pop('step_time', 0)
        for height, length in zip(height_list, length_list):
            ramp_mcu_registers += [height << 16]
            ramp_mcu_registers += [length]
            wait = 65536 / height * length
            ramp_mcu_registers += [wait]
        self.write_register(addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['DTFR'] + 4, ramp_mcu_registers)
        # todo: configre step and width first, and then enable?
        return f"""
           start:
               lui  x1 , 0x100
               lui  x2 , 0x200
               lw   x31, 0xdc(x1)
               sw   x31, 0xb8(x2)
               lw   x28, 0xe0(x1)
               addi x6 , x0, 12
               lw   x7 , 0xe8(x1)
           ensemble_loop:
               addi x7 , x7, -1
               addi x8 , x1, 0
               lw   x5 , 0xe4(x1)
           ramp_loop:
               addi x5,  x5, -1
               lw   x31, 0xec(x8)
               lw   x30, 0xf0(x8)
               lw   x29, 0xf4(x8)
               add  x8 , x8 , x6
               sw   x30, 0xc0(x2)
               sw   x31, 0xbc(x2)
               sw   x28, 0xc4(x2)
               wait x0 , x29, -30
               bne  x5 , x0 , ramp_loop
               bne  x7 , x0 , ensemble_loop
               sw   x0, 0xc4(x2)
               exit x0, x0, 0                      
           """

    def _ramp_mcu_fixed_template(self, **kwargs):
        ramp_mcu_registers = []
        ramp_mcu_registers += [1 << 31]  # RAMPENR
        ensemble_num = kwargs.pop('ensemble_num')
        config_param_num = kwargs.pop('config_param_num')
        ramp_mcu_registers.append(ensemble_num)
        ramp_mcu_registers.append(config_param_num)
        fixed_value_list = kwargs.pop('fixed_value')
        wait_list = kwargs.pop('wait_clk')
        for fixed_value, wait_clk in zip(fixed_value_list, wait_list):
            ramp_mcu_registers += [fixed_value << 16 | 1 << 15]
            ramp_mcu_registers += [wait_clk]
        self.write_register(addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['DTFR'] + 4, ramp_mcu_registers)
        return f"""
           start:
               lui  x1 , 0x100
               lui  x2 , 0x200
               lw   x31, 0xdc(x1)
               sw   x31, 0xc4(x2)
               addi x3 , x0, 8
               lw   x4 , 0xe0(x1)
           ensemble_loop:
               addi x4 , x4, -1
               addi x6 , x1, 0
               lw   x5 , 0xe4(x1)
           ramp_loop:
               addi x5,  x5, -1
               lw   x31, 0xe8(x6)
               lw   x30, 0xec(x6)
               sw   x31, 0xb8(x2)
               add  x6 , x6 , x3
               bne  x5 , x0 , ramp_loop_wait
               jal  x0 , ensemble_loop_wait
           ramp_loop_wait:
               wait x0 , x30, -24
               jal  x0 , ramp_loop
           ensemble_loop_wait:
               wait x0 , x30, -36
               bne  x4 , x0 , ensemble_loop
           exit:
               wait x0 , x0, 15
               sw   x0, 0xb8(x2)
               sw   x0, 0xc4(x2)
               exit x0, x0, 0                      
           """

        
    def _write_machine_codes_to_chip(self, machine_codes: str, **kwargs):
        channel_id = kwargs.get('channel_id', 0)
        self.mk_instr.write(machine_codes, self.config_file, chip_id = channel_id, show=kwargs.pop('instr_show', False))
        if 'inner_sync' in kwargs:
            inner_sync = kwargs.pop('inner_sync')
            if inner_sync:
                self.write_register(addr_base['SYST_BASE']+reg_define['sys_reg']['SYNCR'], 3<<16 | 1)
                self.mk.rw_once('r', addr_base['SYST_BASE']+reg_define['pll_reg']['INTPLL_CLKRXPD'], [0]*20, self.config_file)
                self.mk.rw_once('r', addr_base['DBGM_BASE'], [0]*2048, self.config_file)

def instruction_config(mk_instance, mk_instr, **kwargs):
    asm_templates = AssemblyTemplateManager(mk_instance, mk_instr, **kwargs)
    machine_codes = asm_templates.create_instructions(**kwargs)
    asm_templates._write_machine_codes_to_chip(machine_codes, **kwargs)
