#!/usr/bin/env python3
"""Batch fix: add null check before aclrtFree, set to nullptr after free, for all non_contig CPPs."""
import re
import os
import glob

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

def fix_aclrtfree_in_line(line):
    """Fix all aclrtFree calls in a single line."""
    if 'aclrtFreeFuzz' in line:
        return line, False
    if 'aclrtFree' not in line:
        return line, False

    changed = False

    # Pattern 1: if (workspaceSize > 0 && workspaceAddr) { aclrtFree(workspaceAddr); }
    # Already has null check, just add nullptr assignment
    m = re.search(r'if \(workspaceSize > 0 && workspaceAddr\) \{ aclrtFree\(workspaceAddr\); \}', line)
    if m:
        line = line.replace(
            m.group(0),
            'if (workspaceSize > 0 && workspaceAddr != nullptr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }'
        )
        changed = True
        return line, changed

    # Pattern 2: if(workspaceSize>0)aclrtFree(workspaceAddr);
    m = re.search(r'if\(workspaceSize>0\)aclrtFree\(workspaceAddr\);', line)
    if m:
        line = line.replace(
            m.group(0),
            'if (workspaceSize > 0 && workspaceAddr != nullptr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }'
        )
        changed = True
        return line, changed

    # Pattern 3: if (workspaceSize > 0) { aclrtFree(workspaceAddr); }
    m = re.search(r'if \(workspaceSize > 0\) \{ aclrtFree\(workspaceAddr\); \}', line)
    if m:
        line = line.replace(
            m.group(0),
            'if (workspaceSize > 0 && workspaceAddr != nullptr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }'
        )
        changed = True
        return line, changed

    # Pattern 4: if (xxx) aclrtFree(xxx); — array patterns, already has null check, add nullptr
    m = re.match(r'^(\s*)if \((\w+\[\w+\])\) aclrtFree\(\2\);(.*)$', line)
    if m:
        indent = m.group(1)
        var = m.group(2)
        rest = m.group(3)
        line = f'{indent}if ({var}) {{ aclrtFree({var}); {var} = nullptr; }}{rest}'
        changed = True
        return line, changed

    # Pattern 5: Simple aclrtFree(varName); — possibly multiple on one line
    # Also handles aclrtFree inside if (workspaceSize > 0) { ... } multi-line blocks
    def replace_simple(match):
        nonlocal changed
        var = match.group(1)
        changed = True
        return f'if ({var} != nullptr) {{ aclrtFree({var}); {var} = nullptr; }}'

    # Match aclrtFree(varName); but NOT aclrtFreeFuzz (negative lookbehind for 'Fuzz')
    line = re.sub(r'(?<!Fuzz)aclrtFree\((\w+)\);', replace_simple, line)

    return line, changed


def fix_file(filepath):
    """Fix a single CPP file."""
    with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
        lines = f.readlines()

    changed = False
    new_lines = []
    for line in lines:
        new_line, line_changed = fix_aclrtfree_in_line(line)
        new_lines.append(new_line)
        if line_changed:
            changed = True

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

    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
                if (i+1) % 50 == 0:
                    print(f"[{i+1}/{total}] ... (no changes needed so far for {filename})")
        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 (no aclrtFree or already fixed): {skipped}")
    print(f"Errors: {errors}")


if __name__ == '__main__':
    main()