import json
import os
import glob
from tqdm import tqdm
from huggingface_hub import snapshot_download

# ==========================================
# 1. 配置
# ==========================================
DATASET_ID = "stepfun-ai/Step-3.5-Flash-SFT"

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.jsonl")

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

def main():
    print(f"📦 正在同步 '{DATASET_ID}' 的所有原始文件 (绕过 Arrow 解析) ...")
    print(f"⏳ 会自动下载所有分片到缓存，请耐心等待...")
    
    try:
        # 1. 直接下载仓库中的所有 JSON 数据文件，不走 Schema 校验
        local_dir = snapshot_download(
            repo_id=DATASET_ID, 
            repo_type="dataset",
            allow_patterns=["*.json", "*.jsonl"]
        )
    except Exception as e:
        print(f"❌ 下载失败: {e}")
        return

    # 2. 找到所有下载下来的 json/jsonl 文件
    data_files = glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True) + \
                 glob.glob(os.path.join(local_dir, "**", "*.jsonl"), recursive=True)
                 
    # 过滤掉明显的非数据文件（比如 huggingface 自动生成的元数据）
    data_files = [f for f in data_files if "dataset_info" not in f.lower()]

    print(f"✅ 成功定位到 {len(data_files)} 个数据文件。开始逐个文件解析...")

    success_count = 0
    total_records = 0
    
    with open(OUTPUT_PATH, 'w', encoding='utf-8') as f_out:
        # 以文件为单位进行进度条显示
        for file_path in tqdm(data_files, desc="Processing Files"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                    if not content:
                        continue
                    
                    # 兼容处理：文件内容可能是完整的 JSON Array，也可能是 JSONL
                    if content.startswith('['):
                        records = json.loads(content)
                    else:
                        records = [json.loads(line) for line in content.split('\n') if line.strip()]
            except Exception as e:
                print(f"\n⚠️ 跳过文件 {os.path.basename(file_path)}: 无法解析 JSON ({e})")
                continue

            # 3. 处理单文件内的所有记录
            for entry in records:
                total_records += 1
                
                msgs = entry.get("conversations") or entry.get("messages")
                
                # 如果这个脏数据的 conversations 不是 list (正是由于这个导致 pyarrow 崩溃)，直接跳过
                if not isinstance(msgs, list):
                    continue

                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", "")
                    if role:
                        role = role.lower()
                    content = m.get("content", "")
                    
                    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

                # 4. 写入输出文件
                output_item = {
                    "id": f"step_{total_records}",
                    "conversations": new_convs
                }
                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"📊 共扫描原始记录: {total_records} 条")
    print(f"🎯 成功提取可用记录: {success_count} 条")
    print(f"📂 结果保存至: {OUTPUT_PATH}")

if __name__ == "__main__":
    main()