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

import json
import sys
from pathlib import Path

def read_jsonl_file(file_path, num_samples=10):
    """
    读取JSONL文件并显示前几行数据
    
    Args:
        file_path: JSONL文件路径
        num_samples: 要显示的样本数量
    """
    file_path = Path(file_path)
    
    if not file_path.exists():
        print(f"错误：文件 {file_path} 不存在")
        return
    
    print(f"正在读取文件: {file_path}")
    print(f"文件大小: {file_path.stat().st_size / (1024**3):.2f} GB")
    print("="*50)
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            line_count = 0
            sample_count = 0
            
            print(f"前 {num_samples} 行数据:")
            print("-" * 50)
            
            for line in f:
                line_count += 1
                
                if sample_count < num_samples:
                    try:
                        data = json.loads(line.strip())
                        print(f"第 {line_count} 行:")
                        print(f"数据键: {list(data.keys())}")
                        
                        # 显示每个字段的前100个字符
                        for key, value in data.items():
                            if isinstance(value, str):
                                preview = value[:100] + "..." if len(value) > 100 else value
                                print(f"  {key}: {preview}")
                            else:
                                print(f"  {key}: {value}")
                        print()
                        sample_count += 1
                        
                    except json.JSONDecodeError as e:
                        print(f"第 {line_count} 行JSON解析错误: {e}")
                        continue
                
                # 每处理10000行显示一次进度
                if line_count % 10000 == 0:
                    print(f"已处理 {line_count:,} 行...")
                    
                # 如果只是查看样本，处理前100000行就够了
                if sample_count >= num_samples and line_count > 100000:
                    break
                    
            print(f"总共处理了 {line_count:,} 行数据")
            
    except Exception as e:
        print(f"读取文件时发生错误: {e}")

def analyze_jsonl_structure(file_path, max_lines=1000):
    """
    分析JSONL文件的结构
    
    Args:
        file_path: JSONL文件路径
        max_lines: 分析的最大行数
    """
    file_path = Path(file_path)
    all_keys = set()
    key_counts = {}
    
    print(f"\n分析文件结构 (最多分析 {max_lines} 行):")
    print("-" * 50)
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f):
                if i >= max_lines:
                    break
                    
                try:
                    data = json.loads(line.strip())
                    for key in data.keys():
                        all_keys.add(key)
                        key_counts[key] = key_counts.get(key, 0) + 1
                except json.JSONDecodeError:
                    continue
                    
        print(f"数据字段统计 (基于前 {min(i+1, max_lines)} 行):")
        for key in sorted(all_keys):
            print(f"  {key}: 出现在 {key_counts[key]} 行中")
            
    except Exception as e:
        print(f"分析文件结构时发生错误: {e}")

if __name__ == "__main__":
    file_path = "/lustre/fsw/portfolios/llmservice/users/haozh/data/chocolate-dolphin/chocolate-dolphin_commercial.shuf.removelong.jsonl"
    
    print("开始读取JSONL文件...")
    read_jsonl_file(file_path, num_samples=5)
    
    print("\n" + "="*50)
    analyze_jsonl_structure(file_path, max_lines=1000) 