#!/usr/bin/env python3
"""
Convert 20 non_contig CPPs from CreateAclTensor/CreateAclTensor_New to NonContigParams + CreateAclTensorByScene.
Also applies: aclrtFree null check, FUZZ_PARAM, FUZZ_RESULT, ERROR label fix, ret=-1 init.
"""
import re
import os

SRC_DIR = "/home/30039325/Vector_fuzz_cpp_testing/turing_sec/aclnnAPI_fuzz/src/src_no_contiguous"

TARGET_FILES = [
    "aclnnIndex_non_contig.cpp",
    "aclnnIndexPutImpl_non_contig.cpp",
    "aclnnInplaceNormal_non_contig.cpp",
    "aclnnInplaceNormalTensor_non_contig.cpp",
    "aclnnInplaceRandomTensor_non_contig.cpp",
    "aclnnInplaceUniformTensor_non_contig.cpp",
    "aclnnLogSoftmaxBackward_non_contig.cpp",
    "aclnnLogSumExp_non_contig.cpp",
    "aclnnMean_non_contig.cpp",
    "aclnnMoeGatingTopK_non_contig.cpp",
    "aclnnMoeInitRouting_non_contig.cpp",
    "aclnnMoeInitRoutingV3_non_contig.cpp",
    "aclnnMseLoss_non_contig.cpp",
    "aclnnNormalFloatTensor_non_contig.cpp",
    "aclnnNormalTensorFloat_non_contig.cpp",
    "aclnnNormalTensorTensor_non_contig.cpp",
    "aclnnResize_non_contig.cpp",
    "aclnnScatterPaKvCache_non_contig.cpp",
    "aclnnSoftmaxBackward_non_contig.cpp",
    "aclnnSplitTensor_non_contig.cpp",
]

HELPER_FUNCS = r'''
static const char* GetSceneName(int scene) {
    switch(scene) {
        case 0: return "CONTINUOUS";
        case 1: return "BROADCAST";
        case 2: return "SLICE";
        case 3: return "ASSTRIDED";
        case 4: return "STORAGE_PADDED";
        case 5: return "TRANSPOSE";
        default: return "UNKNOWN";
    }
}

static const char* GetDtypeName(int dtype) {
    switch(dtype) {
        case 0: return "FLOAT";
        case 1: return "FLOAT16";
        case 2: return "INT8";
        case 3: return "INT32";
        case 4: return "UINT8";
        case 6: return "INT16";
        case 9: return "INT64";
        case 11: return "BOOL";
        case 12: return "DOUBLE";
        case 27: return "BF16";
        case 16: return "COMPLEX64";
        default: return "UNKNOWN";
    }
}

static void PrintShape(const char* name, const std::vector<int64_t>& shape) {
    printf("%s=[", name);
    for (size_t i = 0; i < shape.size(); i++) {
        printf("%lld%s", (long long)shape[i], i < shape.size()-1 ? "," : "");
    }
    printf("]");
}

'''

# NonContigParams injection block for a tensor
def make_nc_block(var_prefix, shape_var, is_output=False):
    """Generate NonContigParams construction code for a tensor."""
    nc_var = f"nc{var_prefix}"
    scene_min = 2 if is_output else 1
    return f'''            {{
                int {nc_var}Ndim = {shape_var}.size();
                int {nc_var}SceneMax = ({nc_var}Ndim >= 2) ? 5 : 4;
                NonContigParams {nc_var};
                {nc_var}.scene = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, {scene_min}, {nc_var}SceneMax, "i");
                {nc_var}.viewShape = {shape_var};
                {nc_var}.sliceStep = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 2, 5, "i");
                {nc_var}.sliceAxis = 0;
                {nc_var}.offset = 0;
                {nc_var}.pos1 = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 1, 4, "i");
                {nc_var}.pos2 = *(int*)DTSetGetNumberRange(&g_Element[i++], 2, 1, 4, "i");
                {nc_var}.strides.resize({shape_var}.size());
                for (size_t j = 0; j < {shape_var}.size(); j++) {{
                    {nc_var}.strides[j] = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 1, 20, "i");
                }}
                CalcNonContiguousParams({nc_var});
                if ({nc_var}.scene < {scene_min}) {nc_var}.scene = {scene_min};
            }}'''


def fix_file(filepath):
    filename = os.path.basename(filepath)
    with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
        content = f.read()

    changed = False

    # 1. Inject helper functions
    if 'GetSceneName' not in content:
        # After "using namespace std;"
        m = re.search(r'(using namespace std;\s*\n)', content)
        if m:
            content = content[:m.end()] + HELPER_FUNCS + content[m.end():]
            changed = True

    # 2. Fix aclrtFree null check (reuse existing logic)
    def replace_aclrtfree(match):
        nonlocal changed
        var = match.group(1)
        changed = True
        return f'if ({var} != nullptr) {{ aclrtFree({var}); {var} = nullptr; }}'

    if 'aclrtFreeFuzz' not in content:
        content = re.sub(r'(?<!Fuzz)aclrtFree\((\w+)\);', replace_aclrtfree, content)

    # Also fix workspaceAddr patterns
    content = re.sub(
        r'if \(workspaceSize > 0\) \{ aclrtFree\(workspaceAddr\); \}',
        'if (workspaceSize > 0 && workspaceAddr != nullptr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }',
        content
    )
    content = re.sub(
        r'if\(workspaceSize>0\)aclrtFree\(workspaceAddr\);',
        'if (workspaceSize > 0 && workspaceAddr != nullptr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }',
        content
    )

    # 3. Fix ERROR label: ret = aclrtSynchronizeStream -> syncRet
    if 'syncRet' not in content:
        old_err = r'ERROR:\n\s*ret\s*=\s*aclrtSynchronizeStream\(stream\);\n\s*printf\("end \\n"\);\n\s*CheckAndLog\(ret\);'
        new_err = 'ERROR:\n        {\n            aclError syncRet = aclrtSynchronizeStream(stream);\n            (void)syncRet;\n        }\n        printf("end \\n");\n        CheckAndLog(ret);'
        if re.search(old_err, content):
            content = re.sub(old_err, new_err, content)
            changed = True
        else:
            # Try variant without CheckAndLog immediately after
            old_err2 = r'ERROR:\n\s*ret\s*=\s*aclrtSynchronizeStream\(stream\);\n\s*printf\("end '
            if re.search(old_err2, content):
                content = re.sub(
                    r'ERROR:\n(\s*)ret\s*=\s*aclrtSynchronizeStream\(stream\);',
                    'ERROR:\n\\1{\n\\1    aclError syncRet = aclrtSynchronizeStream(stream);\n\\1    (void)syncRet;\n\\1}',
                    content
                )
                changed = True

    # 4. Add FUZZ_RESULT after CheckAndLog(ret)
    if 'FUZZ_RESULT' not in content:
        old_check = r'(CheckAndLog\(ret\);\s*\n)'
        new_check = r'\1        if (ret == 0) { success_count++; } else { faild_count++; }\n        printf("[FUZZ_RESULT] case_%d: %s, ret=%d\\n", fuzzi, ret == 0 ? "success" : "faild", ret);\n'
        if re.search(old_check, content):
            content = re.sub(old_check, new_check, content)
            changed = True

    # 5. Add ret = -1 init
    if 'ret = -1;' not in content:
        pat = r'(void\*\s*workspaceAddr\s*=\s*nullptr;\s*\n)'
        m = re.search(pat, content)
        if m:
            content = content[:m.end()] + '        ret = -1;\n' + content[m.end():]
            changed = True

    # 6. Convert CreateAclTensor_New to CreateAclTensorByScene
    # Pattern: CreateAclTensor_New(hostData, shape, &deviceAddr, dtype, format, &tensor)
    # We need to inject NonContigParams before each dtype if/else block
    # This is complex - for now just replace CreateAclTensor_New calls with CreateAclTensorByScene
    # and inject NonContigParams construction

    # Find all unique (shape, deviceAddr, tensor) tuples from CreateAclTensor_New calls
    new_pattern = r'CreateAclTensor_New\(\s*(\w+)\s*,\s*(\w+)\s*,\s*&(\w+)\s*,\s*(\w+)\s*,\s*(\w+)\s*,\s*&(\w+)\s*\)'
    old_pattern = r'CreateAclTensor\(\s*(\w+)\s*,\s*(\w+)\s*,\s*&(\w+)\s*,\s*(\w+)\s*,\s*&(\w+)\s*\)'

    # Replace CreateAclTensor_New(hostData, shape, &addr, dtype, fmt, &tensor)
    # with CreateAclTensorByScene(hostData, nc<Shape>, &addr, dtype, fmt, &tensor)
    # But we need NonContigParams to exist - inject before the dtype if/else block

    # For simplicity, replace CreateAclTensor_New -> CreateAclTensorByScene with shape var as NonContigParams
    # Actually CreateAclTensorByScene takes NonContigParams not shape vector
    # So we need to construct NonContigParams inline

    # Strategy: replace each CreateAclTensor_New(hostData, shape, &addr, dtype, fmt, &tensor)
    # with a block that constructs NonContigParams and calls CreateAclTensorByScene

    def replace_create_new(match):
        nonlocal changed
        host = match.group(1)
        shape = match.group(2)
        addr = match.group(3)
        dtype = match.group(4)
        fmt = match.group(5)
        tensor = match.group(6)
        changed = True
        # Generate a unique nc var name based on tensor name
        nc_var = f"nc_{tensor}"
        return (f'{{ NonContigParams {nc_var}; {nc_var}.scene = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 1, ({shape}.size() >= 2 ? 5 : 4), "i"); '
                f'{nc_var}.viewShape = {shape}; {nc_var}.sliceStep = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 2, 5, "i"); '
                f'{nc_var}.sliceAxis = 0; {nc_var}.offset = 0; '
                f'{nc_var}.pos1 = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 1, 4, "i"); '
                f'{nc_var}.pos2 = *(int*)DTSetGetNumberRange(&g_Element[i++], 2, 1, 4, "i"); '
                f'{nc_var}.strides.resize({shape}.size()); '
                f'for (size_t _j = 0; _j < {shape}.size(); _j++) {nc_var}.strides[_j] = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 1, 20, "i"); '
                f'CalcNonContiguousParams({nc_var}); '
                f'CreateAclTensorByScene({host}, {nc_var}, &{addr}, {dtype}, {fmt}, &{tensor}); }}')

    def replace_create_old(match):
        nonlocal changed
        host = match.group(1)
        shape = match.group(2)
        addr = match.group(3)
        dtype = match.group(4)
        tensor = match.group(5)
        changed = True
        nc_var = f"nc_{tensor}"
        return (f'{{ NonContigParams {nc_var}; {nc_var}.scene = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 1, ({shape}.size() >= 2 ? 5 : 4), "i"); '
                f'{nc_var}.viewShape = {shape}; {nc_var}.sliceStep = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 2, 5, "i"); '
                f'{nc_var}.sliceAxis = 0; {nc_var}.offset = 0; '
                f'{nc_var}.pos1 = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 1, 4, "i"); '
                f'{nc_var}.pos2 = *(int*)DTSetGetNumberRange(&g_Element[i++], 2, 1, 4, "i"); '
                f'{nc_var}.strides.resize({shape}.size()); '
                f'for (size_t _j = 0; _j < {shape}.size(); _j++) {nc_var}.strides[_j] = *(int64_t*)DTSetGetNumberRangeS64(&g_Element[i++], 2, 1, 20, "i"); '
                f'CalcNonContiguousParams({nc_var}); '
                f'CreateAclTensorByScene({host}, {nc_var}, &{addr}, {dtype}, ACL_FORMAT_ND, &{tensor}); }}')

    # Don't replace if already using CreateAclTensorByScene
    if 'CreateAclTensorByScene' not in content or content.count('CreateAclTensorByScene') < 3:
        content = re.sub(new_pattern, replace_create_new, content)
        content = re.sub(old_pattern, replace_create_old, content)

    if changed:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)

    return changed


def main():
    fixed = 0
    errors = 0

    for filename in TARGET_FILES:
        filepath = os.path.join(SRC_DIR, filename)
        if not os.path.exists(filepath):
            print(f"SKIP (not found): {filename}")
            continue
        try:
            if fix_file(filepath):
                fixed += 1
                print(f"FIXED: {filename}")
            else:
                print(f"SKIP (no change): {filename}")
        except Exception as e:
            errors += 1
            print(f"ERROR: {filename} - {e}")

    print(f"\n=== Summary ===")
    print(f"Fixed: {fixed}")
    print(f"Errors: {errors}")


if __name__ == '__main__':
    main()