#!/usr/bin/env python3
import os
import re

def add_initialize_function(file_path):
    """为协议文件添加initialize函数"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 检查是否已经有initialize函数
    if 'fn initialize()' in content:
        return False
    
    # 找到impl块中hash函数之后的位置
    # 在闭合}之前插入新函数
    pattern = r'(    pub fn hash\(data: &\[u8\]\) -> Hash \{[^}]+\})\n(\})\n\n(impl Default)'
    
    initialize_func = '''
    
    /// 初始化协议实例（带错误处理）
    /// 
    /// # Returns
    /// 
    /// 返回Result包装的协议实例
    pub fn initialize() -> Result<Self> {
        Ok(Self::new())
    }'''
    
    replacement = r'\1' + initialize_func + r'\n\2\n\n\3'
    
    new_content = re.sub(pattern, replacement, content)
    
    if new_content != content:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(new_content)
        return True
    
    return False

def process_all_protocols(base_dir):
    """处理所有协议文件"""
    modified_count = 0
    
    for layer in range(10):  # Layer 0-9
        layer_dir = os.path.join(base_dir, f'layer{layer}')
        if not os.path.exists(layer_dir):
            continue
        
        for file_name in os.listdir(layer_dir):
            if file_name.endswith('.rs') and file_name != 'mod.rs':
                file_path = os.path.join(layer_dir, file_name)
                if add_initialize_function(file_path):
                    modified_count += 1
                    print(f'✅ 修复: {file_path}')
    
    return modified_count

if __name__ == '__main__':
    base_dir = '/home/ubuntu/NAC_Clean_Dev/nac-protocols/src'
    count = process_all_protocols(base_dir)
    print(f'\n总共修复了 {count} 个协议文件')