//! ModelManager — AI 模型资产管理 (M2.1)
//!
//! 管理 LLM/ASR/TTS 模型文件，支持清单查看、下载、切换、删除。
//! 参照项目已有 plugin_repo + version_manager 模式设计。
//!
//! # 模型存储结构
//! ```text
//! model_store/
//! ├── registry.json          # 本地模型注册表（已下载的模型清单 + active 标记）
//! ├── llm/
//! │   ├── qwen2.5-0.5b-q4_k_m.gguf
//! │   └── gemma3-1b-q4_k_m.gguf
//! ├── asr/
//! │   └── ggml-tiny.bin
//! └── tts/
//!     └── zh_CN-huayan-medium.onnx
//! ```

use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

/// 模型类型
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "lowercase")]
pub enum ModelKind {
    Llm,
    Asr,
    Tts,
}

impl ModelKind {
    pub fn dir_name(&self) -> &'static str {
        match self {
            ModelKind::Llm => "llm",
            ModelKind::Asr => "asr",
            ModelKind::Tts => "tts",
        }
    }
}

/// 模型元信息（仓库清单 + 本地注册表共用）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    /// 模型 ID (如 "qwen2.5-0.5b-q4_k_m")
    pub id: String,
    /// 模型类型
    pub kind: ModelKind,
    /// 显示名称 (如 "Qwen2.5 0.5B Q4")
    pub name: String,
    /// 版本
    pub version: String,
    /// 文件尺寸 (字节)
    pub size: u64,
    /// 运行内存需求 (字节)
    pub memory_required: u64,
    /// 推荐设备档位 (如 "4G 内存档可用")
    pub recommended_tier: String,
    /// 下载 URL（国内镜像源）
    pub url: String,
    /// SHA256 校验和
    pub sha256: String,
    /// 文件名（存储到本地的文件名）
    pub filename: String,
}

/// 本地注册表条目
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalModelEntry {
    #[serde(flatten)]
    pub info: ModelInfo,
    /// 是否已下载
    pub downloaded: bool,
    /// 下载进度 (0-100，100=完成)
    pub download_progress: u8,
}

/// 本地注册表
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ModelRegistry {
    /// 所有已知模型（仓库清单 + 本地状态）
    pub models: Vec<LocalModelEntry>,
    /// 各类型当前激活的模型 ID
    pub active: HashMap<ModelKind, String>,
}

/// 模型管理器
pub struct ModelManager {
    /// 模型存储根目录
    store_dir: PathBuf,
    /// 注册表（内存缓存，持久化到 registry.json）
    registry: ModelRegistry,
    /// 存储配额上限（字节，默认 4GB）
    quota: u64,
}

impl ModelManager {
    pub fn new(store_dir: PathBuf) -> Self {
        let quota = 4 * 1024 * 1024 * 1024; // 4GB
        let mut mgr = Self {
            store_dir,
            registry: ModelRegistry::default(),
            quota,
        };
        let _ = mgr.load_registry();
        mgr
    }

    /// 存储目录
    pub fn store_dir(&self) -> &Path {
        &self.store_dir
    }

    /// 配额
    pub fn quota(&self) -> u64 {
        self.quota
    }

    /// 已用空间
    pub fn used_space(&self) -> u64 {
        self.registry
            .models
            .iter()
            .filter(|m| m.downloaded)
            .map(|m| m.info.size)
            .sum()
    }

    /// 加载本地注册表
    fn load_registry(&mut self) -> Result<()> {
        let path = self.store_dir.join("registry.json");
        if path.exists() {
            let content = fs::read_to_string(&path)?;
            self.registry = serde_json::from_str(&content)?;
        }
        Ok(())
    }

    /// 持久化注册表
    fn save_registry(&self) -> Result<()> {
        fs::create_dir_all(&self.store_dir)?;
        let path = self.store_dir.join("registry.json");
        let content = serde_json::to_string_pretty(&self.registry)?;
        fs::write(&path, content)?;
        Ok(())
    }

    /// 初始化默认模型清单（首次启动或清单缺失时）
    /// 每次启动都重新检查文件存在性，更新 downloaded 状态（符号链接可能后建）
    pub fn ensure_default_models(&mut self) -> Result<()> {
        // 首次：注册表为空，填充默认清单
        if self.registry.models.is_empty() {
            let defaults = vec![
                ModelInfo {
                    id: "qwen2.5-0.5b-q4_k_m".to_string(),
                    kind: ModelKind::Llm,
                    name: "Qwen2.5 0.5B Q4 (推荐默认)".to_string(),
                    version: "0.5b".to_string(),
                    size: 469 * 1024 * 1024,
                    memory_required: 1024 * 1024 * 1024,
                    recommended_tier: "4G 内存档可用".to_string(),
                    url: "https://hf-mirror.com/Qwen/Qwen2.5-0.5B-Instruct-GGUF/resolve/main/qwen2.5-0.5b-instruct-q4_k_m.gguf".to_string(),
                    sha256: String::new(),
                    filename: "qwen2.5-0.5b-q4_k_m.gguf".to_string(),
                },
                       ModelInfo {
                id: "gemma3-1b-q4_k_m".to_string(),
                kind: ModelKind::Llm,
                name: "Gemma3 1B Q4 (可选档)".to_string(),
                version: "1b".to_string(),
                size: 800 * 1024 * 1024,
                memory_required: 2 * 1024 * 1024 * 1024,
                recommended_tier: "建议 8G 以上".to_string(),
                url: "https://hf-mirror.com/google/gemma-3-1b-it-GGUF/resolve/main/gemma-3-1b-it-q4_k_m.gguf".to_string(),
                sha256: String::new(),
                filename: "gemma3-1b-q4_k_m.gguf".to_string(),
            },
            ModelInfo {
                id: "gemma-4-e2b-q4_k_m".to_string(),
                kind: ModelKind::Llm,
                name: "Gemma 4 E2B Q4 (最新代,多模态)".to_string(),
                version: "e2b".to_string(),
                size: 3269 * 1024 * 1024,
                memory_required: 3 * 1024 * 1024 * 1024,
                recommended_tier: "需 4G+ 内存,可能触发 swap".to_string(),
                url: "https://huggingface.co/dahus/gemma-4-e2b-it-Q4_K_M-GGUF/resolve/main/gemma-4-e2b-Q4_K_M.gguf".to_string(),
                sha256: "2de6e5baaea40a8d54fb37ff418730ba4e5c002b".to_string(),
                filename: "gemma-4-e2b-Q4_K_M.gguf".to_string(),
            },
                ModelInfo {
                    id: "whisper-tiny".to_string(),
                    kind: ModelKind::Asr,
                    name: "Whisper Tiny (中文)".to_string(),
                    version: "tiny".to_string(),
                    size: 75 * 1024 * 1024,
                    memory_required: 200 * 1024 * 1024,
                    recommended_tier: "4G 内存档可用".to_string(),
                    url: "https://hf-mirror.com/ggerganov/whisper.cpp/resolve/main/ggml-tiny.bin".to_string(),
                    sha256: String::new(),
                    filename: "ggml-tiny.bin".to_string(),
                },
                ModelInfo {
                    id: "piper-zh-huayan".to_string(),
                    kind: ModelKind::Tts,
                    name: "Piper zh_CN huayan (中文女声)".to_string(),
                    version: "medium".to_string(),
                    size: 63 * 1024 * 1024,
                    memory_required: 150 * 1024 * 1024,
                    recommended_tier: "4G 内存档可用".to_string(),
                    url: "https://hf-mirror.com/rhasspy/piper-voices/resolve/main/zh/zh_CN/huayan/medium/zh_CN-huayan-medium.onnx".to_string(),
                    sha256: String::new(),
                    filename: "zh_CN-huayan-medium.onnx".to_string(),
                },
            ];

            for info in defaults {
                self.registry.models.push(LocalModelEntry {
                    info,
                    downloaded: false,
                    download_progress: 0,
                });
            }

            // 设置默认 active 模型
            if !self.registry.active.contains_key(&ModelKind::Llm) {
                self.registry.active.insert(ModelKind::Llm, "qwen2.5-0.5b-q4_k_m".to_string());
            }
            if !self.registry.active.contains_key(&ModelKind::Asr) {
                self.registry.active.insert(ModelKind::Asr, "whisper-tiny".to_string());
            }
            if !self.registry.active.contains_key(&ModelKind::Tts) {
                self.registry.active.insert(ModelKind::Tts, "piper-zh-huayan".to_string());
            }
        }

        // 每次启动都重新检查文件存在性（符号链接可能后建，或文件被手动删除）
        for entry in &mut self.registry.models {
            let path = self.store_dir.join(entry.info.kind.dir_name()).join(&entry.info.filename);
            let exists = path.exists();
            entry.downloaded = exists;
            entry.download_progress = if exists { 100 } else { 0 };
        }

        self.save_registry()?;
        Ok(())
    }

    /// 列出所有模型
    pub fn list_models(&self) -> &[LocalModelEntry] {
        &self.registry.models
    }

    /// 获取当前激活的模型 ID
    pub fn active_model_id(&self, kind: ModelKind) -> Option<&String> {
        self.registry.active.get(&kind)
    }

    /// 获取已激活 LLM 模型文件路径
    pub fn active_llm_model_path(&self) -> Result<PathBuf> {
        self.active_model_path(ModelKind::Llm)
    }

    /// 获取已激活 ASR 模型文件路径
    pub fn active_asr_model_path(&self) -> Result<PathBuf> {
        self.active_model_path(ModelKind::Asr)
    }

    /// 获取已激活 TTS 模型文件路径
    pub fn active_tts_model_path(&self) -> Result<PathBuf> {
        self.active_model_path(ModelKind::Tts)
    }

    fn active_model_path(&self, kind: ModelKind) -> Result<PathBuf> {
        let id = self
            .registry
            .active
            .get(&kind)
            .ok_or_else(|| anyhow::anyhow!("没有激活的 {kind:?} 模型"))?;
        let entry = self
            .registry
            .models
            .iter()
            .find(|m| &m.info.id == id)
            .ok_or_else(|| anyhow::anyhow!("激活的模型 {id} 不在注册表中"))?;
        if !entry.downloaded {
            bail!("激活的模型 {id} 尚未下载");
        }
        Ok(self.store_dir.join(kind.dir_name()).join(&entry.info.filename))
    }

    /// 下载模型
    pub fn download_model(&mut self, model_id: &str) -> Result<()> {
        let entry = self
            .registry
            .models
            .iter()
            .find(|m| m.info.id == model_id)
            .ok_or_else(|| anyhow::anyhow!("模型 {model_id} 不在清单中"))?
            .clone();

        if entry.downloaded {
            return Ok(()); // 已下载
        }

        // 配额检查
        if self.used_space() + entry.info.size > self.quota {
            bail!(
                "磁盘配额不足: 已用 {} + 需要 {} > 配额 {}",
                self.used_space(),
                entry.info.size,
                self.quota
            );
        }

        // 确保目录存在
        let dir = self.store_dir.join(entry.info.kind.dir_name());
        fs::create_dir_all(&dir)?;
        let dest = dir.join(&entry.info.filename);

        // 下载（简单实现，不带断点续传；生产环境应加）
        println!("[ModelManager] 下载模型 {} from {}", model_id, entry.info.url);
        let resp = ureq::get(&entry.info.url)
            .timeout(std::time::Duration::from_secs(600))
            .call()
            .context("下载模型失败")?;

        let mut file = fs::File::create(&dest)?;
        std::io::copy(&mut resp.into_reader(), &mut file)?;

        // 标记为已下载
        if let Some(m) = self.registry.models.iter_mut().find(|m| m.info.id == model_id) {
            m.downloaded = true;
            m.download_progress = 100;
        }
        self.save_registry()?;
        println!("[ModelManager] 模型 {} 下载完成", model_id);
        Ok(())
    }

    /// 切换激活模型
    pub fn switch_model(&mut self, kind: ModelKind, model_id: &str) -> Result<()> {
        let entry = self
            .registry
            .models
            .iter()
            .find(|m| m.info.id == model_id && m.info.kind == kind)
            .ok_or_else(|| anyhow::anyhow!("模型 {model_id} 不在 {kind:?} 清单中"))?;

        if !entry.downloaded {
            bail!("模型 {model_id} 尚未下载，无法切换");
        }

        let old = self.registry.active.insert(kind, model_id.to_string());
        self.save_registry()?;
        if let Some(old_id) = old {
            println!("[ModelManager] {kind:?} 模型切换: {old_id} → {model_id}");
        } else {
            println!("[ModelManager] {kind:?} 模型激活: {model_id}");
        }
        Ok(())
    }

    /// 删除模型（当前激活的模型不可删除）
    pub fn delete_model(&mut self, model_id: &str) -> Result<()> {
        let entry = self
            .registry
            .models
            .iter()
            .find(|m| m.info.id == model_id)
            .ok_or_else(|| anyhow::anyhow!("模型 {model_id} 不在清单中"))?;

        // 保护当前激活模型
        if let Some(active_id) = self.registry.active.get(&entry.info.kind) {
            if active_id == model_id {
                bail!("模型 {model_id} 正在使用中，无法删除（请先切换到其他模型）");
            }
        }

        if !entry.downloaded {
            bail!("模型 {model_id} 未下载，无需删除");
        }

        // 删除文件
        let path = self
            .store_dir
            .join(entry.info.kind.dir_name())
            .join(&entry.info.filename);
        if path.exists() {
            fs::remove_file(&path)?;
        }

        // 更新注册表
        if let Some(m) = self.registry.models.iter_mut().find(|m| m.info.id == model_id) {
            m.downloaded = false;
            m.download_progress = 0;
        }
        self.save_registry()?;
        println!("[ModelManager] 模型 {} 已删除", model_id);
        Ok(())
    }

    /// 获取注册表（用于 HTTP API 返回）
    pub fn registry(&self) -> &ModelRegistry {
        &self.registry
    }
}

/// 仓库清单索引（远程静态清单文件）
#[derive(Debug, Clone, Deserialize)]
pub struct RepoModelIndex {
    pub models: Vec<ModelInfo>,
}

impl ModelManager {
    /// 从远程仓库拉取最新清单（更新本地注册表，不触发下载）
    pub fn fetch_repo_index(repo_url: &str) -> Result<RepoModelIndex> {
        let resp = ureq::get(&format!("{repo_url}/models/index.json"))
            .timeout(std::time::Duration::from_secs(30))
            .call()
            .context("拉取模型清单失败")?;
        let index: RepoModelIndex = serde_json::from_reader(resp.into_reader())?;
        Ok(index)
    }
}
