#!/usr/bin/env python3
"""
gen_multithread_from_nc.py
从已有 non_contig CPP 机械转换生成 multithread CPP（ascendc-fuzz-multithread-cpp-gen 规则）。

转换规则（与框架已验证模板一致）：
  M1 入口函数：extern "C" int {api}_fuzz(int argc, char *argv[], aclrtContext context)（整块替换旧入口）
  M2 Normal 函数签名加 aclrtContext context；Init(deviceId,...) 替换为
     aclrtSetCurrentContext(context) + aclrtCreateStream(&stream)
  M4 scene 范围 0~5（含 CONTINUOUS/BROADCAST）：scene 下界 2 → 0
  M5 函数末尾 aclrtDestroyStream(stream)，不销毁 context/device（main_multithread 统一管理）
  M6 日志名 non_contig → multithread

用法：
  python gen_multithread_from_nc.py <nc.cpp> <out.cpp> <apiName>
"""
import re
import sys


def convert(nc_path, out_path, api):
    c = open(nc_path, errors='ignore').read()

    # 1) Normal 函数定义与调用（兼容 {api}_non_contig / {api}_non_contig_Normal 两种命名）
    c = re.sub(rf'int {re.escape(api)}_non_contig(_Normal)?\(int fuzz_num\)',
               f'int {api}_multithread_Normal(int fuzz_num, aclrtContext context)', c)
    c = re.sub(rf'{re.escape(api)}_non_contig(_Normal)?\(fuzz_num\)',
               f'{api}_multithread_Normal(fuzz_num, context)', c)

    # 2) Init → SetCurrentContext + CreateStream（lambda 防 repl 转义）
    c = re.sub(r'auto ret = Init\(deviceId, &context, &stream\);',
               lambda m: ('auto ret = aclrtSetCurrentContext(context);\n'
                          '    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d' + r'\n' + '", ret); return ret);\n'
                          '    ret = aclrtCreateStream(&stream);\n'
                          '    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d' + r'\n' + '", ret); return ret);'),
               c)

    # 3) scene 下界 2 → 0（CONTINUOUS/BROADCAST 可达）
    c = re.sub(r'DTSetGetNumberRange\(&g_Element\[i\+\+\], 2, 2, \(ndim >= 2 \? 5 : 4\), "i"\)',
               'DTSetGetNumberRange(&g_Element[i++], 2, 0, (ndim >= 2 ? 5 : 4), "i")', c)

    # 4) 日志名
    c = c.replace(f'"{api}_non_contig"', f'"{api}_multithread"')

    # 5) 末尾：DestroyStream 后移除 DestroyContext/ResetDevice/Finalize（context 由 main_multithread 管理）
    c = re.sub(r'aclrtDestroyStream\(stream\);\s*aclrtDestroyContext\(context\);\s*aclrtResetDevice\(deviceId\);\s*aclFinalize\(\);',
               'aclrtDestroyStream(stream);', c)

    # 6) 整块替换旧入口（extern "C" 可选前缀，从入口签名到文件尾由新入口接管）
    new_entry = f'''extern "C" int {api}_fuzz(int argc, char *argv[], aclrtContext context)
{{
    int fuzz_num = MAXFUZZNUMBER;
    if (argc > 2) {{ fuzz_num = atoi(argv[2]); }}
    printf("fuzz_num %d \\n", fuzz_num);
    printf("{api}_multithread_Normal :%d \\n", {api}_multithread_Normal(fuzz_num, context));
    return 0;
}}
'''
    m = re.search(rf'(extern "C"\s+)?int {re.escape(api)}_non_contig_fuzz\(', c)
    if m:
        c = c[:m.start()] + new_entry
    else:
        c = c.rstrip() + '\n\n' + new_entry
    open(out_path, 'w').write(c)
    print(f'Transformed: {nc_path} -> {out_path}')


if __name__ == '__main__':
    convert(sys.argv[1], sys.argv[2], sys.argv[3])