#!/usr/bin/env python3
"""
Batch fix for all non_contig CPPs:
1. aclrtFree null check + set nullptr (already done, skip if already fixed)
2. Output tensor scene excludes BROADCAST (min=2)
3. [FUZZ_PARAM] print: dtype, viewShape, storageShape, strides, offset, scene
4. [FUZZ_RESULT] print + ERROR label not overwriting ret + ret=-1 init per case
"""
import re
import os
import glob

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

# Helper functions to inject at top of file (after kDtypeEnum/kDtypeNum)
HELPER_FUNCS = '''
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("]");
}

'''


def has_helper_funcs(content):
    return 'GetSceneName' in content and 'GetDtypeName' in content


def inject_helper_funcs(content):
    """Inject helper functions after kDtypeNum line."""
    if has_helper_funcs(content):
        return content, False
    # Find the kDtypeNum line
    pattern = r'(static int kDtypeNum\s*=\s*sizeof\(kDtypeEnum\)\s*/\s*sizeof\(kDtypeEnum\[0\]\)\s*;\s*\n)'
    m = re.search(pattern, content)
    if m:
        content = content[:m.end()] + '\n' + HELPER_FUNCS + content[m.end():]
        return content, True
    # Alternative: inject after "using namespace std;"
    pattern2 = r'(using namespace std;\s*\n)'
    m2 = re.search(pattern2, content)
    if m2:
        content = content[:m2.end()] + '\n' + HELPER_FUNCS + content[m2.end():]
        return content, True
    return content, False


def fix_output_scene(content):
    """Fix output tensor scene to exclude BROADCAST (min=2)."""
    changed = False
    # Pattern: ncOut.scene = ... DTSetGetNumberRange(..., 1, 1, sceneMax/outSceneMax/(dim_num...), "i")
    # Change min from 1 to 2
    patterns = [
        (r'ncOut\.scene\s*=\s*\*\(int\*\)\s*DTSetGetNumberRange\(&g_Element\[i\+\+\],\s*1,\s*1,\s*sceneMax,\s*"i"\)',
         'ncOut.scene = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 2, sceneMax, "i")'),
        (r'ncOut\.scene\s*=\s*\*\(int\*\)\s*DTSetGetNumberRange\(&g_Element\[i\+\+\],\s*1,\s*1,\s*\(dim_num\s*>=\s*2\s*\?\s*5\s*:\s*4\),\s*"i"\)',
         'ncOut.scene = *(int*)DTSetGetNumberRange(&g_Element[i++], 1, 2, (dim_num >= 2 ? 5 : 4), "i")'),
    ]
    for pat, repl in patterns:
        if re.search(pat, content):
            content = re.sub(pat, repl, content)
            changed = True

    # Also handle outNC.scene variant
    pat_outnc = r'outNC\.scene\s*=\s*\*\(int\s*\*\)\s*DTSetGetNumberRange\(&g_Element\[i\+\+\],\s*1,\s*1,'
    if re.search(pat_outnc, content):
        content = re.sub(pat_outnc, 'outNC.scene = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 2,', content)
        changed = True

    # Add scene clamp after CalcNonContiguousParams(ncOut) if not present
    # Look for CalcNonContiguousParams(ncOut) or CalcNonContiguousParams(outNC) not followed by scene clamp
    for nc_var in ['ncOut', 'outNC', 'outputOutNC']:
        calc_pat = f'CalcNonContiguousParams\\({nc_var}\\);'
        clamp_pat = f'if ({nc_var}.scene < 2) {nc_var}.scene = 2;'
        if re.search(calc_pat, content) and not re.search(clamp_pat, content):
            content = re.sub(
                f'(CalcNonContiguousParams\\({nc_var}\\);)',
                f'\\1\n            {clamp_pat}',
                content
            )
            changed = True

    return content, changed


def fix_error_label_and_result(content):
    """Fix ERROR label to not overwrite ret, add FUZZ_RESULT, add ret=-1 init."""
    changed = False

    # Check if this is checkRet mode (13 files)
    is_checkret = 'checkRet = CheckAndLog' in content

    if is_checkret:
        # These files already have FUZZ_RESULT and proper ret handling
        # Just add ret=-1 init if missing
        if 'ret = -1;' not in content:
            # Find the workspaceAddr = nullptr; line and add ret = -1; after it
            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
        return content, changed

    # Standard mode: fix ERROR label
    # Pattern: ERROR:\n        ret = aclrtSynchronizeStream(stream);\n        printf("end \n");\n        CheckAndLog(ret);
    # Replace with: ERROR:\n        { aclError syncRet = aclrtSynchronizeStream(stream); (void)syncRet; }\n        printf("end \n");\n        CheckAndLog(ret);

    # Check if already fixed with syncRet
    if 'syncRet' not in content:
        old_error = r'ERROR:\n\s*ret\s*=\s*aclrtSynchronizeStream\(stream\);\n\s*printf\("end \\n"\);\n\s*CheckAndLog\(ret\);'
        new_error = 'ERROR:\n        {\n            aclError syncRet = aclrtSynchronizeStream(stream);\n            (void)syncRet;\n        }\n        printf("end \\n");\n        CheckAndLog(ret);'
        if re.search(old_error, content):
            content = re.sub(old_error, new_error, content)
            changed = True

    # Add FUZZ_RESULT after CheckAndLog(ret) if missing
    if 'FUZZ_RESULT' not in content:
        # Add after CheckAndLog(ret);
        old_check = r'(CheckAndLog\(ret\);\n)'
        new_check = '\\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

    # Add ret = -1 init per case if missing
    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

    return content, changed


def add_fuzz_param_print(content):
    """Add [FUZZ_PARAM] print after CalcNonContiguousParams for output tensor."""
    changed = False

    # Check if already has our FUZZ_PARAM with viewShape/storageShape
    if 'FUZZ_PARAM' in content and 'viewShape' in content and 'storageShape' in content:
        return content, False

    # Find the pattern: CalcNonContiguousParams(ncOut);\n ... long long outStorageSize = GetShapeSize(ncOut.storageShape);
    # Insert FUZZ_PARAM print before the dtype if/else block

    # Pattern for ncOut
    pat_ncout = r'(CalcNonContiguousParams\(ncOut\);\s*\n\s*(?:if \(ncOut\.scene < 2\) ncOut\.scene = 2;\s*\n\s*)?long long outStorageSize = GetShapeSize\(ncOut\.storageShape\);\s*\n)'
    m = re.search(pat_ncout, content)
    if m and 'FUZZ_PARAM' not in content[:m.end()+100]:
        # Find the inputDtype variable name - look for patterns like "inputDtype" or "dtype0"
        # Try to find CalcNonContiguousParams(nc); ... selfStorageSize
        pat_nc = r'CalcNonContiguousParams\(nc\);\s*\n\s*long long selfStorageSize = GetShapeSize\(nc\.storageShape\);'
        m_nc = re.search(pat_nc, content)

        if m_nc:
            # Build the print statement
            param_print = (
                '\n'
                '            printf("[FUZZ_PARAM] dtype=%s, ", GetDtypeName((int)inputDtype));\n'
                '            printf("self: "); PrintShape("viewShape", nc.viewShape); PrintShape(",storageShape", nc.storageShape);\n'
                '            printf(", strides=["); for (size_t j = 0; j < nc.strides.size(); j++) { printf("%lld%s", (long long)nc.strides[j], j < nc.strides.size()-1 ? "," : ""); } printf("]");\n'
                '            printf(", offset=%lld, scene=%s(%d)", (long long)nc.offset, GetSceneName(nc.scene), nc.scene);\n'
                '            printf(" | out: "); PrintShape("viewShape", ncOut.viewShape); PrintShape(",storageShape", ncOut.storageShape);\n'
                '            printf(", strides=["); for (size_t j = 0; j < ncOut.strides.size(); j++) { printf("%lld%s", (long long)ncOut.strides[j], j < ncOut.strides.size()-1 ? "," : ""); } printf("]");\n'
                '            printf(", offset=%lld, scene=%s(%d)\\n", (long long)ncOut.offset, GetSceneName(ncOut.scene), ncOut.scene);\n'
            )
            content = content[:m.end()] + param_print + content[m.end():]
            changed = True

    return content, changed


def fix_file(filepath):
    """Fix a single CPP file with all 4 modifications."""
    with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
        content = f.read()

    changed = False

    # 1. Inject helper functions (GetSceneName, GetDtypeName, PrintShape)
    content, c = inject_helper_funcs(content)
    if c:
        changed = True

    # 2. Fix output tensor scene to exclude BROADCAST
    content, c = fix_output_scene(content)
    if c:
        changed = True

    # 3. Add FUZZ_PARAM print
    content, c = add_fuzz_param_print(content)
    if c:
        changed = True

    # 4. Fix ERROR label + FUZZ_RESULT + ret=-1 init
    content, c = fix_error_label_and_result(content)
    if c:
        changed = True

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

    return changed


def main():
    cpp_files = sorted(glob.glob(os.path.join(SRC_DIR, "*_non_contig.cpp")))
    total = len(cpp_files)
    fixed = 0
    skipped = 0
    errors = 0

    print(f"Total files: {total}")
    print(f"Source dir: {SRC_DIR}")
    print()

    for i, filepath in enumerate(cpp_files):
        filename = os.path.basename(filepath)
        try:
            if fix_file(filepath):
                fixed += 1
                print(f"[{i+1}/{total}] FIXED: {filename}")
            else:
                skipped += 1
        except Exception as e:
            errors += 1
            print(f"[{i+1}/{total}] ERROR: {filename} - {e}")

    print()
    print(f"=== Summary ===")
    print(f"Total: {total}")
    print(f"Fixed: {fixed}")
    print(f"Skipped (already fixed): {skipped}")
    print(f"Errors: {errors}")


if __name__ == '__main__':
    main()