#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import datasets
import json
from pathlib import Path
import os

def convert_ds_to_json(dataset_path, output_file, max_samples=None, split_files=False, instances_per_file=1000, output_dir=None, max_file_size_mb=None):
    """
    将Hugging Face数据集转换为指定的JSON格式
    
    Args:
        dataset_path: 数据集路径
        output_file: 输出的JSON文件路径（如果不拆分）或基础文件名（如果拆分）
        max_samples: 最大处理样本数，None表示处理所有样本
        split_files: 是否拆分成多个文件
        instances_per_file: 每个文件包含的对话实例数量（当max_file_size_mb为None时使用）
        output_dir: 输出目录（拆分文件时使用）
        max_file_size_mb: 每个文件的最大大小（MB），如果设置了此参数，会按文件大小拆分
    """
    
    print(f"开始加载数据集: {dataset_path}")
    
    try:
        # 加载数据集
        ds = datasets.load_from_disk(dataset_path)
        
        print(f"数据集加载成功，共有 {len(ds):,} 个样本")
        
        # 如果需要拆分文件，设置输出目录
        if split_files:
            if output_dir is None:
                output_dir = "split_conversations"
            output_dir_path = Path(output_dir)
            output_dir_path.mkdir(exist_ok=True)
            print(f"输出目录: {output_dir_path}")
        else:
            output_path = Path(output_file)
        
        if max_samples:
            print(f"最多处理 {max_samples:,} 个样本")
        if split_files:
            if max_file_size_mb:
                print(f"将按文件大小拆分，每个文件最大 {max_file_size_mb} MB")
            else:
                print(f"将拆分成多个文件，每个文件包含 {instances_per_file:,} 个对话实例")
        
        processed_count = 0
        current_instances = []
        file_index = 1
        
        # 遍历数据集
        for i, sample in enumerate(ds):
            # 如果设置了最大样本数限制，检查是否达到
            if max_samples and i >= max_samples:
                break
            
            try:
                # 检查是否包含messages字段
                if 'messages' not in sample:
                    continue
                
                messages = sample['messages']
                
                # 验证messages格式
                if not messages or not isinstance(messages, list) or len(messages) != 2:
                    continue
                
                # 验证每个message的格式
                valid_messages = []
                for msg in messages:
                    if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                        # 确保role是有效的值
                        if msg['role'] in ['assistant', 'user']:
                            valid_messages.append({
                                "role": msg['role'],
                                "content": msg['content']
                            })
                
                # 如果有有效的messages，添加到结果中
                if valid_messages:
                    instance = {
                        "messages": valid_messages
                    }
                    
                    # 可选：添加其他元数据
                    # for key in sample.keys():
                    #     if key != 'messages':
                    #         instance[key] = sample[key]
                    
                    current_instances.append(instance)
                    processed_count += 1
                    
                    # 如果需要拆分文件，检查是否需要保存当前文件
                    if split_files:
                        should_save = False
                        
                        if max_file_size_mb:
                            # 按文件大小拆分：估算当前文件大小
                            estimated_size = len(json.dumps({"type": "conversation", "instances": current_instances}, ensure_ascii=False)) / (1024 * 1024)
                            if estimated_size >= max_file_size_mb:
                                should_save = True
                        else:
                            # 按实例数量拆分
                            if len(current_instances) >= instances_per_file:
                                should_save = True
                        
                        if should_save:
                            _save_split_file(current_instances, output_dir_path, file_index, output_file)
                            current_instances = []
                            file_index += 1
            
            except Exception as e:
                print(f"第 {i+1} 个样本处理错误: {e}")
                continue
            
            # 每处理10000个样本显示一次进度
            if (i + 1) % 10000 == 0:
                print(f"已处理 {i+1:,} 个样本，转换了 {processed_count:,} 个对话")
        
        print(f"处理完成！总共处理了 {i+1:,} 个样本，转换了 {processed_count:,} 个对话")
        
        # 保存结果
        if split_files:
            # 保存剩余的实例
            if current_instances:
                _save_split_file(current_instances, output_dir_path, file_index, output_file)
            
            print(f"已拆分成 {file_index} 个文件，保存在 {output_dir_path} 目录中")
            
            # 计算总文件大小
            total_size = sum(f.stat().st_size for f in output_dir_path.glob("*.json"))
            print(f"总文件大小: {total_size / (1024**2):.2f} MB")
        else:
            # 保存为单个文件
            result = {
                "type": "conversation",
                "instances": current_instances
            }
            
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            
            print(f"结果已保存到 {output_path}")
            print(f"输出文件大小: {output_path.stat().st_size / (1024**2):.2f} MB")
        
    except Exception as e:
        print(f"转换过程中发生错误: {e}")

def _save_split_file(instances, output_dir_path, file_index, base_filename):
    """
    保存拆分的文件
    """
    # 创建文件名
    base_name = Path(base_filename).stem
    filename = f"{base_name}_{file_index:04d}.json"
    filepath = output_dir_path / filename
    
    # 创建JSON结构
    result = {
        "type": "conversation",
        "instances": instances
    }
    
    # 保存文件
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    file_size = filepath.stat().st_size / (1024**2)
    print(f"已保存文件 {filename} ({len(instances):,} 个对话实例, {file_size:.2f} MB)")

def preview_dataset(dataset_path, num_samples=5):
    """
    预览数据集转换结果
    """
    print("=== 数据集预览 ===")
    
    try:
        ds = datasets.load_from_disk(dataset_path)
        print(f"数据集总样本数: {len(ds):,}")
        
        # 显示数据集的列信息
        print(f"数据集列: {list(ds.features.keys())}")
        
        samples = []
        for i, sample in enumerate(ds):
            if i >= num_samples:
                break
            
            if 'messages' in sample and sample['messages']:
                # 截断长文本用于预览
                messages = []
                for msg in sample['messages']:
                    if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                        content = msg['content']
                        if len(content) > 200:
                            content = content[:200] + "..."
                        messages.append({
                            "role": msg['role'],
                            "content": content
                        })
                
                instance = {"messages": messages}
                # 添加其他字段（如果有的话）
                for key in sample.keys():
                    if key != 'messages':
                        instance[key] = sample[key]
                
                samples.append(instance)
        
        # 显示预览结果
        preview_result = {
            "type": "conversation",
            "instances": samples
        }
        
        print(json.dumps(preview_result, ensure_ascii=False, indent=2))
        
    except Exception as e:
        print(f"预览时发生错误: {e}")

if __name__ == "__main__":
    # 数据集路径
    dataset_path = '/home/chengyuew/chengyuew/Data/tulu-3-sft-mixture-processed'
    output_file = "converted_dataset.json"
    
    print("=== Hugging Face 数据集转换工具 ===")
    print("可用的转换选项：")
    print("1. 按对话实例数量拆分")
    print("2. 按文件大小拆分")
    print("3. 不拆分，生成单个文件")
    print("4. 预览数据集")
    print()
    
    # 首先预览数据集
    print("数据集预览：")
    preview_dataset(dataset_path, num_samples=2)
    
    print("\n" + "="*50)
    
    # 示例1：按实例数量拆分（每个文件50000个对话实例）
    print("开始转换：按实例数量拆分")
    convert_ds_to_json(
        dataset_path,
        output_file,
        split_files=True,
        instances_per_file=50000,
        max_samples=None,  # 处理所有样本
        output_dir="split_dataset_by_instances"
    )
    
    # 其他示例（可按需取消注释）
    
    # 示例2：按文件大小拆分（每个文件最大100MB）
    # print("\n示例2：按文件大小拆分")
    # convert_ds_to_json(
    #     dataset_path,
    #     output_file,
    #     split_files=True,
    #     max_file_size_mb=100,
    #     output_dir="split_dataset_by_size"
    # )
    
    # 示例3：生成单个文件（适用于小数据集）
    # print("\n示例3：生成单个文件")
    # convert_ds_to_json(dataset_path, output_file, max_samples=10000)
    
    print("\n" + "="*50)
    print("使用说明：")
    print("1. 按实例数量拆分：")
    print("   convert_ds_to_json(dataset_path, output_file, split_files=True, instances_per_file=50000)")
    print("2. 按文件大小拆分：")
    print("   convert_ds_to_json(dataset_path, output_file, split_files=True, max_file_size_mb=100)")
    print("3. 限制处理样本数：")
    print("   convert_ds_to_json(dataset_path, output_file, max_samples=100000)")
    print("4. 预览数据集：")
    print("   preview_dataset(dataset_path, num_samples=5)")
    print()
    print("输出文件命名规则：")
    print("- 拆分文件：原文件名_0001.json, 原文件名_0002.json, ...")
    print("- 单个文件：按指定的output_file名称保存")