#!/usr/bin/env python3
"""
Batch optimize fuzz CPP files:
1. shape=[] support: dim_num range 1,8 -> 0,8 + add dim==0 branch
2. Tensor size distribution: add sizeBucket logic + totalSize check

Only modifies files where dim_num min=1 (supports 1D => should support 0D).
Skips files with fixed dimensions (dim_num min>=2 or no dim_num).
"""

import os
import re
import sys

# Cube/matmul operators to exclude
CUBE_EXCLUDE = {
    "aclnnMoeComputeExpertTokens",
}

def get_api_name(filename):
    """Extract api name from filename, stripping _non_contig suffix."""
    base = filename.replace(".cpp", "")
    base = base.replace("_non_contig", "")
    return base

def should_exclude(api_name):
    """Check if operator should be excluded (cube/matmul)."""
    for excl in CUBE_EXCLUDE:
        if excl in api_name:
            return True
    return False

def modify_dim_num_zero(content):
    """
    1. Find dim_num declarations with range min=1, change to min=0
    2. Add dim_num==0 branch (shape=[]) before dim_num==1 branch
    
    Only modifies dim_num variables with range starting at 1.
    Skips dim_num with range starting at 2+ (fixed dimension operators).
    """
    modified = content
    changes = []
    
    # Pattern to find dim_num declarations with min=1
    # Matches: int dim_num0 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 8, "i");
    # Also matches variations with different spacing, int*, etc.
    dim_num_pattern = re.compile(
        r'(int\s+(dim_num\w*)\s*=\s*\*\(int\s*\*?\)\s*DTSetGetNumberRange\s*'
        r'\(\s*&g_Element\s*\[\s*i\+\+\s*\]\s*,\s*'
        r'\d+\s*,\s*'      # byte width
        r'1\s*,\s*'         # min=1 (the key: only modify when min is 1)
        r'(\d+)\s*,\s*'     # max
        r'"i"\s*\))'
    )
    
    # Find all dim_num declarations with min=1
    dim_num_vars = []  # list of (var_name, max_val)
    for m in dim_num_pattern.finditer(modified):
        var_name = m.group(2)
        max_val = m.group(3)
        dim_num_vars.append((var_name, max_val))
        changes.append(f"  dim_num range: {var_name} 1,{max_val} -> 0,{max_val}")
    
    if not dim_num_vars:
        return modified, []
    
    # Change range from 1 to 0 for matched dim_num declarations
    modified = dim_num_pattern.sub(
        lambda m: m.group(1).replace(
            m.group(0)[m.group(0).find(', 1,') if ', 1,' in m.group(0) else m.group(0).find(',1,'):],
            ', 0,' + m.group(0)[m.group(0).find(', 1,') + 4 if ', 1,' in m.group(0) else m.group(0).find(',1,') + 3:]
        ) if False else m.group(0).replace(', 1,', ', 0,').replace(',1,', ',0,'),
        modified
    )
    
    # Actually, let's do the replacement more carefully
    # Re-read content and do it properly
    modified = content
    
    # Step 1: Replace dim_num range min=1 with min=0
    def replace_dim_num_range(m):
        full_match = m.group(0)
        # Replace the min value from 1 to 0
        # The pattern has: ..., <bytewidth>, 1, <max>, "i")
        # We need to replace just the 1 that's the min
        parts = re.split(r'(\s*,\s*)', full_match)
        # Find the position of "1" that's between bytewidth and max
        # parts would be like: ['int dim_num0 = *(int *)DTSetGetNumberRange(&g_Element[i++]', ', ', '2', ', ', '1', ', ', '8', ', ', '"i")']
        # We need to find the '1' that comes after the bytewidth (2) and before max (8)
        result_parts = list(parts)
        found_bytewidth = False
        for idx, part in enumerate(parts):
            if found_bytewidth and part.strip() == '1':
                result_parts[idx] = '0'
                break
            if part.strip().isdigit() and not found_bytewidth:
                found_bytewidth = True
        return ''.join(result_parts)
    
    modified = dim_num_pattern.sub(replace_dim_num_range, modified)
    
    # Step 2: Add dim_num==0 branch before dim_num==1
    # For each dim_num variable that was modified, find its if/else chain
    for var_name, max_val in dim_num_vars:
        # Pattern to find: if (dim_numX == 1) { someShape = {shape1}; }
        # Or: if(dim_numX == 1){
        # We need to insert before it: if (dim_numX == 0) { someShape = {}; } else 
        # But we need to know the shape variable name
        
        # Find the pattern: if(dim_numX == 1) or if (dim_numX == 1)
        # And look at what shape is being assigned
        if_pattern = re.compile(
            r'(\s*)(if\s*\(\s*' + re.escape(var_name) + r'\s*==\s*1\s*\)\s*\{?\s*)'
            r'(\w+)\s*=\s*\{',
            re.MULTILINE
        )
        
        m = if_pattern.search(modified)
        if m:
            indent = m.group(1)
            shape_var = m.group(3)
            # Insert dim_num==0 branch before the if(dim_num==1)
            insert_text = f"{indent}if ({var_name} == 0) {{ {shape_var} = {{}}; }} else "
            # Insert before the if(dim_num==1)
            insert_pos = m.start()
            modified = modified[:insert_pos] + insert_text + modified[insert_pos:]
            changes.append(f"  added dim==0 branch for {var_name} -> {shape_var} = {{}}")
        else:
            changes.append(f"  WARNING: could not find if({var_name}==1) pattern to insert dim==0 branch")
    
    return modified, changes


def add_size_distribution(content, filename):
    """
    Add sizeBucket distribution control logic:
    1. Insert sizeBucket DT variation + totalSizeMin/Max after 'try {' or 'try{'
    2. Insert totalSize check after shape size checks
    
    Must be careful with DT element alignment - the sizeBucket consumes one DT element.
    """
    modified = content
    changes = []
    
    # Check if sizeBucket already exists
    if 'sizeBucket' in modified:
        changes.append("  sizeBucket already exists, skipping")
        return modified, changes
    
    # Find the try block opening
    # Pattern: try\n{ or try { or try{
    try_pattern = re.compile(r'(\s*)(try\s*\n?\s*\{)')
    try_match = try_pattern.search(modified)
    if not try_match:
        changes.append("  WARNING: could not find try block, skipping size distribution")
        return modified, changes
    
    # Find the position right after 'try {'
    try_end = try_match.end()
    indent = try_match.group(1)
    inner_indent = indent + "    "
    
    # The sizeBucket code to insert
    sizebucket_code = f"""
{inner_indent}// === Tensor Size distribution control ===
{inner_indent}int sizeBucket = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 3, "i");
{inner_indent}long long totalSizeMin = 0, totalSizeMax = 0;
{inner_indent}switch (sizeBucket) {{
{inner_indent}    case 0: totalSizeMin = 1;          totalSizeMax = 1000;          break;
{inner_indent}    case 1: totalSizeMin = 1001;       totalSizeMax = 1073741824LL;  break;
{inner_indent}    case 2: totalSizeMin = 1073741825LL; totalSizeMax = 2147483648LL; break;
{inner_indent}    case 3: totalSizeMin = 2147483649LL; totalSizeMax = 4294967296LL; break;
{inner_indent}}}
"""
    
    # Insert sizeBucket code after try {
    modified = modified[:try_end] + sizebucket_code + modified[try_end:]
    changes.append("  added sizeBucket DT variation after try block")
    
    # Now find where to insert the totalSize check
    # It should go after all shape construction and size checks, but before tensor creation
    # 
    # Strategy: Find the first ACTUAL tensor creation call (CreateAclTensor/CreateAclTensorByScene)
    # This is more reliable than dtype if-blocks which may be constraint adjustments in some files.
    # Search for shape size variables (GetShapeSize) in the region between try and tensor creation.
    
    # Priority 1: CreateAclTensor / CreateAclTensorByScene / CreateAclTensor_New calls
    create_tensor_pattern = re.compile(
        r'(\s*)(ret\s*=\s*CreateAclTensor\w*\s*\(|CreateAclTensor\w*\s*\()',
        re.MULTILINE
    )
    
    # Priority 2: aclCreateScalar (e.g. alpha = aclCreateScalar in aclnnAdd)
    scalar_pattern = re.compile(r'(\s*)(alpha\s*=\s*aclCreateScalar|=\s*aclCreateScalar)', re.MULTILINE)
    
    # Priority 3: aclCreateIntArray
    intarray_pattern = re.compile(r'(\s*)(=\s*aclCreateIntArray)', re.MULTILINE)
    
    # Priority 4: dtype if-block (fallback for files that construct tensors inside dtype if-else)
    dtype_if_pattern = re.compile(
        r'(\s*)(if\s*\(\s*(?:dtype\w*|inputDtype|dtype0)\s*==\s*aclDataType::ACL_FLOAT\s*\))',
        re.MULTILINE
    )
    
    # Find the first tensor creation point after try block
    # Try each pattern in priority order
    tensor_creation_pos = None
    tensor_indent = inner_indent
    
    ct_match = create_tensor_pattern.search(modified, try_end)
    sc_match = scalar_pattern.search(modified, try_end)
    ia_match = intarray_pattern.search(modified, try_end)
    dt_match = dtype_if_pattern.search(modified, try_end)
    
    if ct_match:
        tensor_creation_pos = ct_match.start()
        tensor_indent = ct_match.group(1)
    elif sc_match:
        tensor_creation_pos = sc_match.start()
        tensor_indent = sc_match.group(1)
    elif ia_match:
        tensor_creation_pos = ia_match.start()
        tensor_indent = ia_match.group(1)
    elif dt_match:
        tensor_creation_pos = dt_match.start()
        tensor_indent = dt_match.group(1)
    
    if tensor_creation_pos is None:
        changes.append("  WARNING: could not find tensor creation point for totalSize check")
        return modified, changes
    
    # Find all shape_size variables defined before tensor_creation_pos
    # Pattern: long long shape_sizeN = GetShapeSize(...)
    # Or: long long xxxStorageSize = GetShapeSize(...)
    # Also: xxxShapeSize = GetShapeSize(...) (declared earlier without long long prefix)
    shape_size_vars = []
    
    # Search between try_end and tensor_creation_pos for shape size variables
    region = modified[try_end:tensor_creation_pos]
    
    # Match: long long <varname> = GetShapeSize(...)
    size_var_pattern = re.compile(
        r'long\s+long\s+(\w+)\s*=\s*GetShapeSize\s*\('
    )
    for m in size_var_pattern.finditer(region):
        var = m.group(1)
        if var not in shape_size_vars:
            shape_size_vars.append(var)
    
    # Also match: <varname> = GetShapeSize(...) without long long prefix
    # (variable was declared earlier as long long)
    # Match any variable assigned from GetShapeSize, not just ones with "Size" in name
    bare_size_pattern = re.compile(
        r'(?<!\w)(\w+)\s*=\s*GetShapeSize\s*\('
    )
    for m in bare_size_pattern.finditer(region):
        var = m.group(1)
        if var not in shape_size_vars and var not in ('GetShapeSize', 'auto'):
            shape_size_vars.append(var)
    
    # Also look for shape_size variables that might be computed differently
    # e.g., long long shape_size = GetShapeSize(selfShape);
    if not shape_size_vars:
        # Try alternative patterns
        alt_pattern = re.compile(r'long\s+long\s+(\w+Size\w*)\s*=')
        for m in alt_pattern.finditer(region):
            var = m.group(1)
            if var not in shape_size_vars:
                shape_size_vars.append(var)
    
    if not shape_size_vars:
        # If we can't find shape size variables, try a wider search
        # Look for any GetShapeSize call and extract the result variable
        wider_region = modified[try_end:tensor_creation_pos]
        getsize_pattern = re.compile(r'(?:long\s+long\s+|auto\s+|int\s+)(\w+)\s*=\s*(?:static_cast<\w+>\()?\s*GetShapeSize\s*\(')
        for m in getsize_pattern.finditer(wider_region):
            var = m.group(1)
            if var not in shape_size_vars:
                shape_size_vars.append(var)
    
    if not shape_size_vars:
        # Last resort: find GetShapeSize calls and create a sum
        # Pattern: GetShapeSize(someShapeVar)
        shape_var_pattern = re.compile(r'GetShapeSize\s*\(\s*(\w+)\s*\)')
        shape_vars_found = []
        for m in shape_var_pattern.finditer(region):
            sv = m.group(1)
            if sv not in shape_vars_found and sv not in ('shape', 'selfShape'):
                shape_vars_found.append(sv)
        
        if shape_vars_found:
            # Create totalSize by summing GetShapeSize of each shape variable
            total_expr = " + ".join(["GetShapeSize(" + sv + ")" for sv in shape_vars_found])
            total_check = f"""
{tensor_indent}// === Tensor Size distribution check ===
{tensor_indent}long long totalSize = {total_expr};
{tensor_indent}if (totalSize < totalSizeMin || totalSize > totalSizeMax) {{ goto ERROR; }}
"""
            modified = modified[:tensor_creation_pos] + total_check + modified[tensor_creation_pos:]
            changes.append(f"  added totalSize check with GetShapeSize calls: {shape_vars_found}")
            return modified, changes
        
        changes.append("  WARNING: no shape_size vars found, skipping totalSize check")
        return modified, changes
    
    # Build the totalSize check
    # Sum all shape_size variables
    total_expr = " + ".join(shape_size_vars)
    total_check = f"""
{tensor_indent}// === Tensor Size distribution check ===
{tensor_indent}long long totalSize = {total_expr};
{tensor_indent}if (totalSize < totalSizeMin || totalSize > totalSizeMax) {{ goto ERROR; }}
"""
    
    # Insert before tensor creation
    modified = modified[:tensor_creation_pos] + total_check + modified[tensor_creation_pos:]
    changes.append(f"  added totalSize check with vars: {shape_size_vars}")
    
    return modified, changes


def process_file(filepath):
    """Process a single CPP file."""
    try:
        with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
            content = f.read()
    except Exception as e:
        return False, [f"  ERROR reading file: {e}"]
    
    original = content
    all_changes = []
    
    # Step 1: shape=[] support (dim_num 1,8 -> 0,8)
    content, changes1 = modify_dim_num_zero(content)
    all_changes.extend(changes1)
    
    # Step 2: Size distribution
    content, changes2 = add_size_distribution(content, os.path.basename(filepath))
    all_changes.extend(changes2)
    
    # Only write if content changed
    if content != original:
        try:
            with open(filepath, 'w', encoding='utf-8', errors='replace') as f:
                f.write(content)
            return True, all_changes
        except Exception as e:
            return False, all_changes + [f"  ERROR writing file: {e}"]
    else:
        return False, ["  no changes needed"]


def main():
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    op_dir = os.path.join(base_dir, "turing_sec", "aclnnAPI_fuzz", "src", "op")
    nc_dir = os.path.join(base_dir, "turing_sec", "aclnnAPI_fuzz", "src", "src_no_contiguous")
    
    # Get the list of non-contiguous files (this is the reference list)
    nc_files = []
    if os.path.isdir(nc_dir):
        for f in os.listdir(nc_dir):
            if f.endswith(".cpp"):
                api_name = get_api_name(f)
                if not should_exclude(api_name):
                    nc_files.append(f)
    
    print(f"Non-contiguous files to process: {len(nc_files)}")
    
    # Build the list of (directory, filename) pairs to process
    # 1. All nc files in src_no_contiguous/
    # 2. Corresponding op files in op/ (matching by api name)
    
    files_to_process = []
    
    # Add nc files
    for nc_file in nc_files:
        nc_path = os.path.join(nc_dir, nc_file)
        files_to_process.append(("nc", nc_path, nc_file))
    
    # Add corresponding op files
    op_files_processed = set()
    for nc_file in nc_files:
        api_name = get_api_name(nc_file)
        op_file = api_name + ".cpp"
        op_path = os.path.join(op_dir, op_file)
        if os.path.isfile(op_path) and op_file not in op_files_processed:
            files_to_process.append(("op", op_path, op_file))
            op_files_processed.add(op_file)
    
    print(f"Total files to process: {len(files_to_process)}")
    print(f"  nc files: {sum(1 for f in files_to_process if f[0] == 'nc')}")
    print(f"  op files: {sum(1 for f in files_to_process if f[0] == 'op')}")
    
    # Process files
    modified_count = 0
    skipped_count = 0
    error_count = 0
    
    for dir_type, filepath, filename in files_to_process:
        try:
            success, changes = process_file(filepath)
            if success:
                modified_count += 1
                if dir_type == "nc":
                    print(f"[MODIFIED] nc/{filename}")
                else:
                    print(f"[MODIFIED] op/{filename}")
                for ch in changes:
                    print(f"  {ch}")
            else:
                skipped_count += 1
                if any("ERROR" in ch for ch in changes):
                    error_count += 1
                    print(f"[ERROR] {dir_type}/{filename}")
                    for ch in changes:
                        print(f"  {ch}")
        except Exception as e:
            error_count += 1
            print(f"[ERROR] {dir_type}/{filename}: {e}")
    
    print(f"\n=== Summary ===")
    print(f"Modified: {modified_count}")
    print(f"Skipped: {skipped_count}")
    print(f"Errors: {error_count}")


if __name__ == "__main__":
    main()