#!/usr/bin/env python3
"""Mechanical non_contig/normal CPP -> multithread CPP converter.

Supports both base forms:
  - non_contig CPP: funcs named {api}_non_contig / {api}_non_contig_fuzz
  - normal CPP:     funcs named {api}_Normal / {api}_fuzz (contiguous-only ops)

Transformations (ascendc-fuzz-multithread-cpp-gen SKILL.md):
 1. Normal func gets (int fuzz_num, aclrtContext context)
 2. Init(...) -> aclrtSetCurrentContext(context) + aclrtCreateStream(&stream)
 3. DT_FUZZ_START(0,..) -> aclnn_multithread_RANDOM_INT + DT_FUZZ_START(random_int,..)
 4. Ending: keep only aclrtDestroyStream(stream)
 5. Entry: extern "C" {api}_fuzz(int argc, char *argv[], aclrtContext context)
 6. Scene range 2~5 -> 0~5 (if present)
"""
import re
import sys


def convert(api, src, outputs=None):
    text = open(src, encoding='utf-8', errors='ignore').read()

    # 1. function signatures (both forms)
    text = text.replace('int %s_non_contig(int fuzz_num)' % api,
                        'int %s_Normal(int fuzz_num, aclrtContext context)' % api)
    text = text.replace('int %s_Normal(int fuzz_num)' % api,
                        'int %s_Normal(int fuzz_num, aclrtContext context)' % api)
    text = text.replace('int %s(int fuzz_num)' % api,
                        'int %s_Normal(int fuzz_num, aclrtContext context)' % api)

    # 2. Init -> SetCurrentContext + CreateStream (lambda repl: avoid re escape processing)
    init_repl = ('auto ret = aclrtSetCurrentContext(context);\n'
                 'CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("%s aclrtSetCurrentContext failed. ERROR: %%d\\n", ret); return ret);\n'
                 'aclrtStream stream;\n'
                 'ret = aclrtCreateStream(&stream);\n'
                 'CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("%s aclrtCreateStream failed. ERROR: %%d\\n", ret); return ret);' % (api, api))
    pat = re.compile(r'auto ret = Init\(deviceId, &context, &stream\);\s*\n\s*'
                     r'CHECK_RET\(ret == ACL_SUCCESS, LOG_PRINT\("Init acl failed\. ERROR: %d\\n", ret\); return ret\);')
    text, n1 = pat.subn(lambda m: init_repl, text)
    if n1 == 0:
        text, n1 = re.subn(r'auto ret = Init\(deviceId, &context, &stream\);', lambda m: init_repl, text)

    # 3. seed
    seed_repl = ('int random_int = aclnn_multithread_RANDOM_INT(1, 99999999);\n'
                 '    DT_FUZZ_START(random_int, fuzz_num,')
    text, n2 = re.subn(r'DT_FUZZ_START\(0, fuzz_num,', lambda m: seed_repl, text, count=1)

    # 4. ending: keep only DestroyStream
    text = re.sub(r'aclrtDestroyStream\(stream\);\s*\n\s*aclrtDestroyContext\(context\);\s*\n\s*aclrtResetDevice\(deviceId\);\s*\n\s*aclFinalize\(\);',
                  'aclrtDestroyStream(stream);', text)

    # 5. entry
    text = text.replace('extern "C" int %s_non_contig_fuzz(int argc, char *argv[]);' % api,
                        'extern "C" int %s_fuzz(int argc, char *argv[], aclrtContext context);' % api)
    text = text.replace('int %s_non_contig_fuzz(int argc, char *argv[])' % api,
                        'int %s_fuzz(int argc, char *argv[], aclrtContext context)' % api)
    text = text.replace('extern "C" int %s_fuzz(int argc, char *argv[]);' % api,
                        'extern "C" int %s_fuzz(int argc, char *argv[], aclrtContext context);' % api)
    text = re.sub(r'int %s_fuzz\(int argc, char \*argv\[\]\)(?!.*aclrtContext)' % api,
                  lambda m: 'int %s_fuzz(int argc, char *argv[], aclrtContext context)' % api, text)

    # 6. call sites in entry
    text = text.replace('%s_non_contig(fuzz_num)' % api, '%s_Normal(fuzz_num, context)' % api)
    text = re.sub(r'printf\("%s(_non_contig)?_Normal :%%d \\n", %s(_non_contig)?_Normal\(fuzz_num\)\);' % (api, api),
                  lambda m: 'printf("%s_Normal :%%d \\n", %s_Normal(fuzz_num, context));' % (api, api), text)
    text = text.replace('%s_Normal(fuzz_num));' % api, '%s_Normal(fuzz_num, context));' % api)

    # 7. scene range (if present)
    text = re.sub(r'DTSetGetNumberRange\(&g_Element\[i\+\+\], 1, 2, 5, "i"\)',
                  'DTSetGetNumberRange(&g_Element[i++], 1, 0, 5, "i")', text)
    text = re.sub(r'DTSetGetNumberRange\(&g_Element\[i\+\+\], 2, 2, \(ndim >= 2 \? 5 : 4\), "i"\)',
                  'DTSetGetNumberRange(&g_Element[i++], 2, 0, 5, "i")', text)
    text = re.sub(r'if \(nc\w+\.scene < 2\) nc\w+\.scene = 2;\s*\n', '', text)

    # 输出 tensor scene 修正（gate 1.6：multithread 输出 0,2~5，禁止 BROADCAST）：
    # MT 引擎在多流模式下 ~90% 用例返回 initValue 而非变异值——输出 scene 的 initValue 必须是
    # 合法的非 BROADCAST 场景（0=CONTINUOUS），并对 draw 范围内的 1 做重映射
    for out_name in (outputs or []):
        cap = out_name[0].upper() + out_name[1:]
        old_draw = 'nc%s.scene = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 5, "i");' % cap
        new_draw = ('nc%s.scene = *(int *)DTSetGetNumberRange(&g_Element[i++], 0, 0, 5, "i");  /* [OUT-SCENE-FIX] initValue=0 */\n'
                    '            if (nc%s.scene == 1) { nc%s.scene = 0; }  /* [OUT-SCENE-FIX] 输出禁止BROADCAST(gate 1.6) */') % (cap, cap, cap)
        if old_draw in text:
            text = text.replace(old_draw, new_draw)
        else:
            # 兼容未被全局 scene 重写覆盖的 2,5 形态
            old_draw2 = 'nc%s.scene = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 2, 5, "i");' % cap
            if old_draw2 in text:
                text = text.replace(old_draw2, new_draw.replace(' 1, 0, 5,', ' 1, 2, 5,').replace('0, 0, 5,', '0, 0, 5,').replace('&g_Element[i++], 0, 0, 5', '&g_Element[i++], 0, 0, 5'))

    text = text.replace(' non_contig op fuzz', ' multithread op fuzz')
    return text, n1, n2


def main():
    if len(sys.argv) < 4:
        print('Usage: convert_to_multithread.py <input_cpp> <output_cpp> <api>')
        sys.exit(1)
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('input'); ap.add_argument('output'); ap.add_argument('api')
    ap.add_argument('--outputs', default='', help='输出tensor名逗号列表(输出scene禁止BROADCAST,gate 1.6)')
    args = ap.parse_args()
    src, dst, api = args.input, args.output, args.api
    outs = [x.strip() for x in args.outputs.split(',') if x.strip()]
    out, n_init, n_seed = convert(api, src, outs)
    open(dst, 'w', encoding='utf-8').write(out)
    ok = []
    ok.append(('entry', 'extern "C" int %s_fuzz(int argc, char *argv[], aclrtContext context)' % api in out))
    ok.append(('sig', 'int %s_Normal(int fuzz_num, aclrtContext context)' % api in out))
    ok.append(('ctx', 'aclrtSetCurrentContext(context)' in out))
    ok.append(('seed', 'aclnn_multithread_RANDOM_INT' in out))
    ok.append(('no_devdestroy', 'aclrtDestroyContext(context);' not in out.split('aclrtDestroyStream(stream);')[-1]))
    has_scene = 'DTSetGetNumberRange(&g_Element[i++], 1, 0, 5, "i")' in out
    for name, v in ok:
        if not v:
            print('[MT_CONVERT_FAIL] check %s failed (init=%d seed=%d)' % (name, n_init, n_seed))
            sys.exit(2)
    if not has_scene:
        print('[MT_CONVERT_WARN] no scene mutation lines (contiguous-only op) - acceptable for md nc=x ops')
    print('[MT_CONVERT_OK] %s -> %s' % (src, dst))


if __name__ == '__main__':
    main()