import json
import argparse
import os
from datetime import datetime

def clean_null_answers(input_file, output_file=None):
    """
    清理JSON文件中answer为null的数据条目
    
    参数:
    input_file (str): 输入JSON文件路径
    output_file (str, optional): 输出JSON文件路径，默认为在原文件名基础上添加_cleaned后缀
    
    返回:
    int: 被删除的条目数量
    """
    if output_file is None:
        # 如果没有指定输出文件，则在原文件名基础上添加_cleaned后缀
        filename, ext = os.path.splitext(input_file)
        output_file = f"{filename}_cleaned{ext}"
    
    # 创建备份文件
    backup_file = f"{input_file}.bak.{datetime.now().strftime('%Y%m%d%H%M%S')}"
    
    try:
        # 读取输入文件
        with open(input_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 创建备份
        with open(backup_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        print(f"已创建备份文件: {backup_file}")
        print(f"原始数据条目数: {len(data)}")
        
        # 过滤掉answer为null的数据
        original_count = len(data)
        filtered_data = [item for item in data if item.get("answer") is not None]
        removed_count = original_count - len(filtered_data)
        
        # 保存处理后的数据
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(filtered_data, f, ensure_ascii=False, indent=2)
        
        print(f"处理完成！删除了 {removed_count} 条空值答案数据")
        print(f"保留的数据条目数: {len(filtered_data)}")
        print(f"清理后的数据已保存至: {output_file}")
        
        return removed_count
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return -1

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="清理JSON文件中answer为null的数据")
    parser.add_argument("--input", "-i", type=str, required=True, help="输入JSON文件路径")
    parser.add_argument("--output", "-o", type=str, help="输出JSON文件路径（可选）")
    
    args = parser.parse_args()
    
    clean_null_answers(args.input, args.output)

# python clean_null_answers.py --input ./generated_problems.json --output ./cleaned_problems.json