#!/usr/bin/env python3
"""
Fix files where totalSize check was incorrectly inserted inside a lambda function.
These files use lambdas for tensor creation (auto mk = [&](...){ ... return CreateAclTensorByScene(...); })
and the totalSize check was inserted inside the lambda body, breaking the return statement.

Fix strategy:
1. Remove the incorrectly placed totalSize check block
2. Find the correct insertion point: after the lambda definition, before the first lambda call
3. Re-insert the totalSize check at the correct location
"""

import os
import re
import sys

def fix_lambda_files(filepath):
    """Fix a single file with lambda-based tensor creation."""
    with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
        content = f.read()
    
    original = content
    changes = []
    
    # Step 1: Remove the incorrectly placed totalSize check block
    # Pattern: the block was inserted after 'return' inside a lambda
    # It looks like:
    #   return
    #  // === Tensor Size distribution check ===
    #  long long totalSize = ...;
    #  if (totalSize < totalSizeMin || totalSize > totalSizeMax) { goto ERROR; }
    
    bad_block_pattern = re.compile(
        r'(return\s*)\n\s*// === Tensor Size distribution check ===\s*\n'
        r'\s*long long totalSize = [^;]+;\s*\n'
        r'\s*if \(totalSize < totalSizeMin \|\| totalSize > totalSizeMax\) \{ goto ERROR; \}\s*\n',
        re.MULTILINE
    )
    
    # Also handle the case where there's a newline between return and the check
    bad_block_pattern2 = re.compile(
        r'(return\s*)\n\s*\n\s*// === Tensor Size distribution check ===\s*\n'
        r'\s*long long totalSize = [^;]+;\s*\n'
        r'\s*if \(totalSize < totalSizeMin \|\| totalSize > totalSizeMax\) \{ goto ERROR; \}\s*\n',
        re.MULTILINE
    )
    
    # Extract the totalSize expression before removing
    total_expr = None
    m = bad_block_pattern.search(content)
    if not m:
        m = bad_block_pattern2.search(content)
    
    if m:
        # Extract the totalSize expression
        expr_match = re.search(r'long long totalSize = ([^;]+);', m.group(0))
        if expr_match:
            total_expr = expr_match.group(1).strip()
        # Remove the bad block, keep just 'return'
        content = bad_block_pattern.sub(r'\1\n', content)
        content = bad_block_pattern2.sub(r'\1\n', content)
        changes.append("  removed incorrectly placed totalSize check from lambda")
    else:
        # Check if there are other patterns of bad insertion
        # Maybe the return is on the same line
        bad_block_pattern3 = re.compile(
            r'(return\s*)// === Tensor Size distribution check ===\s*\n'
            r'\s*long long totalSize = [^;]+;\s*\n'
            r'\s*if \(totalSize < totalSizeMin \|\| totalSize > totalSizeMax\) \{ goto ERROR; \}',
            re.MULTILINE
        )
        m3 = bad_block_pattern3.search(content)
        if m3:
            expr_match = re.search(r'long long totalSize = ([^;]+);', m3.group(0))
            if expr_match:
                total_expr = expr_match.group(1).strip()
            content = bad_block_pattern3.sub(r'\1', content)
            changes.append("  removed incorrectly placed totalSize check from lambda (pattern3)")
        else:
            changes.append("  WARNING: could not find bad block to remove")
            return content, changes
    
    if not total_expr:
        changes.append("  WARNING: could not extract totalSize expression")
        return content, changes
    
    # Step 2: Find the correct insertion point
    # The lambda definition ends with '};' (semicolon after closing brace)
    # After that, the first lambda call is typically: ret=mk(... or similar
    
    # Find the lambda definition pattern: auto <name> = [&](...) -> int { ... };
    # The lambda ends with '};' at the start of a line (after indentation)
    lambda_end_pattern = re.compile(r'^(\s*)\};\s*\n', re.MULTILINE)
    
    # Find all lambda ends and look for the one after which tensor calls begin
    insertion_pos = None
    insertion_indent = "            "
    
    for lm in lambda_end_pattern.finditer(content):
        pos = lm.end()
        # Check if the next non-empty line is a lambda call (contains 'mk(' or 'ret=' or similar)
        next_lines = content[pos:pos+200]
        if re.search(r'(ret\s*=|mk\s*\(|makeTensor|createTensor)', next_lines):
            insertion_pos = pos
            insertion_indent = lm.group(1)
            break
    
    # Alternative: find the first call to a lambda function after the try block
    if insertion_pos is None:
        # Look for patterns like: ret=mk(&... or ret = mk(&...
        lambda_call_pattern = re.compile(r'(\s*)(ret\s*=\s*\w+\s*\()', re.MULTILINE)
        # Find the first one after the sizeBucket code
        sb_pos = content.find('sizeBucket')
        if sb_pos >= 0:
            lc_match = lambda_call_pattern.search(content, sb_pos)
            if lc_match:
                insertion_pos = lc_match.start()
                insertion_indent = lc_match.group(1)
    
    if insertion_pos is None:
        changes.append("  WARNING: could not find correct insertion point for totalSize check")
        return content, changes
    
    # Step 3: Insert the totalSize check at the correct position
    total_check = f"""{insertion_indent}// === Tensor Size distribution check ===
{insertion_indent}long long totalSize = {total_expr};
{insertion_indent}if (totalSize < totalSizeMin || totalSize > totalSizeMax) {{ goto ERROR; }}
"""
    
    content = content[:insertion_pos] + total_check + content[insertion_pos:]
    changes.append(f"  re-inserted totalSize check before lambda call with expr: {total_expr}")
    
    return content, changes


def main():
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    
    dirs = [
        os.path.join(base_dir, "turing_sec", "aclnnAPI_fuzz", "src", "op"),
        os.path.join(base_dir, "turing_sec", "aclnnAPI_fuzz", "src", "src_no_contiguous"),
    ]
    
    fixed_count = 0
    error_count = 0
    
    for d in dirs:
        for f in sorted(os.listdir(d)):
            if not f.endswith(".cpp"):
                continue
            filepath = os.path.join(d, f)
            with open(filepath, 'r', encoding='utf-8', errors='replace') as fh:
                content = fh.read()
            
            # Check if this file has the lambda issue
            if re.search(r'return\s*\n\s*// === Tensor Size distribution check ===', content) or \
               re.search(r'return\s*// === Tensor Size distribution check ===', content):
                
                fixed_content, changes = fix_lambda_files(filepath)
                
                if fixed_content != content:
                    with open(filepath, 'w', encoding='utf-8', errors='replace') as fh:
                        fh.write(fixed_content)
                    fixed_count += 1
                    print("[FIXED] " + d.split('/')[-1] + "/" + f)
                    for ch in changes:
                        print(ch)
                else:
                    error_count += 1
                    print("[ERROR] " + d.split('/')[-1] + "/" + f)
                    for ch in changes:
                        print(ch)
    
    print()
    print("=== Summary ===")
    print("Fixed: " + str(fixed_count))
    print("Errors: " + str(error_count))


if __name__ == "__main__":
    main()