import json
import os
from datasets import load_dataset
from tqdm import tqdm

# ==========================================
# 1. 全局配置与常量
# ==========================================
# 默认输出目录，可根据需要修改为 "data_hf"
OUTPUT_DIR = "shell/playground/data/hf_data_0330"
# 自动创建输出目录，防止报错
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ==========================================
# 2. 数据格式转换策略
# ==========================================
def process_direct_conversations(conversations):
    """处理已经是 from/value 格式的 conversations"""
    valid_convs = []
    for c in conversations:
        val = c.get("value", "")
        if val is not None and val != "":
            valid_convs.append({
                "from": c.get("from"),
                "value": val
            })
    return valid_convs

# ==========================================
# 3. 核心写入引擎
# ==========================================
def process_and_save(dataset_iterator, output_filename, desc, convert_func):
    """统一的数据遍历、转换与保存逻辑"""
    output_path = os.path.join(OUTPUT_DIR, output_filename)
    print(f"\n🚀 开始处理: {desc}")
    
    ix = 0
    with open(output_path, 'w', encoding='utf-8') as f:
        for item in tqdm(dataset_iterator, desc=desc):
            conversations = convert_func(item)
            
            # 跳过无效或空的对话
            if not conversations:
                continue
                
            conv = {
                'id': ix,
                'conversations': conversations
            }
            
            # 兼容 OpenCoder 的 source 字段需求
            if 'source' in item:
                conv['source'] = item['source']
                
            f.write(json.dumps(conv, ensure_ascii=False) + '\n')
            ix += 1
            
    print(f"✅ [{desc}] 转换完成! 总计写入 {ix} 条，保存至 {output_path}")

# ==========================================
# 4. 主执行流
# ==========================================
def main():
    # OpenDCAI/dataflow-instruct-10k
    # 数据集结构中包含 'conversations' 字段，其值已经是 from 和 value 的格式
    process_and_save(
        dataset_iterator=load_dataset("OpenDCAI/dataflow-instruct-10k", split="train"),
        output_filename="dataflow-instruct-10k_train.jsonl",
        desc="OpenDCAI/dataflow-instruct-10k",
        convert_func=lambda x: process_direct_conversations(x['conversations'])
    )
    
    print("\n🎉 数据集下载与转换任务已全部完成！")

if __name__ == "__main__":
    main()