use tauri::{command, Emitter};
use serde::{Deserialize, Serialize};
use std::collections::{HashSet, HashMap};
use std::fs::File;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
use anyhow::Result;
use comfyui_sdk::{
    ComfyUIClient, ComfyUIClientConfig, ExecutionOptions,
    AI_MODEL_FACE_HAIR_FIX_TEMPLATE
};
use comfyui_sdk::utils::SimpleCallbacks;
use serde_json::json;

#[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)
}

// ============================================================================
// 图片增强相关结构体和命令
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageEnhancementProgress {
    pub current: usize,
    pub total: usize,
    pub percentage: f64,
    pub status: String,
    pub current_file: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageEnhancementResult {
    pub success: bool,
    pub message: String,
    pub input_path: String,
    pub output_path: Option<String>,
    pub execution_time: f64,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchImageEnhancementResult {
    pub success: bool,
    pub message: String,
    pub total_processed: usize,
    pub successful_count: usize,
    pub failed_count: usize,
    pub results: Vec<ImageEnhancementResult>,
    pub total_execution_time: f64,
}

/// 单张图片AI模型面部头发修复增强处理
#[command]
pub async fn ai_model_face_hair_fix_single_image(
    input_image: String,
    server_url: String,
    face_prompt: String,
    face_denoise: String,
    output_image: String,
    window: tauri::Window,
) -> Result<ImageEnhancementResult, String> {
    let start_time = std::time::Instant::now();

    // 发送开始处理的进度
    let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
        current: 0,
        total: 1,
        percentage: 0.0,
        status: "开始图片增强处理...".to_string(),
        current_file: Some(input_image.clone()),
    });

    match ai_model_face_hair_fix_internal(&input_image, &server_url, &face_prompt, &face_denoise, &output_image, &window).await {
        Ok(_) => {
            let execution_time = start_time.elapsed().as_secs_f64();

            let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
                current: 1,
                total: 1,
                percentage: 100.0,
                status: "图片增强完成".to_string(),
                current_file: Some(input_image.clone()),
            });

            Ok(ImageEnhancementResult {
                success: true,
                message: "图片增强成功".to_string(),
                input_path: input_image,
                output_path: Some(output_image),
                execution_time,
                error: None,
            })
        }
        Err(e) => {
            let execution_time = start_time.elapsed().as_secs_f64();

            let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
                current: 1,
                total: 1,
                percentage: 100.0,
                status: format!("图片增强失败: {}", e),
                current_file: Some(input_image.clone()),
            });

            Ok(ImageEnhancementResult {
                success: false,
                message: format!("图片增强失败: {}", e),
                input_path: input_image,
                output_path: None,
                execution_time,
                error: Some(e),
            })
        }
    }
}

/// 批量图片AI模型面部头发修复增强处理
#[command]
pub async fn ai_model_face_hair_fix_batch_images(
    input_dir: String,
    server_url: String,
    face_prompt: String,
    face_denoise: String,
    output_dir: String,
    window: tauri::Window,
) -> Result<BatchImageEnhancementResult, String> {
    let start_time = std::time::Instant::now();

    // 发送开始处理的进度
    let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
        current: 0,
        total: 0,
        percentage: 0.0,
        status: "扫描图片文件...".to_string(),
        current_file: None,
    });

    // 扫描输入目录中的所有图片文件
    let image_files = match scan_image_files(&input_dir).await {
        Ok(files) => files,
        Err(e) => return Err(format!("扫描图片文件失败: {}", e)),
    };

    if image_files.is_empty() {
        return Ok(BatchImageEnhancementResult {
            success: true,
            message: "未找到图片文件".to_string(),
            total_processed: 0,
            successful_count: 0,
            failed_count: 0,
            results: vec![],
            total_execution_time: start_time.elapsed().as_secs_f64(),
        });
    }

    let total_files = image_files.len();
    let mut results = Vec::new();
    let mut successful_count = 0;
    let mut failed_count = 0;

    // 确保输出目录存在
    if let Err(e) = fs::create_dir_all(&output_dir).await {
        return Err(format!("创建输出目录失败: {}", e));
    }

    // 处理每个图片文件
    for (index, (input_path, relative_path)) in image_files.iter().enumerate() {
        let current = index + 1;
        let percentage = (current as f64 / total_files as f64) * 100.0;

        let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
            current,
            total: total_files,
            percentage,
            status: format!("处理图片 {}/{}", current, total_files),
            current_file: Some(input_path.clone()),
        });

        // 构建输出路径，保持目录结构
        let output_path = Path::new(&output_dir).join(relative_path);

        // 确保输出文件的目录存在
        if let Some(parent) = output_path.parent() {
            if let Err(e) = fs::create_dir_all(parent).await {
                let error_msg = format!("创建输出目录失败: {}", e);
                results.push(ImageEnhancementResult {
                    success: false,
                    message: error_msg.clone(),
                    input_path: input_path.clone(),
                    output_path: None,
                    execution_time: 0.0,
                    error: Some(error_msg),
                });
                failed_count += 1;
                continue;
            }
        }

        let output_path_str = output_path.to_string_lossy().to_string();
        let file_start_time = std::time::Instant::now();

        match ai_model_face_hair_fix_internal(input_path, &server_url, &face_prompt, &face_denoise, &output_path_str, &window).await {
            Ok(_) => {
                let execution_time = file_start_time.elapsed().as_secs_f64();
                println!("✅ 图片处理成功: {} -> {}", input_path, output_path_str);
                results.push(ImageEnhancementResult {
                    success: true,
                    message: "图片增强成功".to_string(),
                    input_path: input_path.clone(),
                    output_path: Some(output_path_str),
                    execution_time,
                    error: None,
                });
                successful_count += 1;
            }
            Err(e) => {
                let execution_time = file_start_time.elapsed().as_secs_f64();
                println!("❌ 图片处理失败: {} - 错误: {}", input_path, e);
                eprintln!("详细错误信息: {}", e);
                results.push(ImageEnhancementResult {
                    success: false,
                    message: format!("图片增强失败: {}", e),
                    input_path: input_path.clone(),
                    output_path: None,
                    execution_time,
                    error: Some(e),
                });
                failed_count += 1;
            }
        }
    }

    let total_execution_time = start_time.elapsed().as_secs_f64();

    // 打印最终统计结果
    println!("📊 批量处理完成统计:");
    println!("   总计: {}", total_files);
    println!("   成功: {}", successful_count);
    println!("   失败: {}", failed_count);
    println!("   总耗时: {:.2}秒", total_execution_time);

    let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
        current: total_files,
        total: total_files,
        percentage: 100.0,
        status: format!("批量处理完成: 成功 {}, 失败 {}", successful_count, failed_count),
        current_file: None,
    });

    let batch_result = BatchImageEnhancementResult {
        success: failed_count == 0,
        message: format!("批量处理完成: 总计 {}, 成功 {}, 失败 {}", total_files, successful_count, failed_count),
        total_processed: total_files,
        successful_count,
        failed_count,
        results,
        total_execution_time,
    };

    println!("🎯 返回批量处理结果: {:?}", batch_result);
    Ok(batch_result)
}

/// AI模型面部头发修复内部处理函数
async fn ai_model_face_hair_fix_internal(
    input_image: &str,
    server_url: &str,
    face_prompt: &str,
    face_denoise: &str,
    output_image: &str,
    window: &tauri::Window,
) -> Result<(), String> {
    println!("🔄 开始处理图片: {}", input_image);
    println!("📋 处理参数: server_url={}, face_prompt={}, face_denoise={}", server_url, face_prompt, face_denoise);

    // 验证输入图片文件
    if !Path::new(input_image).exists() {
        return Err(format!("输入图片文件不存在: {}", input_image));
    }

    // 验证图片格式
    validate_image_file(input_image)?;

    // 初始化 ComfyUI 客户端
    println!("🔌 初始化ComfyUI客户端...");
    let mut client = ComfyUIClient::new(ComfyUIClientConfig {
        base_url: server_url.to_string(),
        ..Default::default()
    }).map_err(|e| format!("创建ComfyUI客户端失败: {}", e))?;

    // 连接到服务器
    println!("🌐 连接到ComfyUI服务器: {}", server_url);
    client.connect().await
        .map_err(|e| format!("连接ComfyUI服务器失败: {}", e))?;

    // 注册模板
    println!("📝 注册AI模型面部头发修复模板...");
    client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())
        .map_err(|e| format!("注册模板失败: {}", e))?;

    let template = client.templates_ref()
        .get_by_id("ai-model-face-hair-fix")
        .ok_or("模板未找到")?;

    // 上传图片
    println!("📤 上传图片: {}", input_image);
    let upload_response = client.upload_image(input_image, false).await
        .map_err(|e| format!("上传图片失败: {}", e))?;
    println!("✅ 图片上传成功，服务器文件名: {}", upload_response.name);

    // 设置进度回调
    let callbacks = Arc::new(
        SimpleCallbacks::new()
            .with_progress(|progress| {
                let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32;
                // 这里可以发送更详细的进度信息
            })
            .with_executing(|_node_id| {
                // 节点执行回调
            })
            .with_error(|error| {
                eprintln!("ComfyUI执行错误: {}", error.message);
            })
    );

    // 设置参数
    let mut parameters = HashMap::new();
    parameters.insert("input_image".to_string(), json!(upload_response.name));
    parameters.insert("face_prompt".to_string(), json!(face_prompt));
    parameters.insert("face_denoise".to_string(), json!(face_denoise));

    // 执行增强
    let result = client.execute_template_with_callbacks(
        template,
        parameters,
        ExecutionOptions {
            timeout: Some(std::time::Duration::from_secs(180)),
            priority: None,
        },
        callbacks,
    ).await.map_err(|e| format!("执行图片增强失败: {}", e))?;

    if !result.success {
        let error_msg = result.error
            .map(|e| e.message)
            .unwrap_or_else(|| "未知错误".to_string());
        return Err(format!("图片增强失败: {}", error_msg));
    }

    // 获取输出图片URL并下载
    if let Some(outputs) = &result.outputs {
        let image_urls = client.outputs_to_urls(outputs);

        if let Some(first_url) = image_urls.first() {
            // 下载增强后的图片
            download_image_from_url(first_url, output_image).await
                .map_err(|e| format!("下载增强图片失败: {}", e))?;
        } else {
            return Err("未找到输出图片URL".to_string());
        }
    } else {
        return Err("未找到输出结果".to_string());
    }

    // 断开连接
    client.disconnect().await
        .map_err(|e| format!("断开连接失败: {}", e))?;

    Ok(())
}

/// 扫描目录中的所有图片文件
async fn scan_image_files(dir: &str) -> Result<Vec<(String, String)>, String> {
    let mut image_files = Vec::new();
    let base_path = Path::new(dir);

    if !base_path.exists() {
        return Err(format!("目录不存在: {}", dir));
    }

    scan_directory_recursive(base_path, base_path, &mut image_files).await?;

    Ok(image_files)
}

/// 递归扫描目录
fn scan_directory_recursive<'a>(
    current_dir: &'a Path,
    base_dir: &'a Path,
    image_files: &'a mut Vec<(String, String)>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send + 'a>> {
    Box::pin(async move {
    let mut entries = fs::read_dir(current_dir).await
        .map_err(|e| format!("读取目录失败: {}", e))?;

    while let Some(entry) = entries.next_entry().await
        .map_err(|e| format!("读取目录项失败: {}", e))? {

        let path = entry.path();

        if path.is_dir() {
            // 递归处理子目录
            scan_directory_recursive(&path, base_dir, image_files).await?;
        } else if path.is_file() {
            // 检查是否为图片文件
            if is_image_file(&path) {
                let absolute_path = path.to_string_lossy().to_string();
                let relative_path = path.strip_prefix(base_dir)
                    .map_err(|e| format!("计算相对路径失败: {}", e))?
                    .to_string_lossy()
                    .to_string();

                image_files.push((absolute_path, relative_path));
            }
        }
    }

    Ok(())
    })
}

/// 检查文件是否为图片文件
fn is_image_file(path: &Path) -> bool {
    if let Some(extension) = path.extension() {
        let ext = extension.to_string_lossy().to_lowercase();
        matches!(ext.as_str(), "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp")
    } else {
        false
    }
}

/// 验证图片文件
fn validate_image_file(path: &str) -> Result<(), String> {
    let path = Path::new(path);

    if !path.exists() {
        return Err(format!("文件不存在: {}", path.display()));
    }

    if !path.is_file() {
        return Err(format!("路径不是文件: {}", path.display()));
    }

    // 检查文件扩展名
    if let Some(extension) = path.extension() {
        let ext = extension.to_string_lossy().to_lowercase();
        match ext.as_str() {
            "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp" => Ok(()),
            _ => Err(format!("不支持的图片格式: {}", ext)),
        }
    } else {
        Err("文件没有扩展名".to_string())
    }
}

/// 从URL下载图片到本地文件
async fn download_image_from_url(url: &str, output_path: &str) -> Result<(), String> {
    let response = reqwest::get(url).await
        .map_err(|e| format!("请求图片失败: {}", e))?;

    if !response.status().is_success() {
        return Err(format!("下载图片失败，状态码: {}", response.status()));
    }

    let bytes = response.bytes().await
        .map_err(|e| format!("读取图片数据失败: {}", e))?;

    fs::write(output_path, bytes).await
        .map_err(|e| format!("保存图片失败: {}", e))?;

    Ok(())
}
