#!/usr/bin/env python3
"""
Inject [FUZZ_PARAM] print into files that don't have it yet.
Handles various NonContigParams variable naming: nc/ncSelf/ncInput, ncOut/ncOutput/outNC
"""
import re
import os
import glob

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

# Files to process
TARGET_FILES = [
    "aclnnBinaryCrossEntropyWithLogits_non_contig.cpp",
    "aclnnFmodScalar_non_contig.cpp",
    "aclnnFmodTensor_non_contig.cpp",
    "aclnnForeachAddcmulScalarV2_non_contig.cpp",
    "aclnnForeachAddScalarV2_non_contig.cpp",
    "aclnnForeachDivScalarList_non_contig.cpp",
    "aclnnForeachLerpScalar_non_contig.cpp",
    "aclnnForeachMulScalarV2_non_contig.cpp",
    "aclnnForeachSqrt_non_contig.cpp",
    "aclnnGatherV3_non_contig.cpp",
    "aclnnIndex_non_contig.cpp",
    "aclnnIndexPutImpl_non_contig.cpp",
    "aclnnInplaceLogicalOr_non_contig.cpp",
    "aclnnInplaceNormal_non_contig.cpp",
    "aclnnInplaceNormalTensor_non_contig.cpp",
    "aclnnInplaceRandomTensor_non_contig.cpp",
    "aclnnInplaceUniformTensor_non_contig.cpp",
    "aclnnLogicalOr_non_contig.cpp",
    "aclnnLogSoftmaxBackward_non_contig.cpp",
    "aclnnLogSumExp_non_contig.cpp",
    "aclnnMean_non_contig.cpp",
    "aclnnMoeFinalizeRoutingV2_non_contig.cpp",
    "aclnnMoeGatingTopK_non_contig.cpp",
    "aclnnMoeInitRouting_non_contig.cpp",
    "aclnnMoeInitRoutingV3_non_contig.cpp",
    "aclnnMoeTokenPermuteGrad_non_contig.cpp",
    "aclnnMoeTokenPermute_non_contig.cpp",
    "aclnnMseLoss_non_contig.cpp",
    "aclnnNormalFloatFloat_non_contig.cpp",
    "aclnnNormalFloatTensor_non_contig.cpp",
    "aclnnNormalTensorFloat_non_contig.cpp",
    "aclnnNormalTensorTensor_non_contig.cpp",
    "aclnnResize_non_contig.cpp",
    "aclnnRmsNormQuant_non_contig.cpp",
    "aclnnScatterPaKvCache_non_contig.cpp",
    "aclnnSoftmaxBackward_non_contig.cpp",
    "aclnnSplitTensor_non_contig.cpp",
]


def find_noncontig_vars(content):
    """Find all NonContigParams variable names and their CalcNonContiguousParams calls."""
    # Find all CalcNonContiguousParams(varName); lines
    calc_calls = re.findall(r'CalcNonContiguousParams\((\w+)\)\s*;', content)
    return calc_calls


def inject_fuzz_param(content):
    """Inject [FUZZ_PARAM] print after the last CalcNonContiguousParams call."""
    if '[FUZZ_PARAM]' in content:
        return content, False

    # Find all CalcNonContiguousParams calls
    calc_calls = find_noncontig_vars(content)
    if not calc_calls:
        return content, False

    # Find the last CalcNonContiguousParams line and its position
    last_calc_var = calc_calls[-1]
    # Find the line: CalcNonContiguousParams(last_calc_var); possibly followed by scene clamp and storageSize
    pattern = rf'(CalcNonContiguousParams\({last_calc_var}\);\s*\n(?:\s*if \({last_calc_var}\.scene < 2\) {last_calc_var}\.scene = 2;\s*\n)?\s*long long \w+StorageSize = GetShapeSize\({last_calc_var}\.storageShape\);\s*\n)'

    m = re.search(pattern, content)
    if not m:
        # Try simpler pattern
        pattern2 = rf'(CalcNonContiguousParams\({last_calc_var}\);\s*\n)'
        m = re.search(pattern2, content)
        if not m:
            return content, False

    # Build print statement for all NonContigParams variables
    print_lines = ['\n            printf("[FUZZ_PARAM] ");']
    for var in calc_calls:
        print_lines.append(f'            printf("{var}: scene=%s(%d), ", GetSceneName({var}.scene), {var}.scene);')
        print_lines.append(f'            PrintShape("{var}_viewShape", {var}.viewShape); PrintShape(", {var}_storageShape", {var}.storageShape);')
        print_lines.append(f'            printf(", strides=["); for (size_t j = 0; j < {var}.strides.size(); j++) printf("%lld%s", (long long){var}.strides[j], j < {var}.strides.size()-1 ? "," : ""); printf("], offset=%lld | ", (long long){var}.offset);')
    print_lines.append('            printf("\\n");\n')

    param_print = '\n'.join(print_lines)
    content = content[:m.end()] + param_print + content[m.end():]
    return content, True


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:
            with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
                content = f.read()

            content, changed = inject_fuzz_param(content)

            if changed:
                with open(filepath, 'w', encoding='utf-8') as f:
                    f.write(content)
                fixed += 1
                print(f"FIXED: {filename}")
            else:
                print(f"SKIP (no match): {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()