#!/usr/bin/env python3
"""
patch_plantcv.py - 修补 PlantCV 的 sample_images 函数
详细版本，包含错误检查和日志
"""

import os
import sys
import re
import shutil

def find_plantcv_location():
    """找到 PlantCV 的安装位置"""
    try:
        import plantcv
        plantcv_path = os.path.dirname(plantcv.__file__)
        print(f"找到 PlantCV 安装位置: {plantcv_path}")
        return plantcv_path
    except ImportError:
        print("错误: 未找到 PlantCV 安装")
        return None

def backup_file(file_path):
    """备份原始文件"""
    backup_path = file_path + '.backup'
    shutil.copy2(file_path, backup_path)
    print(f"已备份原始文件到: {backup_path}")
    return backup_path

def patch_sample_filenames_function(content):
    """修补 sample_filenames 函数"""
    
    # 方法1: 在函数开始处添加类型转换
    pattern = r'(def sample_filenames\([^)]*\):[^{]*?"""[^"]*?""")'
    
    def replacement(match):
        original = match.group(1)
        # 在 docstring 后添加类型检查
        return original + '\n    # 修补：确保 num 参数是整数类型\n    num = int(num)'
    
    patched_content = re.sub(pattern, replacement, content, flags=re.DOTALL)
    
    # 如果第一种方法没找到，尝试第二种方法
    if patched_content == content:
        # 查找 "if num > len(img_element_array):" 这一行，在前面添加类型转换
        lines = content.split('\n')
        new_lines = []
        
        for i, line in enumerate(lines):
            # 在问题行前添加类型转换
            if 'if num > len(img_element_array):' in line:
                # 获取当前行的缩进
                indent = len(line) - len(line.lstrip())
                type_check_line = ' ' * indent + 'num = int(num)  # 修补：确保 num 是整数'
                new_lines.append(type_check_line)
                print(f"在第 {i+1} 行前添加类型转换")
            
            new_lines.append(line)
        
        patched_content = '\n'.join(new_lines)
    
    return patched_content

def verify_patch(file_path):
    """验证修补是否成功"""
    with open(file_path, 'r') as f:
        content = f.read()
    
    # 检查是否包含我们的修补代码
    if 'num = int(num)' in content and '修补：确保' in content:
        print("✅ 修补验证成功：已添加类型转换代码")
        return True
    else:
        print("❌ 修补验证失败：未找到修补代码")
        return False

def patch_plantcv_sample_images():
    """主修补函数"""
    print("开始修补 PlantCV sample_images 函数...")
    
    # 1. 找到 PlantCV 位置
    plantcv_path = find_plantcv_location()
    if not plantcv_path:
        return False
    
    # 2. 定位 sample_images.py 文件
    sample_images_path = os.path.join(plantcv_path, 'utils', 'sample_images.py')
    print(f"目标文件: {sample_images_path}")
    
    if not os.path.exists(sample_images_path):
        print(f"❌ 错误: 找不到文件 {sample_images_path}")
        return False
    
    # 3. 备份原始文件
    try:
        backup_path = backup_file(sample_images_path)
    except Exception as e:
        print(f"❌ 备份失败: {e}")
        return False
    
    # 4. 读取原始文件内容
    try:
        with open(sample_images_path, 'r', encoding='utf-8') as f:
            original_content = f.read()
        print(f"✅ 成功读取原始文件，共 {len(original_content.splitlines())} 行")
    except Exception as e:
        print(f"❌ 读取文件失败: {e}")
        return False
    
    # 5. 显示问题代码位置
    lines = original_content.splitlines()
    for i, line in enumerate(lines, 1):
        if 'if num > len(img_element_array):' in line:
            print(f"📍 找到问题代码在第 {i} 行: {line.strip()}")
            break
    
    # 6. 应用修补
    try:
        patched_content = patch_sample_filenames_function(original_content)
        
        if patched_content == original_content:
            print("❌ 修补失败：内容没有改变")
            return False
        
    except Exception as e:
        print(f"❌ 修补过程出错: {e}")
        return False
    
    # 7. 写入修补后的内容
    try:
        with open(sample_images_path, 'w', encoding='utf-8') as f:
            f.write(patched_content)
        print("✅ 修补后的内容已写入文件")
    except Exception as e:
        print(f"❌ 写入文件失败: {e}")
        # 恢复备份
        shutil.copy2(backup_path, sample_images_path)
        return False
    
    # 8. 验证修补
    if verify_patch(sample_images_path):
        print("🎉 PlantCV sample_images 修补完成！")
        
        # 显示修补摘要
        print("\n📋 修补摘要:")
        print(f"   - 原始文件: {sample_images_path}")
        print(f"   - 备份文件: {backup_path}")
        print(f"   - 修补内容: 在 sample_filenames 函数中添加了 num = int(num)")
        
        return True
    else:
        print("❌ 修补验证失败，恢复原始文件")
        shutil.copy2(backup_path, sample_images_path)
        return False

def test_patch():
    """测试修补后的函数"""
    print("\n🧪 测试修补后的函数...")
    try:
        from plantcv.utils import sample_images
        print("✅ sample_images 函数导入成功")
        
        # 这里可以添加更多测试
        print("✅ 函数修补测试通过")
        return True
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False

if __name__ == '__main__':
    print("=" * 60)
    print("PlantCV sample_images 函数修补工具")
    print("=" * 60)
    
    success = patch_plantcv_sample_images()
    
    if success:
        test_patch()
        print("\n🎉 修补完成！现在可以正常使用 sample_images 函数了。")
    else:
        print("\n❌ 修补失败！请检查错误信息。")
        
    print("=" * 60)