//! Model Registry — GPT 模型元数据管理 + 权重存储 (Phase-69: T69.1 + T69.2 + T69.6)
//!
//! 提供模型的完整生命周期管理：
//! - 创建/删除/列出模型 (T69.1)
//! - 保存/加载 checkpoint — safetensors 格式 + 版本化 (T69.2)
//! - 端到端训练管线 — text → tokenize → train → checkpoint (T69.6)
//! - 文本生成推理
//!
//! ## 目录结构
//!
//! ```text
//! base_dir/
//! ├── model_a/
//! │   ├── model_meta.json
//! │   ├── checkpoint_v1.safetensors
//! │   └── checkpoint_v2.safetensors
//! └── model_b/
//!     ├── model_meta.json
//!     └── checkpoint_v1.safetensors
//! ```
//!
//! ## SQL 使用示例
//!
//! ```sql
//! SELECT gpt_create('my_model', '{"n_embd":64,"n_head":4,"n_layer":2}');
//! SELECT gpt_train('my_model', 'training text...', 100, 0.001);
//! SELECT gpt_generate('my_model', 'Hello', 50, 0.8);
//! SELECT gpt_status('my_model');
//! ```

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};

use candle_core::Device;
use serde::{Deserialize, Serialize};

use super::{
    GenerateConfig, GptConfig, GptTrainer, MicroGpt, TokenizerManager, TrainConfig,
    TrainingReport,
};
use crate::device::DeviceSelector;

// ============================================================================
// Types
// ============================================================================

/// 模型状态
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ModelStatus {
    /// 已创建但未训练
    Created,
    /// 正在训练中
    Training {
        current_step: usize,
        current_loss: f64,
    },
    /// 训练完成，可用于推理
    Ready {
        final_loss: f64,
        total_steps: usize,
    },
    /// 训练或加载失败
    Failed {
        error: String,
    },
}

impl std::fmt::Display for ModelStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Created => write!(f, "created"),
            Self::Training {
                current_step,
                current_loss,
            } => write!(f, "training(step={}, loss={:.4})", current_step, current_loss),
            Self::Ready {
                final_loss,
                total_steps,
            } => write!(f, "ready(loss={:.4}, steps={})", final_loss, total_steps),
            Self::Failed { error } => write!(f, "failed({})", error),
        }
    }
}

/// Checkpoint 元数据
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointInfo {
    /// 版本号（从 1 开始递增）
    pub version: usize,
    /// 保存时的损失值
    pub loss: f64,
    /// 保存时的训练步数
    pub step: usize,
    /// safetensors 文件名
    pub filename: String,
    /// 创建时间戳
    pub created_at: u64,
    /// 文件大小（字节）
    pub size_bytes: u64,
}

/// 模型元数据
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMeta {
    /// 模型名称（唯一标识符）
    pub name: String,
    /// GPT 配置
    pub config: GptConfig,
    /// 当前状态
    pub status: ModelStatus,
    /// 创建时间戳
    pub created_at: u64,
    /// 最后更新时间戳
    pub updated_at: u64,
    /// Checkpoint 列表
    pub checkpoints: Vec<CheckpointInfo>,
    /// 累计训练步数
    pub total_training_steps: usize,
    /// 历史最佳损失
    pub best_loss: Option<f64>,
}

/// Registry 错误类型
#[derive(Debug, thiserror::Error)]
pub enum RegistryError {
    #[error("Model not found: {0}")]
    ModelNotFound(String),
    #[error("Model already exists: {0}")]
    ModelAlreadyExists(String),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Candle error: {0}")]
    Candle(String),
    #[error("Serialization error: {0}")]
    Serde(#[from] serde_json::Error),
    #[error("Training error: {0}")]
    Training(String),
}

impl From<candle_core::Error> for RegistryError {
    fn from(e: candle_core::Error) -> Self {
        Self::Candle(e.to_string())
    }
}

pub type RegistryResult<T> = Result<T, RegistryError>;

// ============================================================================
// LoadedModel — wraps MicroGpt + tokenizer + metadata
// ============================================================================

/// 已加载的模型实例（包含模型、tokenizer、元数据）
pub struct LoadedModel {
    /// MicroGPT 模型实例
    pub model: MicroGpt,
    /// Tokenizer 管理器
    pub tokenizer: TokenizerManager,
    /// 模型元数据快照
    pub meta: ModelMeta,
}

// ============================================================================
// ModelRegistry
// ============================================================================

/// 当前时间戳（秒）
fn now_ts() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Model Registry — 管理 GPT 模型的完整生命周期
///
/// 线程安全：所有可变状态通过 `Mutex` 保护，可安全跨线程共享。
/// 持久化：模型元数据保存为 JSON，权重保存为 safetensors。
pub struct ModelRegistry {
    /// 模型存储根目录
    base_dir: PathBuf,
    /// 模型元数据缓存
    models: Mutex<HashMap<String, ModelMeta>>,
    /// 已加载的模型实例缓存（延迟加载）
    instances: Mutex<HashMap<String, Arc<LoadedModel>>>,
}

impl ModelRegistry {
    /// 创建 ModelRegistry
    ///
    /// # Arguments
    /// * `base_dir` - 模型存储根目录（不存在则自动创建）
    pub fn new<P: AsRef<Path>>(base_dir: P) -> RegistryResult<Self> {
        let base_dir = base_dir.as_ref().to_path_buf();
        std::fs::create_dir_all(&base_dir)?;

        let registry = Self {
            base_dir,
            models: Mutex::new(HashMap::new()),
            instances: Mutex::new(HashMap::new()),
        };

        // 扫描磁盘上已有的模型
        registry.scan_models()?;
        Ok(registry)
    }

    /// 扫描 base_dir 目录，加载已有模型的元数据
    fn scan_models(&self) -> RegistryResult<()> {
        let mut models = self.models.lock().unwrap();

        if let Ok(entries) = std::fs::read_dir(&self.base_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_dir() {
                    let meta_path = path.join("model_meta.json");
                    if meta_path.exists() {
                        if let Ok(content) = std::fs::read_to_string(&meta_path) {
                            if let Ok(meta) = serde_json::from_str::<ModelMeta>(&content) {
                                eprintln!(
                                    "[ModelRegistry] loaded model '{}' ({})",
                                    meta.name, meta.status
                                );
                                models.insert(meta.name.clone(), meta);
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// 获取模型存储目录
    fn model_dir(&self, name: &str) -> PathBuf {
        self.base_dir.join(name)
    }

    /// 将模型元数据保存到磁盘
    fn persist_meta(&self, meta: &ModelMeta) -> RegistryResult<()> {
        let dir = self.model_dir(&meta.name);
        std::fs::create_dir_all(&dir)?;
        let meta_path = dir.join("model_meta.json");
        let content = serde_json::to_string_pretty(meta)?;
        std::fs::write(meta_path, content)?;
        Ok(())
    }

    // ========================================================================
    // T69.1 — Model CRUD
    // ========================================================================

    /// 创建新模型
    ///
    /// 仅注册元数据，不创建模型实例（延迟到首次使用）。
    pub fn create_model(&self, name: &str, config: GptConfig) -> RegistryResult<ModelMeta> {
        let mut models = self.models.lock().unwrap();
        if models.contains_key(name) {
            return Err(RegistryError::ModelAlreadyExists(name.to_string()));
        }

        let now = now_ts();
        let meta = ModelMeta {
            name: name.to_string(),
            config,
            status: ModelStatus::Created,
            created_at: now,
            updated_at: now,
            checkpoints: Vec::new(),
            total_training_steps: 0,
            best_loss: None,
        };

        self.persist_meta(&meta)?;
        models.insert(name.to_string(), meta.clone());

        eprintln!(
            "[ModelRegistry] created model '{}' (~{}K params)",
            name,
            meta.config.estimate_params() / 1000
        );
        Ok(meta)
    }

    /// 获取模型元数据
    pub fn get_model(&self, name: &str) -> RegistryResult<ModelMeta> {
        let models = self.models.lock().unwrap();
        models
            .get(name)
            .cloned()
            .ok_or_else(|| RegistryError::ModelNotFound(name.to_string()))
    }

    /// 列出所有模型
    pub fn list_models(&self) -> Vec<ModelMeta> {
        let models = self.models.lock().unwrap();
        let mut list: Vec<_> = models.values().cloned().collect();
        list.sort_by(|a, b| a.created_at.cmp(&b.created_at));
        list
    }

    /// 删除模型及其所有 checkpoint
    pub fn delete_model(&self, name: &str) -> RegistryResult<()> {
        let mut models = self.models.lock().unwrap();
        if !models.contains_key(name) {
            return Err(RegistryError::ModelNotFound(name.to_string()));
        }

        // 清除实例缓存
        self.instances.lock().unwrap().remove(name);
        // 清除元数据
        models.remove(name);
        // 删除磁盘文件
        let dir = self.model_dir(name);
        if dir.exists() {
            std::fs::remove_dir_all(dir)?;
        }

        eprintln!("[ModelRegistry] deleted model '{}'", name);
        Ok(())
    }

    // ========================================================================
    // T69.2 — Checkpoint Management (safetensors)
    // ========================================================================

    /// 保存模型 checkpoint
    ///
    /// 将 VarMap 中所有参数序列化为 safetensors 格式并保存到磁盘。
    /// 自动递增版本号，更新元数据中的 checkpoints 列表。
    pub fn save_checkpoint(
        &self,
        name: &str,
        model: &MicroGpt,
        loss: f64,
        step: usize,
    ) -> RegistryResult<CheckpointInfo> {
        let mut models = self.models.lock().unwrap();
        let meta = models
            .get_mut(name)
            .ok_or_else(|| RegistryError::ModelNotFound(name.to_string()))?;

        let version = meta.checkpoints.len() + 1;
        let filename = format!("checkpoint_v{}.safetensors", version);
        let filepath = self.model_dir(name).join(&filename);

        // 使用 VarMap::save 将权重序列化为 safetensors
        model.varmap().save(&filepath)?;

        // 获取文件大小
        let size_bytes = std::fs::metadata(&filepath).map(|m| m.len()).unwrap_or(0);

        let checkpoint = CheckpointInfo {
            version,
            loss,
            step,
            filename,
            created_at: now_ts(),
            size_bytes,
        };

        meta.checkpoints.push(checkpoint.clone());
        meta.total_training_steps = step;
        meta.updated_at = now_ts();

        // 更新历史最佳损失
        if meta.best_loss.map_or(true, |best| loss < best) {
            meta.best_loss = Some(loss);
        }

        self.persist_meta(meta)?;

        eprintln!(
            "[ModelRegistry] saved checkpoint '{}' v{} (loss={:.4}, step={}, {:.1}KB)",
            name,
            checkpoint.version,
            loss,
            step,
            size_bytes as f64 / 1024.0
        );

        Ok(checkpoint)
    }

    /// 加载模型实例（含 checkpoint 恢复）
    ///
    /// 1. 按 config 创建 MicroGpt（随机初始化 VarMap）
    /// 2. 从 safetensors 加载权重更新 VarMap
    /// 3. 缓存实例以供后续使用
    ///
    /// # Arguments
    /// * `name` - 模型名称
    /// * `version` - checkpoint 版本（None = 最新版）
    /// * `device` - 计算设备
    pub fn load_model(
        &self,
        name: &str,
        version: Option<usize>,
        device: &Device,
    ) -> RegistryResult<Arc<LoadedModel>> {
        // 检查实例缓存（仅最新版本）
        if version.is_none() {
            let instances = self.instances.lock().unwrap();
            if let Some(instance) = instances.get(name) {
                return Ok(Arc::clone(instance));
            }
        }

        let meta = self.get_model(name)?;

        // 创建模型实例（随机初始化）
        let model = MicroGpt::new(&meta.config, device)?;

        // 如果有 checkpoint，加载权重
        if !meta.checkpoints.is_empty() {
            let checkpoint = if let Some(ver) = version {
                meta.checkpoints
                    .iter()
                    .find(|c| c.version == ver)
                    .ok_or_else(|| {
                        RegistryError::ModelNotFound(format!("{}@v{}", name, ver))
                    })?
            } else {
                meta.checkpoints.last().unwrap()
            };

            let filepath = self.model_dir(name).join(&checkpoint.filename);

            // 从 safetensors 加载张量并更新 VarMap 中的变量
            let tensors = candle_core::safetensors::load(&filepath, device)?;
            let var_data = model.varmap().data().lock().unwrap();
            for (var_name, var) in var_data.iter() {
                if let Some(tensor) = tensors.get(var_name) {
                    var.set(tensor).map_err(|e| {
                        RegistryError::Candle(format!("set var '{}': {}", var_name, e))
                    })?;
                }
            }
            drop(var_data);

            eprintln!(
                "[ModelRegistry] loaded '{}' v{} → {} (loss={:.4})",
                name,
                checkpoint.version,
                if device.is_cuda() { "GPU" } else { "CPU" },
                checkpoint.loss,
            );
        } else {
            eprintln!(
                "[ModelRegistry] loaded '{}' (no checkpoint, random weights)",
                name
            );
        }

        let tokenizer = TokenizerManager::char_level();

        let loaded = Arc::new(LoadedModel {
            model,
            tokenizer,
            meta: meta.clone(),
        });

        // 缓存最新版本
        if version.is_none() {
            self.instances
                .lock()
                .unwrap()
                .insert(name.to_string(), Arc::clone(&loaded));
        }

        Ok(loaded)
    }

    /// 清除模型实例缓存（下次使用时重新加载）
    pub fn evict_model(&self, name: &str) {
        self.instances.lock().unwrap().remove(name);
    }

    // ========================================================================
    // T69.6 — Training Pipeline (text → tokenize → train → checkpoint)
    // ========================================================================

    /// 端到端训练管线
    ///
    /// 完整流程：加载模型 → tokenize → 训练循环 → 定期 checkpoint → 更新状态
    ///
    /// # Arguments
    /// * `name` - 模型名称
    /// * `text_data` - 训练文本
    /// * `num_steps` - 训练步数
    /// * `train_config` - 训练超参数
    /// * `batch_size` - 批大小
    /// * `checkpoint_every` - 每隔多少步保存 checkpoint（None = 仅最终保存）
    pub fn train_model(
        &self,
        name: &str,
        text_data: &str,
        num_steps: usize,
        train_config: TrainConfig,
        batch_size: usize,
        checkpoint_every: Option<usize>,
    ) -> RegistryResult<TrainingReport> {
        let selector = DeviceSelector::auto();
        let (device, is_gpu) = selector.candle_device_with_fallback("gpt_train");

        eprintln!(
            "[ModelRegistry] training '{}' on {} ({} steps, batch={})",
            name,
            if is_gpu { "GPU" } else { "CPU" },
            num_steps,
            batch_size,
        );

        // 加载或创建模型实例
        // 先 evict 以确保不使用缓存的只读实例
        self.evict_model(name);
        let loaded = self.load_model(name, None, &device)?;

        // Tokenize 文本
        let tokens = loaded
            .tokenizer
            .encode(text_data)
            .map_err(RegistryError::Training)?;

        let block_size = loaded.model.config().block_size;
        if tokens.len() < block_size + 1 {
            return Err(RegistryError::Training(format!(
                "Text too short: {} tokens, need >= {}",
                tokens.len(),
                block_size + 1
            )));
        }

        // 更新状态为训练中
        {
            let mut models = self.models.lock().unwrap();
            if let Some(meta) = models.get_mut(name) {
                meta.status = ModelStatus::Training {
                    current_step: 0,
                    current_loss: 0.0,
                };
                meta.updated_at = now_ts();
                let _ = self.persist_meta(meta);
            }
        }

        // 创建训练器
        let mut trainer = GptTrainer::new(&loaded.model, train_config)
            .map_err(|e| RegistryError::Candle(e.to_string()))?;

        let checkpoint_interval = checkpoint_every.unwrap_or(num_steps);
        let mut loss_history = Vec::with_capacity(num_steps);
        let start = std::time::Instant::now();
        let mut total_tokens = 0usize;

        // 训练循环
        for step in 0..num_steps {
            let (input_ids, target_ids) = GptTrainer::prepare_batch(
                &tokens, block_size, batch_size, &device,
            )
            .map_err(|e| RegistryError::Candle(e.to_string()))?;

            let result = trainer
                .train_step(&loaded.model, &input_ids, &target_ids)
                .map_err(|e| RegistryError::Candle(e.to_string()))?;

            loss_history.push(result.loss);
            total_tokens += batch_size * block_size;

            // 每 10 步更新训练状态
            if (step + 1) % 10 == 0 || step == num_steps - 1 {
                let mut models = self.models.lock().unwrap();
                if let Some(meta) = models.get_mut(name) {
                    meta.status = ModelStatus::Training {
                        current_step: step + 1,
                        current_loss: result.loss,
                    };
                    meta.updated_at = now_ts();
                }
            }

            // 定期保存 checkpoint
            if (step + 1) % checkpoint_interval == 0 && step < num_steps - 1 {
                let _ = self.save_checkpoint(name, &loaded.model, result.loss, step + 1);
            }
        }

        let elapsed = start.elapsed().as_secs_f64();
        let final_loss = *loss_history.last().unwrap_or(&0.0);
        let initial_loss = *loss_history.first().unwrap_or(&0.0);

        // 保存最终 checkpoint
        let base_steps = loaded.meta.total_training_steps;
        self.save_checkpoint(name, &loaded.model, final_loss, base_steps + num_steps)?;

        // 更新状态为 Ready
        {
            let mut models = self.models.lock().unwrap();
            if let Some(meta) = models.get_mut(name) {
                meta.status = ModelStatus::Ready {
                    final_loss,
                    total_steps: meta.total_training_steps,
                };
                meta.updated_at = now_ts();
                let _ = self.persist_meta(meta);
            }
        }

        // 清除实例缓存（权重已变更）
        self.evict_model(name);

        let report = TrainingReport {
            initial_loss,
            final_loss,
            total_steps: num_steps,
            loss_history,
            elapsed_secs: elapsed,
            avg_tokens_per_sec: total_tokens as f64 / elapsed,
        };

        eprintln!(
            "[ModelRegistry] training '{}' complete: loss {:.4} → {:.4} ({:.0} tok/s)",
            name, report.initial_loss, report.final_loss, report.avg_tokens_per_sec
        );

        Ok(report)
    }

    /// 使用训练好的模型生成文本
    pub fn generate(
        &self,
        name: &str,
        prompt: &str,
        gen_config: &GenerateConfig,
    ) -> RegistryResult<String> {
        let selector = DeviceSelector::auto();
        let (device, _) = selector.candle_device_with_fallback("gpt_generate");

        let loaded = self.load_model(name, None, &device)?;

        let text = loaded
            .model
            .generate_text(prompt, &loaded.tokenizer, gen_config)
            .map_err(|e| RegistryError::Candle(e.to_string()))?;

        Ok(text)
    }

    /// 获取存储根目录
    pub fn base_dir(&self) -> &Path {
        &self.base_dir
    }

    /// 获取注册的模型数量
    pub fn model_count(&self) -> usize {
        self.models.lock().unwrap().len()
    }
}

// ============================================================================
// Global Registry (singleton)
// ============================================================================

static GLOBAL_REGISTRY: OnceLock<Arc<ModelRegistry>> = OnceLock::new();

/// 获取全局 ModelRegistry 单例
///
/// 默认存储目录为 `.aifastdb/models`，可通过环境变量 `AIFASTDB_MODEL_DIR` 覆盖。
pub fn global_model_registry() -> &'static Arc<ModelRegistry> {
    GLOBAL_REGISTRY.get_or_init(|| {
        let base_dir = std::env::var("AIFASTDB_MODEL_DIR")
            .unwrap_or_else(|_| ".aifastdb/models".to_string());
        Arc::new(ModelRegistry::new(base_dir).expect("Failed to create global model registry"))
    })
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_create_model() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        let meta = registry
            .create_model("test_model", GptConfig::tiny())
            .unwrap();
        assert_eq!(meta.name, "test_model");
        assert_eq!(meta.status, ModelStatus::Created);
        assert_eq!(meta.checkpoints.len(), 0);

        // 重复创建应失败
        let result = registry.create_model("test_model", GptConfig::tiny());
        assert!(result.is_err());
    }

    #[test]
    fn test_list_and_get_models() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        registry
            .create_model("model_a", GptConfig::tiny())
            .unwrap();
        registry
            .create_model("model_b", GptConfig::default())
            .unwrap();

        let models = registry.list_models();
        assert_eq!(models.len(), 2);

        let meta = registry.get_model("model_a").unwrap();
        assert_eq!(meta.name, "model_a");

        assert!(registry.get_model("nonexistent").is_err());
    }

    #[test]
    fn test_delete_model() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        registry
            .create_model("to_delete", GptConfig::tiny())
            .unwrap();
        assert!(registry.get_model("to_delete").is_ok());

        registry.delete_model("to_delete").unwrap();
        assert!(registry.get_model("to_delete").is_err());

        // 目录应被删除
        assert!(!dir.path().join("to_delete").exists());
    }

    #[test]
    fn test_save_and_load_checkpoint() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();
        let config = GptConfig::tiny();
        registry.create_model("ckpt_test", config.clone()).unwrap();

        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        // 保存 checkpoint
        let ckpt = registry
            .save_checkpoint("ckpt_test", &model, 5.5, 100)
            .unwrap();
        assert_eq!(ckpt.version, 1);
        assert_eq!(ckpt.step, 100);
        assert!(ckpt.size_bytes > 0);

        // safetensors 文件应存在
        let ckpt_path = dir.path().join("ckpt_test").join(&ckpt.filename);
        assert!(ckpt_path.exists());

        // 加载模型
        let loaded = registry.load_model("ckpt_test", None, &device).unwrap();
        assert_eq!(loaded.meta.checkpoints.len(), 1);
        assert_eq!(loaded.model.num_params(), model.num_params());
    }

    #[test]
    fn test_multiple_checkpoints() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();
        let config = GptConfig::tiny();
        registry
            .create_model("multi_ckpt", config.clone())
            .unwrap();

        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        // 保存多个 checkpoint
        registry
            .save_checkpoint("multi_ckpt", &model, 5.0, 50)
            .unwrap();
        registry
            .save_checkpoint("multi_ckpt", &model, 3.0, 100)
            .unwrap();
        registry
            .save_checkpoint("multi_ckpt", &model, 2.0, 150)
            .unwrap();

        let meta = registry.get_model("multi_ckpt").unwrap();
        assert_eq!(meta.checkpoints.len(), 3);
        assert_eq!(meta.best_loss, Some(2.0));
        assert_eq!(meta.total_training_steps, 150);

        // 加载特定版本
        let loaded_v1 = registry
            .load_model("multi_ckpt", Some(1), &device)
            .unwrap();
        assert_eq!(loaded_v1.meta.checkpoints.len(), 3);
    }

    #[test]
    fn test_persistence_across_restarts() {
        let dir = tempdir().unwrap();

        // 第一次会话：创建模型 + 保存 checkpoint
        {
            let registry = ModelRegistry::new(dir.path()).unwrap();
            let config = GptConfig::tiny();
            registry
                .create_model("persist_test", config.clone())
                .unwrap();
            let model = MicroGpt::new(&config, &Device::Cpu).unwrap();
            registry
                .save_checkpoint("persist_test", &model, 3.0, 50)
                .unwrap();
            assert_eq!(registry.model_count(), 1);
        }

        // 第二次会话：重新扫描磁盘
        {
            let registry = ModelRegistry::new(dir.path()).unwrap();
            assert_eq!(registry.model_count(), 1);
            let meta = registry.get_model("persist_test").unwrap();
            assert_eq!(meta.name, "persist_test");
            assert_eq!(meta.checkpoints.len(), 1);
            assert_eq!(meta.checkpoints[0].step, 50);
        }
    }

    #[test]
    fn test_train_and_generate_pipeline() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        let config = GptConfig::tiny();
        registry.create_model("train_test", config).unwrap();

        // 训练
        let text = "abcdefgh".repeat(50);
        let train_config = TrainConfig {
            learning_rate: 1e-3,
            ..Default::default()
        };

        let report = registry
            .train_model("train_test", &text, 15, train_config, 2, None)
            .unwrap();

        assert_eq!(report.total_steps, 15);
        assert!(report.final_loss < report.initial_loss, "Loss should decrease");

        // 状态应为 Ready
        let meta = registry.get_model("train_test").unwrap();
        assert!(matches!(meta.status, ModelStatus::Ready { .. }));
        assert!(meta.checkpoints.len() >= 1);

        // 生成文本
        let gen_config = GenerateConfig {
            max_new_tokens: 20,
            temperature: 0.5,
            top_k: 10,
        };
        let generated = registry
            .generate("train_test", "abc", &gen_config)
            .unwrap();
        assert!(generated.starts_with("abc"));
        assert!(generated.len() > 3);
    }

    #[test]
    fn test_train_with_periodic_checkpoints() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        registry
            .create_model("periodic_ckpt", GptConfig::tiny())
            .unwrap();

        let text = "Hello World! ".repeat(50);
        let train_config = TrainConfig {
            learning_rate: 1e-3,
            ..Default::default()
        };

        // 每 10 步保存一次
        let report = registry
            .train_model("periodic_ckpt", &text, 25, train_config, 2, Some(10))
            .unwrap();

        assert_eq!(report.total_steps, 25);

        let meta = registry.get_model("periodic_ckpt").unwrap();
        // 应有 checkpoint: step 10, step 20, step 25 (final)
        assert!(
            meta.checkpoints.len() >= 3,
            "Expected >= 3 checkpoints, got {}",
            meta.checkpoints.len()
        );
    }

    #[test]
    fn test_model_not_found_errors() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path()).unwrap();

        assert!(registry.get_model("nope").is_err());
        assert!(registry.delete_model("nope").is_err());
        assert!(registry
            .load_model("nope", None, &Device::Cpu)
            .is_err());
    }
}
