//! Topaz Video AI 模型管理工具
//! 
//! 提供模型下载、检查和管理功能

use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use crate::core::TvaiError;

/// 模型信息
#[derive(Debug, Clone)]
pub struct ModelInfo {
    pub short_name: String,
    pub display_name: Option<String>,
    pub model_type: ModelType,
    pub is_downloaded: bool,
}

/// 模型类型
#[derive(Debug, Clone, PartialEq)]
pub enum ModelType {
    /// 超分辨率模型
    Upscale,
    /// 帧插值模型
    Interpolation,
    /// 其他模型 (稳定化、参数估算等)
    Other,
}

/// 模型管理器
pub struct ModelManager {
    models_dir: PathBuf,
    topaz_exe: Option<PathBuf>,
    topaz_ffmpeg: PathBuf,
}

impl ModelManager {
    /// 创建新的模型管理器
    pub fn new(topaz_path: &Path) -> Result<Self, TvaiError> {
        let topaz_exe = topaz_path.join("Topaz Video AI.exe");
        let topaz_ffmpeg = topaz_path.join("ffmpeg.exe");

        if !topaz_ffmpeg.exists() {
            return Err(TvaiError::TopazNotFound(
                format!("Topaz FFmpeg not found: {}", topaz_ffmpeg.display())
            ));
        }

        // 检查多个可能的模型目录位置
        let possible_models_dirs = vec![
            topaz_path.join("models"),
            PathBuf::from(r"C:\ProgramData\Topaz Labs LLC\Topaz Video AI\models"),
        ];

        let mut models_dir = None;
        for dir in possible_models_dirs {
            if dir.exists() {
                // 检查是否包含模型文件
                if let Ok(entries) = std::fs::read_dir(&dir) {
                    let has_models = entries
                        .filter_map(|e| e.ok())
                        .any(|e| {
                            let name = e.file_name().to_string_lossy().to_lowercase();
                            name.ends_with(".tz") || (name.ends_with(".json") &&
                                !name.contains("audio-codecs") &&
                                !name.contains("benchmarks") &&
                                !name.contains("proxy"))
                        });

                    if has_models {
                        models_dir = Some(dir);
                        break;
                    }
                }
            }
        }

        let models_dir = models_dir.ok_or_else(|| {
            TvaiError::TopazNotFound(
                "No valid models directory found. Please ensure Topaz Video AI is properly installed and models are downloaded.".to_string()
            )
        })?;

        Ok(Self {
            models_dir,
            topaz_exe: if topaz_exe.exists() { Some(topaz_exe) } else { None },
            topaz_ffmpeg,
        })
    }

    /// 获取所有模型信息
    pub fn get_all_models(&self) -> Result<Vec<ModelInfo>, TvaiError> {
        let mut models = Vec::new();
        let downloaded_models = self.get_downloaded_models()?;

        if let Ok(entries) = fs::read_dir(&self.models_dir) {
            for entry in entries {
                if let Ok(entry) = entry {
                    let path = entry.path();
                    if let Some(extension) = path.extension() {
                        if extension == "json" {
                            let filename = path.file_name().unwrap().to_string_lossy();

                            // 跳过非模型文件
                            if filename.contains("audio-codecs") ||
                               filename.contains("benchmarks") ||
                               filename.contains("proxy") ||
                               filename.contains("video-encoders") {
                                continue;
                            }

                            if let Ok(content) = fs::read_to_string(&path) {
                                // 解析 shortName，如过短（家族名）则回退到文件名基名作为更精确的短名
                                let base_stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("").to_string();
                                let mut short_name = extract_short_name(&content).unwrap_or_else(|| base_stem.clone());

                                // 如果 JSON 中的 shortName 是家族名，优先使用文件名
                                if !short_name.contains('-') && base_stem.contains('-') && model_family(&short_name) == model_family(&base_stem) {
                                    short_name = base_stem.clone();
                                }

                                // 如果还是没有短名，直接用文件名（去掉.json）
                                if short_name.is_empty() {
                                    short_name = base_stem;
                                }

                                let display_name = extract_display_name(&content).or_else(|| {
                                    // 如果没有 displayName，尝试从文件名生成友好名称
                                    Some(format!("{} Model", short_name.to_uppercase()))
                                });

                                let model_type = determine_model_type(&short_name);
                                // 改为家族匹配，兼容 chr-v2 / chr-2 / chr 这类命名差异
                                let is_downloaded = downloaded_models.iter().any(|d| model_name_equal(d, &short_name));

                                // 调试输出
                                println!("Found model: {} (file: {}, downloaded: {})", short_name, filename, is_downloaded);

                                models.push(ModelInfo {
                                    short_name,
                                    display_name,
                                    model_type,
                                    is_downloaded,
                                });
                            }
                        }
                    }
                }
            }
        }

        // 去重并排序
        models.sort_by(|a, b| a.short_name.cmp(&b.short_name));
        models.dedup_by(|a, b| a.short_name == b.short_name);

        println!("Total models found: {}", models.len());
        for model in &models {
            println!("  - {} ({}): downloaded={}", model.short_name,
                     model.display_name.as_deref().unwrap_or("N/A"), model.is_downloaded);
        }

        Ok(models)
    }

    /// 获取已下载的模型列表
    pub fn get_downloaded_models(&self) -> Result<Vec<String>, TvaiError> {
        let mut downloaded = Vec::new();

        if let Ok(entries) = fs::read_dir(&self.models_dir) {
            for entry in entries {
                if let Ok(entry) = entry {
                    let filename = entry.file_name().to_string_lossy().to_string();
                    if filename.ends_with(".tz") {
                        // 旧逻辑依赖下划线，实际文件名多为连字符，改为更稳健的提取：
                        // 1) 去扩展名，取 stem
                        // 2) 取家族前缀（连字符前第一段），如 "chr-v2-fp32-..." -> "chr"
                        if let Some(stem) = Path::new(&filename).file_stem().and_then(|s| s.to_str()) {
                            let family = model_family(stem).to_string();
                            if !downloaded.contains(&family) {
                                downloaded.push(family);
                            }
                        }
                    }
                }
            }
        }

        downloaded.sort();
        downloaded.dedup();
        Ok(downloaded)
    }

    /// 获取缺失的模型列表
    pub fn get_missing_models(&self) -> Result<Vec<ModelInfo>, TvaiError> {
        let all_models = self.get_all_models()?;
        Ok(all_models.into_iter()
            .filter(|m| !m.is_downloaded)
            .collect())
    }

    /// 检查特定模型是否已下载
    pub fn is_model_downloaded(&self, model_name: &str) -> Result<bool, TvaiError> {
        let downloaded = self.get_downloaded_models()?;
        Ok(downloaded.iter().any(|d| model_name_equal(d, model_name)))
    }

    /// 生成模型下载指南
    pub fn generate_download_guide(&self, output_path: &Path) -> Result<(), TvaiError> {
        let missing_models = self.get_missing_models()?;
        
        if missing_models.is_empty() {
            return Ok(());
        }

        let mut guide = String::new();
        guide.push_str("# Topaz Video AI 模型下载指南\n\n");
        
        // 统计信息
        let total_models = self.get_all_models()?.len();
        let downloaded_count = total_models - missing_models.len();
        
        guide.push_str(&format!("## 📊 模型状态\n\n"));
        guide.push_str(&format!("- 总计: {} 个模型\n", total_models));
        guide.push_str(&format!("- 已下载: {} 个模型\n", downloaded_count));
        guide.push_str(&format!("- 缺失: {} 个模型\n\n", missing_models.len()));

        // 按类型分组
        let upscale_models: Vec<_> = missing_models.iter()
            .filter(|m| m.model_type == ModelType::Upscale)
            .collect();
        
        let interpolation_models: Vec<_> = missing_models.iter()
            .filter(|m| m.model_type == ModelType::Interpolation)
            .collect();

        let other_models: Vec<_> = missing_models.iter()
            .filter(|m| m.model_type == ModelType::Other)
            .collect();

        if !upscale_models.is_empty() {
            guide.push_str("## 🔍 超分辨率模型 (Enhancement)\n\n");
            for model in upscale_models {
                guide.push_str(&format!("- **{}** - {}\n", 
                    model.short_name, 
                    model.display_name.as_deref().unwrap_or("N/A")));
            }
            guide.push_str("\n");
        }
        
        if !interpolation_models.is_empty() {
            guide.push_str("## 🎬 帧插值模型 (Frame Interpolation)\n\n");
            for model in interpolation_models {
                guide.push_str(&format!("- **{}** - {}\n", 
                    model.short_name, 
                    model.display_name.as_deref().unwrap_or("N/A")));
            }
            guide.push_str("\n");
        }

        if !other_models.is_empty() {
            guide.push_str("## 🔧 其他模型\n\n");
            for model in other_models {
                guide.push_str(&format!("- **{}** - {}\n", 
                    model.short_name, 
                    model.display_name.as_deref().unwrap_or("N/A")));
            }
            guide.push_str("\n");
        }
        
        guide.push_str("## 🚀 下载步骤\n\n");
        guide.push_str("1. 启动 Topaz Video AI 应用程序\n");
        guide.push_str("2. 导入任意视频文件\n");
        guide.push_str("3. 在右侧面板中选择相应的处理类型\n");
        guide.push_str("4. 依次选择上述列出的每个模型\n");
        guide.push_str("5. 点击 Preview 或 Export，应用程序会自动下载模型\n");
        guide.push_str("6. 等待下载完成后取消处理\n");
        guide.push_str("7. 重复步骤4-6直到所有模型下载完成\n\n");
        
        guide.push_str("## 💡 提示\n\n");
        guide.push_str("- 模型文件较大，请确保网络连接稳定\n");
        guide.push_str("- 下载过程中不要关闭应用程序\n");
        guide.push_str("- 某些模型可能需要特定的输入分辨率才能触发下载\n");
        guide.push_str("- 下载完成后可以在模型目录中看到 .tz 文件\n");

        if let Some(topaz_exe) = &self.topaz_exe {
            guide.push_str(&format!("\n## 🚀 启动命令\n\n"));
            guide.push_str(&format!("```\n\"{}\"\n```\n", topaz_exe.display()));
        }

        fs::write(output_path, guide)?;

        Ok(())
    }

    /// 启动 Topaz Video AI 应用程序
    pub fn launch_topaz_app(&self) -> Result<(), TvaiError> {
        if let Some(topaz_exe) = &self.topaz_exe {
            Command::new(topaz_exe)
                .spawn()
                .map_err(|e| TvaiError::ProcessingError(format!("Failed to launch Topaz: {}", e)))?;
            Ok(())
        } else {
            Err(TvaiError::TopazNotFound("Topaz Video AI executable not found".to_string()))
        }
    }

    /// 尝试触发模型下载 (实验性功能)
    pub fn attempt_model_download(&self, model_name: &str, test_video: &Path) -> Result<bool, TvaiError> {
        if !test_video.exists() {
            return Err(TvaiError::FileNotFound(test_video.to_string_lossy().to_string()));
        }

        let output_file = format!("temp_download_test_{}.mp4", model_name);
        
        // 构建测试命令
        let filter = format!("tvai_up=model={}:scale=2:device=-2:vram=0.1:estimate=1", model_name);
        
        let args = vec![
            "-y", "-hide_banner",
            "-i", test_video.to_str().unwrap(),
            "-t", "0.1",  // 只处理 0.1 秒
            "-filter_complex", &filter,
            "-c:v", "h264_nvenc",
            "-preset", "p1",
            "-an",
            &output_file,
        ];

        let output = Command::new(&self.topaz_ffmpeg)
            .args(&args)
            .output()
            .map_err(|e| TvaiError::ProcessingError(format!("Failed to execute FFmpeg: {}", e)))?;

        // 清理临时文件
        let _ = fs::remove_file(&output_file);

        let stderr = String::from_utf8_lossy(&output.stderr);
        
        // 检查结果
        if stderr.contains("Model not found") {
            Ok(false)
        } else if stderr.contains("Downloading") || 
                 stderr.contains("Loading") ||
                 output.status.success() {
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// 获取推荐的模型列表 (基于用途)
    pub fn get_recommended_models(&self, use_case: &str) -> Result<Vec<String>, TvaiError> {
        let models = self.get_all_models()?;

        let recommended = match use_case.to_lowercase().as_str() {
            "general" | "通用" => vec!["iris", "amq", "chr"],
            "high_quality" | "高质量" => vec!["ahq", "nyx", "apo"],
            "fast" | "快速" => vec!["alq", "apf", "chf"],
            "gaming" | "游戏" => vec!["ghq", "gcg"],
            "old_video" | "老视频" => vec!["prob", "dtv", "ddv"],
            "portrait" | "人像" => vec!["nyx", "nxf"],
            _ => vec!["iris", "amq", "chr"], // 默认推荐
        };

        // 过滤出实际存在的模型
        let available_recommended: Vec<String> = recommended.into_iter()
            .filter(|name| models.iter().any(|m| m.short_name == *name))
            .map(|s| s.to_string())
            .collect();

        Ok(available_recommended)
    }

    /// 校验所需模型是否可用，并给出替代建议
    pub fn validate_models(&self, required_models: &[&str]) -> Result<ValidationReport, TvaiError> {
        let all = self.get_all_models()?; // 全部（基于 json）
        let downloaded = self.get_downloaded_models()?; // 已下载（基于 tz）

        let required: Vec<String> = required_models.iter().map(|s| s.to_string()).collect();
        let mut available = Vec::new();
        let mut missing = Vec::new();

        for m in &required {
            if downloaded.iter().any(|d| model_name_equal(d, m)) {
                available.push(m.clone());
            } else {
                missing.push(m.clone());
            }
        }

        let mut suggestions = std::collections::HashMap::new();
        for miss in &missing {
            let suggest = self.suggest_alternatives(miss, &all, &downloaded);
            if !suggest.is_empty() {
                suggestions.insert(miss.clone(), suggest);
            }
        }

        Ok(ValidationReport { required, available, missing, suggestions })
    }

    /// 为缺失模型提供替代建议（基于家族前缀或同类型）
    fn suggest_alternatives(&self, model_name: &str, all: &[ModelInfo], downloaded: &[String]) -> Vec<String> {
        let family = model_family(model_name);
        let mtype = determine_model_type(model_name);

        // downloaded 是“家族前缀”集合，这里需要返回可直接用于命令的 short_name
        let mut candidates: Vec<String> = Vec::new();

        // 1) 同家族且已下载的 short_name 优先
        for m in all {
            if model_family(&m.short_name) == family && downloaded.iter().any(|d| model_family(d) == model_family(&m.short_name)) {
                candidates.push(m.short_name.clone());
            }
        }

        // 2) 若无同家族，按类型选择常见可替代族，并映射为对应 short_name
        if candidates.is_empty() {
            let preferred_families: &[&str] = match mtype {
                ModelType::Interpolation => &["apo", "apf", "chr", "chf"],
                ModelType::Upscale => &["iris", "nyx", "ahq", "prob", "thf"],
                ModelType::Other => &[],
            };
            for fam in preferred_families {
                for m in all {
                    if model_family(&m.short_name) == *fam && downloaded.iter().any(|d| model_family(d) == *fam) {
                        candidates.push(m.short_name.clone());
                    }
                }
            }
        }

        // 去重并限制数量
        candidates.sort();
        candidates.dedup();
        candidates.truncate(5);
        candidates
    }
}

/// 校验结果报告
#[derive(Debug, Clone)]
pub struct ValidationReport {
    pub required: Vec<String>,
    pub available: Vec<String>,
    pub missing: Vec<String>,
    pub suggestions: std::collections::HashMap<String, Vec<String>>,
}

// 辅助函数
fn extract_short_name(content: &str) -> Option<String> {
    if let Some(start) = content.find("\"shortName\"") {
        let after_key = &content[start..];
        if let Some(colon_pos) = after_key.find(':') {
            let after_colon = &after_key[colon_pos + 1..];
            if let Some(quote_start) = after_colon.find('"') {
                let after_quote = &after_colon[quote_start + 1..];
                if let Some(quote_end) = after_quote.find('"') {
                    return Some(after_quote[..quote_end].to_string());
                }
            }
        }
    }
    None
}

fn extract_display_name(content: &str) -> Option<String> {
    if let Some(start) = content.find("\"displayName\"") {
        let after_key = &content[start..];
        if let Some(colon_pos) = after_key.find(':') {
            let after_colon = &after_key[colon_pos + 1..];
            if let Some(quote_start) = after_colon.find('"') {
                let after_quote = &after_colon[quote_start + 1..];
                if let Some(quote_end) = after_quote.find('"') {
                    return Some(after_quote[..quote_end].to_string());
                }
            }
        }
    }
    None
}

fn model_family(name: &str) -> &str {
    // 取短横线前的族名前缀，例如 iris-3 => iris, apo-8 => apo
    name.split('-').next().unwrap_or(name)
}

fn model_name_equal(installed: &str, required: &str) -> bool {
    // 简化比较：只要家族前缀和主要版本前缀相同即可认为可替代
    // 例如 required=iris-3，installed=iris -> 也视为满足（有些 tz 命名可能不带版本）
    let fi = model_family(installed);
    let fr = model_family(required);
    fi == fr || installed == required
}

fn determine_model_type(short_name: &str) -> ModelType {
    if short_name.starts_with("ap") || short_name.starts_with("ch") || short_name == "ifi" {
        ModelType::Interpolation
    } else if short_name.starts_with("cpe") || short_name.starts_with("ref") ||
              short_name.starts_with("prap") || short_name.starts_with("nap") {
        ModelType::Other
    } else {
        ModelType::Upscale
    }
}
