mixvideo-v2/apps/desktop/src-tauri/src/presentation/commands/tools_commands.rs

use tauri::{command, Emitter};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::fs::File;
use std::io::{BufRead, BufReader, Write};
use std::path::Path;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataCleaningProgress {
    pub current: usize,
    pub total: usize,
    pub percentage: f64,
    pub status: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataCleaningResult {
    pub success: bool,
    pub message: String,
    pub original_count: usize,
    pub removed_count: usize,
    pub final_count: usize,
    pub output_file: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct StorageObject {
    pub uri: String,
    // 其他字段可以忽略，我们只关心 uri
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ProcessedData {
    pub json_data: String,
}

/// AI检索图片/数据清洗 - JSONL格式处理
#[command]
pub async fn clean_jsonl_data(
    all_data_file: String,
    remove_data_file: String,
    output_file: String,
    window: tauri::Window,
) -> Result<DataCleaningResult, String> {
    // 验证文件存在
    if !Path::new(&all_data_file).exists() {
        return Err(format!("全部数据文件不存在: {}", all_data_file));
    }
    
    if !Path::new(&remove_data_file).exists() {
        return Err(format!("要去除的数据文件不存在: {}", remove_data_file));
    }

    // 发送开始处理的进度
    let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
        current: 0,
        total: 0,
        percentage: 0.0,
        status: "开始处理...".to_string(),
    });

    // 第一步：读取要去除的数据文件，提取所有URI
    let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
        current: 0,
        total: 0,
        percentage: 10.0,
        status: "读取要去除的数据文件...".to_string(),
    });

    let remove_uris = match extract_remove_uris(&remove_data_file).await {
        Ok(uris) => uris,
        Err(e) => return Err(format!("读取要去除的数据文件失败: {}", e)),
    };

    let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
        current: 0,
        total: 0,
        percentage: 30.0,
        status: format!("找到 {} 个要去除的URI", remove_uris.len()),
    });

    // 第二步：处理全部数据文件
    let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
        current: 0,
        total: 0,
        percentage: 40.0,
        status: "开始处理全部数据文件...".to_string(),
    });

    let result = match process_all_data_file(&all_data_file, &remove_uris, &output_file, &window).await {
        Ok(result) => result,
        Err(e) => return Err(format!("处理全部数据文件失败: {}", e)),
    };

    let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
        current: result.final_count,
        total: result.original_count,
        percentage: 100.0,
        status: "处理完成".to_string(),
    });

    Ok(result)
}

/// 从要去除的数据文件中提取所有URI
async fn extract_remove_uris(file_path: &str) -> Result<HashSet<String>, String> {
    let file = File::open(file_path)
        .map_err(|e| format!("无法打开文件: {}", e))?;
    
    let reader = BufReader::new(file);
    let mut uris = HashSet::new();

    for line in reader.lines() {
        let line = line.map_err(|e| format!("读取行失败: {}", e))?;
        
        if line.trim().is_empty() {
            continue;
        }

        // 解析JSON
        let processed_data: ProcessedData = serde_json::from_str(&line)
            .map_err(|e| format!("解析JSON失败: {} - 行内容: {}", e, line))?;

        // 从json_data字段中解析内部JSON
        let inner_json: serde_json::Value = serde_json::from_str(&processed_data.json_data)
            .map_err(|e| format!("解析内部JSON失败: {}", e))?;

        // 提取uri字段
        if let Some(uri) = inner_json.get("uri").and_then(|v| v.as_str()) {
            uris.insert(uri.to_string());
        }
    }

    Ok(uris)
}

/// 处理全部数据文件，过滤掉要去除的URI
async fn process_all_data_file(
    input_file: &str,
    remove_uris: &HashSet<String>,
    output_file: &str,
    window: &tauri::Window,
) -> Result<DataCleaningResult, String> {
    let input_file_handle = File::open(input_file)
        .map_err(|e| format!("无法打开输入文件: {}", e))?;
    
    let reader = BufReader::new(input_file_handle);
    
    let mut output_file_handle = File::create(output_file)
        .map_err(|e| format!("无法创建输出文件: {}", e))?;

    let mut original_count = 0;
    let mut removed_count = 0;
    let mut final_count = 0;

    // 首先计算总行数用于进度显示
    let total_lines = count_lines(input_file)?;

    for (line_number, line) in reader.lines().enumerate() {
        let line = line.map_err(|e| format!("读取行失败: {}", e))?;
        
        if line.trim().is_empty() {
            continue;
        }

        original_count += 1;

        // 解析JSON
        let storage_object: StorageObject = match serde_json::from_str(&line) {
            Ok(obj) => obj,
            Err(e) => {
                // 如果解析失败，跳过这一行但记录警告
                eprintln!("警告: 解析JSON失败，跳过行 {}: {} - 内容: {}", line_number + 1, e, line);
                continue;
            }
        };

        // 检查URI是否在要去除的列表中
        if remove_uris.contains(&storage_object.uri) {
            removed_count += 1;
        } else {
            // 写入输出文件
            writeln!(output_file_handle, "{}", line)
                .map_err(|e| format!("写入输出文件失败: {}", e))?;
            final_count += 1;
        }

        // 每处理100行发送一次进度更新
        if line_number % 100 == 0 {
            let percentage = 40.0 + (line_number as f64 / total_lines as f64) * 50.0;
            let _ = window.emit("data-cleaning-progress", DataCleaningProgress {
                current: line_number,
                total: total_lines,
                percentage,
                status: format!("已处理 {} 行，已去除 {} 个重复项", line_number, removed_count),
            });
        }
    }

    // 确保文件被正确写入
    output_file_handle.flush()
        .map_err(|e| format!("刷新输出文件失败: {}", e))?;

    Ok(DataCleaningResult {
        success: true,
        message: "数据清洗完成".to_string(),
        original_count,
        removed_count,
        final_count,
        output_file: output_file.to_string(),
    })
}

/// 计算文件行数
fn count_lines(file_path: &str) -> Result<usize, String> {
    let file = File::open(file_path)
        .map_err(|e| format!("无法打开文件计算行数: {}", e))?;
    
    let reader = BufReader::new(file);
    let count = reader.lines().count();
    Ok(count)
}