import json
import os
from tqdm import tqdm

# ==========================================
# 1. 配置
# ==========================================
CACHE_FILE = "/root/.cache/huggingface/datasets/downloads/9ea09f0bc4302fcbd98eaa4a315758ff0c3b8d2f7566170d1450868e29429b2b"
OUTPUT_DIR = "shell/playground/data/hf_data_0330"
os.makedirs(OUTPUT_DIR, exist_ok=True)
OUTPUT_PATH = os.path.join(OUTPUT_DIR, "Step-3.5-Flash-SFT_train_with_reasoning.jsonl")

ROLE_MAP = {
    "user": "human",
    "human": "human",
    "assistant": "gpt",
    "gpt": "gpt",
    "system": "system",
}

def main():
    print(f"📦 检测到大型 JSON 数组格式，正在加载文件 (可能需要几十秒)...")
    
    try:
        # 直接一次性读取整个 JSON 数组
        with open(CACHE_FILE, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)
    except Exception as e:
        print(f"❌ 加载失败: {e}")
        return

    print(f"✅ 加载成功，共计 {len(raw_data)} 条原始记录。开始转换...")

    success_count = 0
    
    with open(OUTPUT_PATH, 'w', encoding='utf-8') as f_out:
        # 此时 raw_data 是一个 list，可以正常 tqdm 迭代
        for i, entry in enumerate(tqdm(raw_data, desc="Processing")):
            
            # 1. 自动兼容字段名: conversations 或 messages
            msgs = entry.get("conversations") or entry.get("messages")
            
            if not isinstance(msgs, list):
                continue

            # 2. 转换格式 & 过滤
            new_convs = []
            is_agent = False
            
            for m in msgs:
                if not isinstance(m, dict): continue
                
                # 过滤工具调用数据
                if m.get("tool_calls") or m.get("name") or m.get("tool_call_id"):
                    is_agent = True
                    break
                
                role = m.get("role", "").lower()
                content = m.get("content", "")
                reasoning_content = m.get("reasoning_content", "")

                if reasoning_content:
                    content = f"{reasoning_content}\n\n{content}"
                
                # 处理 content 可能是 list 的情况（多模态或复杂格式）
                if isinstance(content, list):
                    # 取出文本部分
                    content = " ".join([c.get("text", "") for c in content if isinstance(c, dict) and "text" in c])
                
                if role in ROLE_MAP and isinstance(content, str) and content.strip():
                    new_convs.append({
                        "from": ROLE_MAP[role],
                        "value": content
                    })
            
            if is_agent or not new_convs:
                continue

            # 3. 写入 JSONL 格式 (每行一个对象)
            output_item = {
                "id": f"step_{i}",
                "conversations": new_convs
            }
            # 保留 source 字段（如果原数据有）
            if 'source' in entry:
                output_item['source'] = entry['source']
                
            f_out.write(json.dumps(output_item, ensure_ascii=False) + '\n')
            success_count += 1

    print(f"\n✨ 处理完成！")
    print(f"📊 成功提取: {success_count} 条")
    print(f"📂 结果保存至: {OUTPUT_PATH}")

if __name__ == "__main__":
    main()