import json
import os
import numpy as np
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
import decord
from decord import VideoReader, cpu

# 屏蔽底层烦人的 C++ 报错输出
# decord.bridge.set_bridge('torch')
import logging
decord.logging.set_level(logging.ERROR)

# 🌟 模拟你修复后的训练配置
class DataArgs:
    def __init__(self):
        self.video_fps = 1
        self.frames_upbound = 16  # 保持和你训练时设置的上限一致
        self.force_sample = True

data_args = DataArgs()

def check_single_sample(item):
    """
    核心验证逻辑：完全复刻训练时的视频抽帧代码
    只有通过了真实 get_batch 考验的视频，才算真正的“好数据”
    """
    video_path = item.get("video_abs_path")
    if not video_path or not os.path.exists(video_path):
        return False, item, "文件不存在或路径为空"

    try:
        # 完全照搬你的 process_video_with_decord 逻辑
        vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
        total_frame_num = len(vr)
        
        if total_frame_num == 0:
            return False, item, "视频总帧数为 0"
        
        avg_fps = round(vr.get_avg_fps() / data_args.video_fps)
        if avg_fps <= 0: 
            avg_fps = 1
            
        frame_idx = [i for i in range(0, total_frame_num, avg_fps)]
        
        if data_args.frames_upbound > 0:
            if len(frame_idx) > data_args.frames_upbound or data_args.force_sample:
                uniform_sampled_frames = np.linspace(0, total_frame_num - 1, data_args.frames_upbound, dtype=int)
                frame_idx = uniform_sampled_frames.tolist()
        
        # 🚨 终极考验：尝试真实提取这批帧！坏视频在这里一定会原形毕露崩溃
        video = vr.get_batch(frame_idx).asnumpy()
        
        if video.shape[0] == 0:
            return False, item, "提取到的帧矩阵为空"
            
        return True, item, "OK"
    except Exception as e:
        # 捕获所有 decord 底层抛出的异常 (包括你遇到的 h264 错误)
        return False, item, f"解码崩溃: {str(e)}"

def run_deep_clean(input_json, clean_json, bad_json, num_workers=16):
    print(f"正在读取数据集: {input_json} ...")
    with open(input_json, 'r') as f:
        data_list = json.load(f)
    
    # data_list = data_list[5000:5100]

    clean_data = []
    bad_data = []
    
    print(f"\n🚀 开始深度清洗 {len(data_list)} 条数据 (启用 {num_workers} 个进程并发)...")
    
    # 使用进程池极大加快处理速度
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(check_single_sample, item) for item in data_list]
        
        for future in tqdm(as_completed(futures), total=len(data_list), desc="🔨 质检进度"):
            is_valid, item, error_msg = future.result()
            if is_valid:
                clean_data.append(item)
            else:
                item['error_reason'] = error_msg
                bad_data.append(item)
                
    print("\n" + "="*50)
    print(f"✅ 清洗彻底完成！")
    print(f"🎉 纯净数据: {len(clean_data)} 条 -> 已保存至 {clean_json}")
    print(f"☠️ 损毁数据: {len(bad_data)} 条 -> 已保存至 {bad_json}")
    print("="*50)
    
    with open(clean_json, 'w') as f:
        json.dump(clean_data, f, indent=4)
    with open(bad_json, 'w') as f:
        json.dump(bad_data, f, indent=4)

if __name__ == "__main__":
    # 替换成你当前的 JSON 路径
    INPUT_FILE = "./data/video_reversal_sft_train_cleaned_10k.json" 
    CLEAN_FILE = "./data/video_reversal_sft_train_ultra_clean.json"
    BAD_FILE = "./data/video_reversal_sft_train_bad_samples.json"
    
    # 根据你机器的 CPU 核心数调整 num_workers，核心越多越快
    run_deep_clean(INPUT_FILE, CLEAN_FILE, BAD_FILE, num_workers=16)