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

import json
from pathlib import Path
import os

def check_reasoning_data(conversations):
    """
    检查对话是否为reasoning数据（包含<think>和</think>标签）
    
    Args:
        conversations: 对话列表
    
    Returns:
        bool: 是否为reasoning数据
    """
    for conv in conversations:
        if 'value' in conv:
            content = conv['value']
            if '<think>' in content and '</think>' in content:
                return True
    return False

def add_system_prompt(messages, is_reasoning):
    """
    根据是否为reasoning数据添加system prompt
    
    Args:
        messages: 消息列表
        is_reasoning: 是否为reasoning数据
    
    Returns:
        list: 添加了system prompt的消息列表
    """
    if not messages:
        return messages
    
    # 创建system prompt
    if is_reasoning:
        system_content = "detailed thinking on"
    else:
        system_content = "detailed thinking off"
    
    # 添加system message到开头
    system_message = {
        "role": "system",
        "content": system_content
    }
    
    return [system_message] + messages

def convert_jsonl_to_json(input_file, output_file, max_lines=None, split_files=False, instances_per_file=1000, output_dir=None, max_file_size_mb=None, add_system_prompts=True):
    """
    将JSONL文件转换为指定的JSON格式
    
    Args:
        input_file: 输入的JSONL文件路径
        output_file: 输出的JSON文件路径（如果不拆分）或基础文件名（如果拆分）
        max_lines: 最大处理行数，None表示处理所有行
        split_files: 是否拆分成多个文件
        instances_per_file: 每个文件包含的对话实例数量（当max_file_size_mb为None时使用）
        output_dir: 输出目录（拆分文件时使用）
        max_file_size_mb: 每个文件的最大大小（MB），如果设置了此参数，会按文件大小拆分
        add_system_prompts: 是否添加system prompt
    """
    input_path = Path(input_file)
    
    if not input_path.exists():
        print(f"错误：输入文件 {input_path} 不存在")
        return
    
    # 如果需要拆分文件，设置输出目录
    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)
    
    print(f"开始转换 {input_path}")
    if max_lines:
        print(f"最多处理 {max_lines} 行")
    if split_files:
        if max_file_size_mb:
            print(f"将按文件大小拆分，每个文件最大 {max_file_size_mb} MB")
        else:
            print(f"将拆分成多个文件，每个文件包含 {instances_per_file} 个对话实例")
    
    if add_system_prompts:
        print("将根据reasoning状态添加system prompt")
    
    # 统计变量
    reasoning_count = 0
    non_reasoning_count = 0
    
    try:
        # 如果需要拆分文件，获取下一个文件索引
        if split_files:
            file_index = _get_next_file_index(output_dir_path, output_file)
            print(f"从文件索引 {file_index} 开始")
        
        with open(input_path, 'r', encoding='utf-8') as f:
            line_count = 0
            processed_count = 0
            current_instances = []
            current_file_size = 0
            
            for line in f:
                line_count += 1
                
                # 如果设置了最大行数限制，检查是否达到
                if max_lines and line_count > max_lines:
                    break
                
                try:
                    # 解析JSONL行
                    data = json.loads(line.strip())
                    
                    # 检查是否包含conversations字段
                    if 'conversations' not in data:
                        continue
                    
                    # 转换conversations到目标格式
                    conversations = data['conversations']

                    # if not conversations or len(conversations) != 2:
                    #     continue
                    
                    # 检查是否为reasoning数据
                    is_reasoning = check_reasoning_data(conversations)
                    if is_reasoning:
                        reasoning_count += 1
                    else:
                        non_reasoning_count += 1
                    
                    # 创建messages数组
                    messages = []
                    for conv in conversations:
                        if 'from' in conv and 'value' in conv:
                            # 映射角色名称
                            role = conv['from'].lower()
                            if role == 'user':
                                role = 'user'
                            elif role == 'assistant':
                                role = 'assistant'
                            else:
                                # 跳过未知角色
                                continue
                                
                            message = {
                                "role": role,
                                "content": conv['value']
                            }
                            messages.append(message)
                    
                    # 如果有有效的messages，添加到结果中
                    if messages:
                        # 如果需要添加system prompt
                        if add_system_prompts:
                            messages = add_system_prompt(messages, is_reasoning)
                        
                        instance = {
                            "messages": messages
                        }
                        # 可选：添加原始数据集信息作为元数据
                        if 'dataset' in data:
                            instance['dataset'] = data['dataset']
                        
                        # 添加reasoning标记作为元数据
                        instance['is_reasoning'] = is_reasoning
                        
                        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 json.JSONDecodeError as e:
                    print(f"第 {line_count} 行JSON解析错误: {e}")
                    continue
                except Exception as e:
                    print(f"第 {line_count} 行处理错误: {e}")
                    continue
                
                # 每处理10000行显示一次进度
                if line_count % 10000 == 0:
                    print(f"已处理 {line_count:,} 行，转换了 {processed_count:,} 个对话")
                    print(f"  Reasoning: {reasoning_count:,}, 非Reasoning: {non_reasoning_count:,}")
        
        print(f"处理完成！总共处理了 {line_count:,} 行，转换了 {processed_count:,} 个对话")
        print(f"Reasoning样本: {reasoning_count:,} ({reasoning_count/processed_count*100:.1f}%)")
        print(f"非Reasoning样本: {non_reasoning_count:,} ({non_reasoning_count/processed_count*100:.1f}%)")
        
        # 保存结果
        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 _get_next_file_index(output_dir_path, base_filename):
    """
    获取下一个文件索引，考虑已存在的文件
    """
    base_name = Path(base_filename).stem
    pattern = f"{base_name}_*.json"
    
    existing_files = list(output_dir_path.glob(pattern))
    
    if not existing_files:
        return 1
    
    # 提取现有文件的索引号
    indices = []
    for file_path in existing_files:
        filename = file_path.stem
        if filename.startswith(f"{base_name}_"):
            try:
                index_part = filename[len(f"{base_name}_"):]
                index = int(index_part)
                indices.append(index)
            except ValueError:
                continue
    
    # 返回最大索引号 + 1
    return max(indices) + 1 if indices else 1

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 split_existing_json(input_file, output_dir, instances_per_file=1000, max_file_size_mb=None):
    """
    拆分已存在的JSON文件
    
    Args:
        input_file: 输入的JSON文件路径
        output_dir: 输出目录
        instances_per_file: 每个文件包含的对话实例数量
        max_file_size_mb: 每个文件的最大大小（MB）
    """
    input_path = Path(input_file)
    output_dir_path = Path(output_dir)
    
    if not input_path.exists():
        print(f"错误：输入文件 {input_path} 不存在")
        return
    
    # 创建输出目录
    output_dir_path.mkdir(exist_ok=True)
    
    print(f"开始拆分 {input_path}")
    if max_file_size_mb:
        print(f"按文件大小拆分，每个文件最大 {max_file_size_mb} MB")
    else:
        print(f"按实例数量拆分，每个文件包含 {instances_per_file} 个对话实例")
    
    try:
        # 读取原始JSON文件
        with open(input_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        if 'instances' not in data:
            print("错误：输入文件不包含 'instances' 字段")
            return
        
        instances = data['instances']
        total_instances = len(instances)
        print(f"总共有 {total_instances:,} 个对话实例")
        
        # 拆分实例
        current_instances = []
        file_index = _get_next_file_index(output_dir_path, input_path.name)
        print(f"从文件索引 {file_index} 开始")
        
        for i, instance in enumerate(instances):
            current_instances.append(instance)
            
            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 i == total_instances - 1:
                should_save = True
            
            if should_save:
                base_filename = input_path.stem
                _save_split_file(current_instances, output_dir_path, file_index, base_filename)
                current_instances = []
                file_index += 1
        
        print(f"拆分完成！已生成 {file_index - 1} 个文件，保存在 {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")
        
    except Exception as e:
        print(f"拆分过程中发生错误: {e}")

def preview_conversion(input_file, num_samples=5, add_system_prompts=True):
    """
    预览转换结果
    
    Args:
        input_file: 输入文件路径
        num_samples: 预览样本数量
        add_system_prompts: 是否添加system prompt
    """
    input_path = Path(input_file)
    
    if not input_path.exists():
        print(f"错误：输入文件 {input_path} 不存在")
        return
    
    print("=== 转换预览 ===")
    if add_system_prompts:
        print("将根据reasoning状态添加system prompt")
    
    try:
        with open(input_path, 'r', encoding='utf-8') as f:
            samples = []
            line_count = 0
            reasoning_count = 0
            non_reasoning_count = 0
            
            for line in f:
                line_count += 1
                if len(samples) >= num_samples:
                    break
                
                try:
                    data = json.loads(line.strip())
                    if 'conversations' in data and data['conversations']:
                        # 转换样本
                        conversations = data['conversations']
                        
                        # 检查是否为reasoning数据
                        is_reasoning = check_reasoning_data(conversations)
                        if is_reasoning:
                            reasoning_count += 1
                        else:
                            non_reasoning_count += 1
                        
                        messages = []
                        
                        for conv in conversations:
                            if 'from' in conv and 'value' in conv:
                                role = conv['from'].lower()
                                if role in ['user', 'assistant']:
                                    # 截断长内容用于预览
                                    content = conv['value']
                                    if len(content) > 200:
                                        content = content[:200] + "..."
                                    
                                    message = {
                                        "role": role,
                                        "content": content
                                    }
                                    messages.append(message)
                        
                        if messages:
                            # 如果需要添加system prompt
                            if add_system_prompts:
                                messages = add_system_prompt(messages, is_reasoning)
                            
                            instance = {
                                "messages": messages,
                                "is_reasoning": is_reasoning
                            }
                            if 'dataset' in data:
                                instance['dataset'] = data['dataset']
                            
                            samples.append(instance)
                
                except json.JSONDecodeError:
                    continue
                except Exception as e:
                    print(f"预览时发生错误: {e}")
                    continue
            
            # 显示预览结果
            preview_result = {
                "type": "conversation",
                "instances": samples
            }
            
            print(json.dumps(preview_result, ensure_ascii=False, indent=2))
            
            # 显示统计信息
            total_samples = len(samples)
            if total_samples > 0:
                print(f"\n预览统计:")
                print(f"Reasoning样本: {reasoning_count} ({reasoning_count/total_samples*100:.1f}%)")
                print(f"非Reasoning样本: {non_reasoning_count} ({non_reasoning_count/total_samples*100:.1f}%)")
            
    except Exception as e:
        print(f"预览时发生错误: {e}")

def test_file_indexing(output_dir, base_filename):
    """
    测试文件索引功能
    """
    output_dir_path = Path(output_dir)
    output_dir_path.mkdir(exist_ok=True)
    
    print(f"测试目录: {output_dir_path}")
    print(f"基础文件名: {base_filename}")
    
    # 创建一些测试文件
    test_files = [
        f"{base_filename}_0001.json",
        f"{base_filename}_0003.json", 
        f"{base_filename}_0005.json"
    ]
    
    for filename in test_files:
        filepath = output_dir_path / filename
        with open(filepath, 'w') as f:
            f.write('{"test": "data"}')
        print(f"创建测试文件: {filename}")
    
    # 测试获取下一个索引
    next_index = _get_next_file_index(output_dir_path, base_filename)
    print(f"下一个文件索引应该是: {next_index}")
    
    # 清理测试文件
    for filename in test_files:
        filepath = output_dir_path / filename
        if filepath.exists():
            filepath.unlink()
            print(f"删除测试文件: {filename}")

if __name__ == "__main__":
    input_file = "/lustre/fsw/portfolios/llmservice/users/haozh/data/chocolate-dolphin/chocolate-dolphin_commercial.shuf.removelong.jsonl"
    output_file = "converted_conversations.json"
    output_dir = "shizhe_stage1"
    
    print("=== JSON文件拆分转换工具 ===")
    print("可用的拆分选项：")
    print("1. 按对话实例数量拆分")
    print("2. 按文件大小拆分")
    print("3. 不拆分，生成单个文件")
    print("4. 拆分已存在的JSON文件")
    print("5. 测试文件索引功能")
    print("6. 预览转换效果（包含reasoning检测）")
    print("7. 转换并添加system prompt（根据reasoning状态）")
    print()
    
    # 测试文件索引功能
    print("测试文件索引功能：")
    test_file_indexing("test_output", "converted_conversations")
    print()
    
    # 预览转换效果（包含reasoning检测）
    print("预览转换效果（包含reasoning检测）：")
    preview_conversion(input_file, num_samples=3, add_system_prompts=True)
    print()
    
    # 示例1：按实例数量拆分（每个文件50000个对话实例），包含reasoning检测和system prompt
    print("示例1：按实例数量拆分（包含reasoning检测和system prompt）")
    convert_jsonl_to_json(
        input_file, 
        output_file, 
        split_files=True, 
        instances_per_file=50_000,
        max_lines=None,
        output_dir=output_dir,
        add_system_prompts=True  # 启用reasoning检测和system prompt
    )
    
    # 示例2：按文件大小拆分（每个文件最大10MB）
    # print("\n示例2：按文件大小拆分")
    # convert_jsonl_to_json(
    #     input_file, 
    #     output_file, 
    #     split_files=True, 
    #     max_file_size_mb=10,
    #     output_dir="split_by_size",
    #     add_system_prompts=True
    # )
    
    # 示例3：单个文件（注释掉）
    # print("\n示例3：生成单个文件")
    # convert_jsonl_to_json(input_file, output_file, add_system_prompts=True)
    
    # 示例4：拆分已存在的JSON文件（注释掉）
    # print("\n示例4：拆分已存在的JSON文件")
    # split_existing_json("converted_conversations.json", "split_existing", instances_per_file=500)
    
    print("\n" + "="*50)
    print("使用说明：")
    print("1. 按实例数量拆分：")
    print("   convert_jsonl_to_json(input_file, output_file, split_files=True, instances_per_file=1000, add_system_prompts=True)")
    print("2. 按文件大小拆分：")
    print("   convert_jsonl_to_json(input_file, output_file, split_files=True, max_file_size_mb=10, add_system_prompts=True)")
    print("3. 拆分已存在的JSON文件：")
    print("   split_existing_json('input.json', 'output_dir', instances_per_file=500)")
    print("4. 限制处理行数：")
    print("   convert_jsonl_to_json(input_file, output_file, max_lines=10000, add_system_prompts=True)")
    print("5. 预览转换效果：")
    print("   preview_conversion(input_file, num_samples=5, add_system_prompts=True)")
    print("6. 测试文件索引功能：")
    print("   test_file_indexing('output_dir', 'base_filename')")
    print()
    print("新增功能：")
    print("- reasoning检测：自动检测对话中是否包含<think>和</think>标签")
    print("- system prompt：根据reasoning状态自动添加相应的system prompt")
    print("  * reasoning样本：'You are a helpful assistant. Please provide detailed thinking on the given question.'")
    print("  * 非reasoning样本：'You are a helpful assistant. Please provide detailed thinking off and give direct answers.'")
    print("- 统计信息：显示reasoning和非reasoning样本的数量和比例")
    print()
    print("输出文件命名规则：")
    print("- 拆分文件：原文件名_0001.json, 原文件名_0002.json, ...")
    print("- 单个文件：按指定的output_file名称保存")
    print("- 智能索引：如果输出目录已存在文件，会从最大编号+1开始")
    print("- 元数据：每个实例包含is_reasoning字段标识是否为reasoning样本") 