#!/usr/bin/env python3
import json
import os
from pathlib import Path

def process_file(input_file, output_dir):
    """处理单个文件，去掉system prompt"""
    print(f"处理文件: {input_file}")
    
    # 读取输入文件
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    output_data = {
        "type": "conversation",
        "instances": []
    }
    
    # 处理每个instance的messages
    for instance in data['instances']:
        # 过滤掉role为system的消息
        system_prompt = instance['messages'][0]['content']
        if system_prompt == "detailed thinking on":
            # pop这个instance，continue
            continue

        instance['messages'] = instance['messages'][1:]
        # instance['messages'] = [
        #     msg for msg in instance['messages'] 
        #     if msg.get('role') != 'system'
        # ]
        instance.pop('is_reasoning', None)
        instance.pop('dataset', None)
        output_data['instances'].append(instance)
    
    # 生成输出文件名
    output_filename = os.path.basename(input_file)
    output_path = os.path.join(output_dir, output_filename)
    
    # 保存处理后的文件
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, ensure_ascii=False, indent=2)
    
    print(f"已保存到: {output_path}")

def main():
    # 输入和输出目录
    input_dir = "data/nemotron_subset"
    output_dir = "data/math_sft"
    
    # 要处理的文件范围
    start_num = 175
    end_num = 215
    
    # 处理每个文件
    for i in range(start_num, end_num + 1):
        filename = f"converted_conversations_{i:04d}.json"
        input_file = os.path.join(input_dir, filename)
        
        if os.path.exists(input_file):
            try:
                process_file(input_file, output_dir)
            except Exception as e:
                print(f"处理文件 {filename} 时出错: {e}")
        else:
            print(f"文件不存在: {input_file}")
    
    print("所有文件处理完成！")

if __name__ == "__main__":
    main()
