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

import json
from pathlib import Path
import os
from datasets import load_dataset
from typing import Optional, List
import re

def convert_nemotron_datasets(dataset_path, output_dir, 
                             max_samples_per_split=None,
                             split_files=True, 
                             instances_per_file=50000,
                             max_file_size_mb=None,
                             include_splits=None,
                             exclude_reasoning=True):
    """
    将Nemotron datasets格式转换为LMFlow所需的JSON格式
    
    Args:
        dataset_path: Nemotron数据集路径
        output_dir: 输出目录
        max_samples_per_split: 每个子集最多处理多少样本，None表示全部处理
        split_files: 是否拆分成多个文件
        instances_per_file: 每个文件包含的对话实例数量
        max_file_size_mb: 每个文件的最大大小（MB）
        include_splits: 要处理的子集列表，None表示处理所有子集
        exclude_reasoning: 是否排除reasoning="on"的数据
    """
    
    print(f"开始加载Nemotron数据集: {dataset_path}")
    
    try:
        # 加载数据集
        ds = load_dataset(dataset_path)
        print(f"成功加载数据集，包含 {len(ds)} 个子集")
        
        # 创建输出目录
        output_dir_path = Path(output_dir)
        output_dir_path.mkdir(exist_ok=True)
        
        # 如果指定了包含的子集，则过滤
        if include_splits:
            available_splits = list(ds.keys())
            splits_to_process = [split for split in include_splits if split in available_splits]
            print(f"将处理指定的子集: {splits_to_process}")
        else:
            splits_to_process = list(ds.keys())
            print(f"将处理所有子集: {splits_to_process}")
        
        # 处理每个子集
        for split_name in splits_to_process:
            print(f"\n{'='*50}")
            print(f"开始处理子集: {split_name}")
            print(f"原始数据量: {len(ds[split_name]):,}")
            
            # 转换这个子集
            convert_split(ds[split_name], split_name, output_dir_path,
                         max_samples=max_samples_per_split,
                         split_files=split_files,
                         instances_per_file=instances_per_file,
                         max_file_size_mb=max_file_size_mb,
                         exclude_reasoning=exclude_reasoning)
        
        print(f"\n{'='*50}")
        print("所有子集处理完成！")
        
        # 计算总文件大小
        total_size = sum(f.stat().st_size for f in output_dir_path.glob("*.json"))
        print(f"输出目录: {output_dir_path}")
        print(f"总文件大小: {total_size / (1024**2):.2f} MB")
        
    except Exception as e:
        print(f"加载或处理数据集时发生错误: {e}")

def convert_split(dataset_split, split_name, output_dir_path, 
                 max_samples=None,
                 split_files=True,
                 instances_per_file=50000,
                 max_file_size_mb=None,
                 exclude_reasoning=True):
    """
    转换单个数据子集
    """
    
    # 统计计数器
    total_processed = 0
    reasoning_excluded = 0
    error_count = 0
    
    # 准备输出
    current_instances = []
    file_index = 1
    
    # 处理样本
    for i, sample in enumerate(dataset_split):
        try:
            # 限制样本数量
            if max_samples and i >= max_samples:
                break
            
            # 检查是否排除reasoning数据
            if exclude_reasoning and sample.get('reasoning', '').lower() == 'on':
                reasoning_excluded += 1
                continue
            
            # 转换messages格式
            messages = []
            
            # 处理原始messages
            original_messages = sample.get('messages', [])
            
            for msg in original_messages:
                role = msg.get('role', '').lower()
                content = msg.get('content', '')
                
                # 跳过空的system消息
                if role == 'system' and not content.strip():
                    continue
                
                # 只保留user和assistant消息
                if role in ['user', 'assistant']:
                    # 处理assistant消息中的<think>标签
                    if role == 'assistant':
                        content = clean_assistant_content(content)
                    
                    messages.append({
                        "role": role,
                        "content": content
                    })
            
            # 确保消息格式正确（必须以user开始，assistant结束）
            if not messages:
                continue
                
            if messages[0]['role'] != 'user':
                continue
                
            if messages[-1]['role'] != 'assistant':
                continue
            
            # 创建实例
            instance = {
                "messages": messages
            }
            
            # 添加可选的元数据
            if sample.get('category'):
                instance['category'] = sample['category']
            
            if sample.get('uuid'):
                instance['conversation_id'] = sample['uuid']
            
            current_instances.append(instance)
            total_processed += 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, split_name, file_index)
                    current_instances = []
                    file_index += 1
            
        except Exception as e:
            error_count += 1
            print(f"处理第 {i+1} 个样本时出错: {e}")
            continue
        
        # 进度显示
        if (i + 1) % 10000 == 0:
            print(f"已处理 {i+1:,} 个样本，转换了 {total_processed:,} 个对话")
    
    # 保存剩余的实例
    if current_instances:
        if split_files:
            _save_split_file(current_instances, output_dir_path, split_name, file_index)
        else:
            # 保存为单个文件
            filename = f"{split_name}_conversations.json"
            filepath = output_dir_path / filename
            
            result = {
                "type": "conversation",
                "instances": current_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(current_instances)} 个对话实例, {file_size:.2f} MB)")
    
    print(f"子集 {split_name} 处理完成:")
    print(f"  - 总处理样本: {total_processed:,}")
    print(f"  - 排除reasoning数据: {reasoning_excluded:,}")
    print(f"  - 错误样本: {error_count:,}")
    if split_files:
        print(f"  - 生成文件数: {file_index}")

def clean_assistant_content(content):
    """
    清理assistant内容，移除<think>标签但保留内容
    """
    if not content:
        return content
    
    # 移除<think>和</think>标签，但保留其中的内容
    # 这样可以保留推理过程
    cleaned_content = re.sub(r'<think>', '', content)
    cleaned_content = re.sub(r'</think>', '', cleaned_content)
    
    return cleaned_content.strip()

def _save_split_file(instances, output_dir_path, split_name, file_index):
    """
    保存拆分的文件
    """
    filename = f"{split_name}_{file_index:04d}.json"
    filepath = output_dir_path / filename
    
    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, split_name='math', num_samples=3):
    """
    预览数据集内容
    """
    print(f"预览数据集: {dataset_path}")
    
    try:
        ds = load_dataset(dataset_path)
        
        if split_name not in ds:
            print(f"错误: 子集 '{split_name}' 不存在")
            print(f"可用子集: {list(ds.keys())}")
            return
        
        print(f"\n预览子集 '{split_name}' 的前 {num_samples} 个样本:")
        
        for i in range(min(num_samples, len(ds[split_name]))):
            sample = ds[split_name][i]
            print(f"\n--- 样本 {i+1} ---")
            print(f"UUID: {sample.get('uuid', 'N/A')}")
            print(f"Category: {sample.get('category', 'N/A')}")
            print(f"Reasoning: {sample.get('reasoning', 'N/A')}")
            print(f"Messages count: {len(sample.get('messages', []))}")
            
            # 显示messages内容（截取前200字符）
            for j, msg in enumerate(sample.get('messages', [])):
                role = msg.get('role', 'unknown')
                content = msg.get('content', '')[:200]
                if len(msg.get('content', '')) > 200:
                    content += "..."
                print(f"  {j+1}. {role}: {content}")
        
    except Exception as e:
        print(f"预览数据集时发生错误: {e}")

if __name__ == "__main__":
    # 数据集路径
    dataset_path = "/lustre/fsw/portfolios/llmservice/users/haozh/data/Nemotron-Post-Training-Dataset-v2"
    output_dir = "data/nemotron_v2_converted"
    
    print("=== Nemotron数据集转换工具 ===")
    print()
    
    # 预览数据集
    print("1. 预览数据集内容:")
    preview_dataset(dataset_path, split_name='math', num_samples=2)
    print()
    
    # 转换配置
    print("2. 开始转换数据集:")
    print("配置:")
    print(f"  - 输出目录: {output_dir}")
    print(f"  - 每个文件最多实例数: 50,000")
    print(f"  - 排除reasoning数据: True")
    print(f"  - 处理所有子集")
    print()
    
    # 执行转换
    convert_nemotron_datasets(
        dataset_path=dataset_path,
        output_dir=output_dir,
        max_samples_per_split=None,  # 每个子集最多10万样本用于测试
        split_files=True,
        instances_per_file=50000,
        exclude_reasoning=True,
        include_splits=['stem', 'chat', 'math', 'code']  # 先处理这几个子集
    )
    
    print("\n" + "="*50)
    print("使用说明:")
    print("1. 处理所有子集:")
    print("   convert_nemotron_datasets(dataset_path, output_dir)")
    print("2. 处理指定子集:")
    print("   convert_nemotron_datasets(dataset_path, output_dir, include_splits=['math', 'code'])")
    print("3. 限制样本数量:")
    print("   convert_nemotron_datasets(dataset_path, output_dir, max_samples_per_split=10000)")
    print("4. 预览数据集:")
    print("   preview_dataset(dataset_path, split_name='math', num_samples=5)")
