#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ARTv2 代码风格批量修复工具
自动修复 == True/False 比较问题
"""

import os
import re

def fix_boolean_comparison(content):
    """修复 == True/False 比较"""
    
    # Pattern 1: if xxx == True: -> if xxx:
    # 匹配简单变量和方法调用
    content = re.sub(
        r'(\s+)if\s+([a-zA-Z_][a-zA-Z0-9_\.]*(?:\([^)]*\))?)\s+==\s+True\s*:',
        r'\1if \2:',
        content
    )
    
    # Pattern 2: if xxx == False: -> if not xxx:
    content = re.sub(
        r'(\s+)if\s+([a-zA-Z_][a-zA-Z0-9_\.]*(?:\([^)]*\))?)\s+==\s+False\s*:',
        r'\1if not \2:',
        content
    )
    
    # Pattern 3: while xxx == True: -> while xxx:
    content = re.sub(
        r'(\s+)while\s+([a-zA-Z_][a-zA-Z0-9_\.]*(?:\([^)]*\))?)\s+==\s+True\s*:',
        r'\1while \2:',
        content
    )
    
    # Pattern 4: while xxx == False: -> while not xxx:
    content = re.sub(
        r'(\s+)while\s+([a-zA-Z_][a-zA-Z0-9_\.]*(?:\([^)]*\))?)\s+==\s+False\s*:',
        r'\1while not \2:',
        content
    )
    
    # Pattern 5: 处理复杂表达式，如 cmds.getAttr(...) == True
    content = re.sub(
        r'(\s+)if\s+(cmds\.[a-zA-Z_]+\([^)]+\))\s+==\s+True\s*:',
        r'\1if \2:',
        content
    )
    
    content = re.sub(
        r'(\s+)if\s+(cmds\.[a-zA-Z_]+\([^)]+\))\s+==\s+False\s*:',
        r'\1if not \2:',
        content
    )
    
    return content

def process_file(filepath):
    """处理单个文件"""
    try:
        print(f"Processing: {filepath}")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            original = f.read()
        
        # 应用修复
        fixed = fix_boolean_comparison(original)
        
        # 检查是否有变化
        if fixed != original:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(fixed)
            
            # 统计变化
            original_lines = original.split('\n')
            fixed_lines = fixed.split('\n')
            changes = sum(1 for o, f in zip(original_lines, fixed_lines) if o != f)
            print(f"  ✓ Fixed {changes} lines")
            return changes
        else:
            print(f"  - No changes needed")
            return 0
    except Exception as e:
        print(f"  ✗ Error: {e}")
        return 0

def main():
    """主函数"""
    base_path = r"h:\Workspace\Raw\Tools\Plugins\Maya\plug-ins\ARTv2\Core\Scripts"
    
    # 高优先级文件
    files_to_fix = [
        os.path.join(base_path, "RigModules", "ART_Leg_Standard.py"),
        os.path.join(base_path, "RigModules", "ART_Arm_Standard.py"),
    ]
    
    print("=" * 60)
    print("ARTv2 代码风格批量修复")
    print("=" * 60)
    print(f"\n将修复 {len(files_to_fix)} 个文件\n")
    
    total_changes = 0
    for filepath in files_to_fix:
        if os.path.exists(filepath):
            changes = process_file(filepath)
            total_changes += changes
        else:
            print(f"✗ File not found: {filepath}")
    
    print("\n" + "=" * 60)
    print(f"完成！共修复 {total_changes} 行代码")
    print("=" * 60)

if __name__ == "__main__":
    main()