import argparse
import sys
import xlrd
import os

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Generate SRAM wrapper verilog files')
    parser.add_argument('xls_file', help='Excel configuration file path')
    parser.add_argument('-name', help='Specify sheet name to process')
    return parser.parse_args()

def validate_sheet(rb, sheet_name):
    """验证Sheet是否存在"""
    try:
        return rb.sheet_by_name(sheet_name)
    except xlrd.biffh.XLRDError:
        print(f"[Error] Sheet '{sheet_name}' not found in {rb}")
        sys.exit(1)



def generate_sram_wrapper(template_file, sram_name, width, depth, ref_name,\
                            asmbly_depth_nums,asmbly_width_nums,compile_depth,compile_width,\
                            wtsel,\
                          output_dir):
    """生成单个SRAM包装文件"""
    with open(template_file, 'r') as f:
        template = f.read()
    
    # 执行模板替换（根据您提供的模板结构）
    replaced = template.replace('${SramWrapName}', sram_name)\
                      .replace('${Width}', str(width))\
                      .replace('${Depth}', str(depth))\
                      .replace('${AssemblyDepth}',str(asmbly_depth_nums))\
                      .replace('${AssemblyWidth}',str(asmbly_width_nums))\
                      .replace('${CompileDepth}',str(compile_depth))\
                      .replace('${CompileWidth}',str(compile_width))\
                      .replace('${WTselRange}', str(wtsel))\
                      .replace('${ReferenceName}', ref_name)

    # 写入输出文件
    output_path = f"{output_dir}/{sram_name}.v"
    with open(output_path, 'w') as f:
        f.write(replaced)

def get_mem_ctrl_bus_fields(i_compilerName,mux,width, depth):
    compilerName = i_compilerName.split('_')[0]
    # get WTSEL
    wtsel = '[45:44]'
    if compilerName == 'tsn12ffcllspsram':
        if (mux == 4 and (depth>1024 and depth <= 4096)) or \
            (mux == 8 and (depth>2048 and depth <= 8192)) or \
            (mux == 16 and (depth>4096 and depth <= 16384)):
            wtsel = '[47:46]'
    
    return wtsel

def gen_sram_sh(rb,sheet_name, output_dir):
    """生成SRAM相关的Shell脚本"""
    config_dir = os.path.join(output_dir, "config")
    os.makedirs(config_dir, exist_ok=True)
    try:
        sheet = rb.sheet_by_name(sheet_name)
        headers = sheet.row_values(1)
        
        # 获取关键列索引
        name_col = headers.index('SramWrapName')
        compiler_col = headers.index('CompilerName') if 'CompilerName' in headers else -1
        
        if compiler_col == -1:
            print(f"[Error] Missing 'CompilerName' column in sheet {sheet_name}")
            return

        commands = []
        for row_idx in range(2, sheet.nrows):
            row = sheet.row_values(row_idx)
            sram_name = row[name_col]
            compiler = row[compiler_col]
            
            if sram_name and compiler:  # 空值检查
                cmd = f"../{compiler}.pl -file {sram_name}_{compiler}_config.txt -ColRed -NonBist"
                commands.append(cmd)
            else:
                print(f"[Warning] Missing data in row {row_idx} for SRAM '{sram_name}' with compiler '{compiler}', skipping.")


        # 写入bash文件
        sh_path = os.path.join(config_dir, f"run_{sheet_name}.sh")
        with open(sh_path, 'w') as f:
            f.write("#!/bin/bash\n")
            f.write("\n".join(sorted(commands)))
        os.chmod(sh_path, 0o755)  # 设置可执行权限
            
    except Exception as e:
        print(f"Generate shell script failed: {str(e)}")

def gen_config_txt(sram_name, compiler, compile_depth, compile_width, mux, output_dir):
    """生成SRAM配置文本文件到config子目录"""
    config_dir = os.path.join(output_dir, "config")
    os.makedirs(config_dir, exist_ok=True)
    
    if(compiler == 'tsn12ffcllspsram_130b'):
        config_content = f"{compile_depth}x{compile_width}m{mux}scp ulvt "
    elif(compiler == 'tsn12ffcll2prf_130a'):
        config_content = f"{compile_depth}x{compile_width}m{mux}scp ulvt "
    elif(compiler == 'tsn12ffclluhd2prf_130b'):
        config_content = f"{compile_depth}x{compile_width}m{mux} ulvt "
    elif(compiler == 'tsn12ffcllshcspsbsram_130c'):
        config_content = f"{compile_depth}x{compile_width}m{mux}s ulvt "
    elif(compiler == 'tsn12ffcll1prf_130a'):
        config_content = f"{compile_depth}x{compile_width}m{mux}scp ulvt "
    else:
        config_content = f"{compile_depth}x{compile_width}m{mux}scp ulvt "
    
    config_path = os.path.join(config_dir, f"{sram_name}_{compiler}_config.txt")
    with open(config_path, 'w') as f:
        f.write(config_content)
        f.write("\n")


def generate_files(xls_path, sheet_name=None):
    """主生成函数"""
    rb = xlrd.open_workbook(xls_path)
    
    # 处理Sheet选择逻辑
    sheets = [rb.sheet_by_name(sheet_name)] if sheet_name else rb.sheets()
    
    for sheet in sheets:
        # 获取列索引（根据模板参数）
        headers = sheet.row_values(1)
        col_map = {
            'name' : headers.index('SramWrapName'),
            'width': headers.index('Width'),
            'depth': headers.index('Depth'),
            'ref'  : headers.index('ReferenceName'),
            'port' : headers.index('PortType'),
            'AssemblyDepth' : headers.index('AssemblyDepth'),
            'AssemblyWidth' : headers.index('AssemblyWidth'),
            'CompileDepth'  : headers.index('CompileDepth'),
            'CompileWidth'  : headers.index('CompileWidth'),
            'ASYNC'         : headers.index('Async'),
            'MUX'           : headers.index('MUX') if 'MUX' in headers else -1,
            'CompilerName'   : headers.index('CompilerName') if 'CompilerName' in headers else -1
        }

        # 新增列存在性检查
        required_columns = ['SramWrapName', 'Width', 'Depth', 'ReferenceName', 
                          'PortType', 'AssemblyDepth', 'AssemblyWidth']
        missing_cols = [col for col in required_columns if col not in headers]
        if missing_cols:
            error_msg = f"[Error] Sheet '{sheet.name}' missing columns: {', '.join(missing_cols)}"
            if sheet_name:  # 指定特定sheet时直接退出
                print(error_msg)
                sys.exit(1)
            print(error_msg + ", skipping...")
            continue

        print(f"Processing sheet: {sheet.name} with columns: {col_map}")
        # 创建输出目录
#        output_dir = f"output_{sheet.name}"
        output_dir = f"sram_wrap_{sheet.name}"
        os.makedirs(output_dir, exist_ok=True)
        
        sram_port = ""
        template_file = ""
        # 处理每行数据
        for row_idx in range(2, sheet.nrows):
            row = sheet.row_values(row_idx)
            # 检查必要列是否有值
            if not row[col_map['name']] or not row[col_map['AssemblyWidth']] or not row[col_map['AssemblyDepth']]:
                print(f"[Warning] Missing required data in row {row_idx}, skipping.")
                continue
            # generate sram wrap
            if int(row[col_map['AssemblyDepth']]) + int(row[col_map['AssemblyWidth']]) == 2:
                if row[col_map['port']] == 'SP':
                    sram_port     = 'SRAM_SP'
                    template_file = "template_sram_sp_wrap.v"
                elif row[col_map['port']] == 'TP':
                    sram_port = 'SRAM_SYNC_TP' if (row[col_map['ASYNC']] == 'No' or row[col_map['ASYNC']] == 'NO') else 'SRAM_ASYNC_TP'
                    template_file = "template_sram_tp_wrap.v" if (row[col_map['ASYNC']] == 'No') \
                                    else "template_sram_tp_async_wrap.v"
            else: # Assembly
                if row[col_map['port']] == 'SP':
                    sram_port     = 'SRAM_SP'
                    template_file = "template_sram_sp_wrap_asmbly.v"
                elif row[col_map['port']] == 'TP':
                    sram_port = 'SRAM_SYNC_TP_ASMBLY' if (row[col_map['ASYNC']] == 'No' or row[col_map['ASYNC']] == 'NO') else 'SRAM_ASYNC_TP_ASMBLY'
                    template_file = "template_sram_tp_wrap_asmbly.v" if (row[col_map['ASYNC']] == 'No') \
                                    else "template_sram_tp_async_wrap_asmbly.v"

            print(f"Generated {sram_port} wrapper for: {row[col_map['name']]}")

            # 获取WTSEL
            WTsel = get_mem_ctrl_bus_fields(row[col_map['CompilerName']], int(row[col_map['MUX']]), 
                                                int(row[col_map['CompileWidth']]), int(row[col_map['CompileDepth']]))
            # 调用生成函数
            generate_sram_wrapper(
                # input file_path
                template_file = template_file,
                sram_name = row[col_map['name']],
                width = int(row[col_map['width']]),
                depth = int(row[col_map['depth']]),
                ref_name = row[col_map['ref']],
                asmbly_depth_nums = int(row[col_map['AssemblyDepth']]),
                asmbly_width_nums = int(row[col_map['AssemblyWidth']]),
                compile_depth     = int(row[col_map['CompileDepth']]),
                compile_width     = int(row[col_map['CompileWidth']]),
                wtsel               = WTsel ,
                output_dir = output_dir
            )

            gen_config_txt(
                sram_name = row[col_map['name']],
                compiler = row[col_map['CompilerName']],
                compile_depth = int(row[col_map['CompileDepth']]),
                compile_width = int(row[col_map['CompileWidth']]),
                mux = int(row[col_map['MUX']]) if col_map['MUX'] != -1 else 1,
                output_dir = output_dir
            )
        
        # scan the all files of the directory/ to be inclued into xxx.lst
        lst_path = os.path.join(output_dir, f"{output_dir}.lst")
#        # 获取目录中所有.v文件k按字母顺序排序）
#        v_files = [f for f in os.listdir(output_dir) if f.endswith('.v')]
        # 获取目录中所有.v文件并添加路径前缀
        v_files = [f"$PROJ_ROOT/rtl/common/mem/{output_dir}/{f}" 
                 for f in os.listdir(output_dir) 
                 if f.endswith('.v')]
        with open(lst_path, 'w') as f:
            f.write('\n'.join(sorted(v_files)))

        gen_sram_sh(rb, sheet.name, output_dir)

if __name__ == "__main__":
    args = parse_args()
    
    # 必须参数检查
    if not os.path.exists(args.xls_file):
        print(f"[Error] File not found: {args.xls_file}")
        sys.exit(1)
        
    generate_files(args.xls_file, args.name)