pokeys_lib/
model_manager.rs

1//! Model management API
2//!
3//! This module provides a high-level API for managing device models.
4//! It allows users to list, create, edit, validate, and apply models to devices.
5
6use crate::error::{PoKeysError, Result};
7use crate::models::{
8    DeviceModel, PinModel, copy_default_models_to_user_dir, get_default_model_dir,
9};
10use log::{info, warn};
11use std::collections::HashMap;
12use std::fs;
13use std::path::{Path, PathBuf};
14
15/// Model manager for device models
16pub struct ModelManager {
17    /// Directory containing model files
18    model_dir: PathBuf,
19
20    /// Loaded models
21    models: HashMap<String, DeviceModel>,
22}
23
24impl ModelManager {
25    /// Create a new model manager
26    ///
27    /// # Arguments
28    ///
29    /// * `model_dir` - Optional custom directory for model files
30    ///
31    /// # Returns
32    ///
33    /// * `Result<ModelManager>` - The new model manager or an error
34    pub fn new(model_dir: Option<PathBuf>) -> Result<Self> {
35        let dir = model_dir.unwrap_or_else(get_default_model_dir);
36
37        // Create the directory if it doesn't exist
38        if !dir.exists() {
39            fs::create_dir_all(&dir).map_err(|e| {
40                PoKeysError::ModelDirCreateError(dir.to_string_lossy().to_string(), e.to_string())
41            })?;
42        }
43
44        // Copy default models to the user's directory
45        copy_default_models_to_user_dir(Some(&dir))?;
46
47        let mut manager = Self {
48            model_dir: dir,
49            models: HashMap::new(),
50        };
51
52        // Load all models from the directory
53        manager.reload_models()?;
54
55        Ok(manager)
56    }
57
58    /// Reload all models from the model directory
59    ///
60    /// # Returns
61    ///
62    /// * `Result<()>` - Ok if models were loaded successfully, an error otherwise
63    pub fn reload_models(&mut self) -> Result<()> {
64        self.models.clear();
65
66        // Read all YAML files in the directory
67        let entries = fs::read_dir(&self.model_dir).map_err(|e| {
68            PoKeysError::ModelDirReadError(
69                self.model_dir.to_string_lossy().to_string(),
70                e.to_string(),
71            )
72        })?;
73
74        for entry in entries {
75            let entry = entry.map_err(|e| {
76                PoKeysError::ModelDirReadError(
77                    self.model_dir.to_string_lossy().to_string(),
78                    e.to_string(),
79                )
80            })?;
81
82            let path = entry.path();
83
84            if path.extension().is_some_and(|ext| ext == "yaml") {
85                if let Some(file_name) = path.file_stem() {
86                    let model_name = file_name.to_string_lossy().to_string();
87
88                    match DeviceModel::from_file(&path) {
89                        Ok(model) => {
90                            self.models.insert(model_name, model);
91                        }
92                        Err(e) => {
93                            warn!("Failed to load model from {}: {}", path.display(), e);
94                        }
95                    }
96                }
97            }
98        }
99
100        info!(
101            "Loaded {} models from {}",
102            self.models.len(),
103            self.model_dir.display()
104        );
105        Ok(())
106    }
107
108    /// Get a model by name
109    ///
110    /// # Arguments
111    ///
112    /// * `name` - The name of the model
113    ///
114    /// # Returns
115    ///
116    /// * `Option<&DeviceModel>` - The model if found, None otherwise
117    pub fn get_model(&self, name: &str) -> Option<&DeviceModel> {
118        self.models.get(name)
119    }
120
121    /// Get a mutable model by name
122    ///
123    /// # Arguments
124    ///
125    /// * `name` - The name of the model
126    ///
127    /// # Returns
128    ///
129    /// * `Option<&mut DeviceModel>` - The model if found, None otherwise
130    pub fn get_model_mut(&mut self, name: &str) -> Option<&mut DeviceModel> {
131        self.models.get_mut(name)
132    }
133
134    /// Get all loaded models
135    ///
136    /// # Returns
137    ///
138    /// * `&HashMap<String, DeviceModel>` - Map of model names to models
139    pub fn get_all_models(&self) -> &HashMap<String, DeviceModel> {
140        &self.models
141    }
142
143    /// Create a new model
144    ///
145    /// # Arguments
146    ///
147    /// * `name` - The name of the model
148    /// * `pins` - Map of pin numbers to pin models
149    ///
150    /// # Returns
151    ///
152    /// * `Result<()>` - Ok if the model was created successfully, an error otherwise
153    pub fn create_model(&mut self, name: &str, pins: HashMap<u8, PinModel>) -> Result<()> {
154        // Create the model
155        let model = DeviceModel {
156            name: name.to_string(),
157            pins,
158        };
159
160        // Validate the model
161        model.validate()?;
162
163        // Save the model to a file
164        self.save_model(&model)?;
165
166        // Add the model to the map
167        self.models.insert(name.to_string(), model);
168
169        Ok(())
170    }
171
172    /// Save a model to a file
173    ///
174    /// # Arguments
175    ///
176    /// * `model` - The model to save
177    ///
178    /// # Returns
179    ///
180    /// * `Result<()>` - Ok if the model was saved successfully, an error otherwise
181    pub fn save_model(&self, model: &DeviceModel) -> Result<()> {
182        let file_path = self.model_dir.join(format!("{}.yaml", model.name));
183
184        // Serialize the model to YAML
185        let yaml = serde_yaml::to_string(model).map_err(|e| {
186            PoKeysError::ModelParseError(file_path.to_string_lossy().to_string(), e.to_string())
187        })?;
188
189        // Write the YAML to a file
190        fs::write(&file_path, yaml).map_err(|e| {
191            PoKeysError::ModelLoadError(file_path.to_string_lossy().to_string(), e.to_string())
192        })?;
193
194        info!("Saved model {} to {}", model.name, file_path.display());
195        Ok(())
196    }
197
198    /// Delete a model
199    ///
200    /// # Arguments
201    ///
202    /// * `name` - The name of the model
203    ///
204    /// # Returns
205    ///
206    /// * `Result<()>` - Ok if the model was deleted successfully, an error otherwise
207    pub fn delete_model(&mut self, name: &str) -> Result<()> {
208        // Remove the model from the map
209        self.models.remove(name);
210
211        // Delete the model file
212        let file_path = self.model_dir.join(format!("{}.yaml", name));
213
214        if file_path.exists() {
215            fs::remove_file(&file_path).map_err(|e| {
216                PoKeysError::ModelLoadError(file_path.to_string_lossy().to_string(), e.to_string())
217            })?;
218
219            info!("Deleted model {}", name);
220        }
221
222        Ok(())
223    }
224
225    /// Create a copy of a model with a new name
226    ///
227    /// # Arguments
228    ///
229    /// * `source_name` - The name of the source model
230    /// * `target_name` - The name of the target model
231    ///
232    /// # Returns
233    ///
234    /// * `Result<()>` - Ok if the model was copied successfully, an error otherwise
235    pub fn copy_model(&mut self, source_name: &str, target_name: &str) -> Result<()> {
236        // Get the source model
237        let source_model = self.get_model(source_name).ok_or_else(|| {
238            PoKeysError::ModelLoadError(source_name.to_string(), "Model not found".to_string())
239        })?;
240
241        // Create a copy of the model with the new name
242        let mut target_model = source_model.clone();
243        target_model.name = target_name.to_string();
244
245        // Save the target model
246        self.save_model(&target_model)?;
247
248        // Add the target model to the map
249        self.models.insert(target_name.to_string(), target_model);
250
251        Ok(())
252    }
253
254    /// Validate a model
255    ///
256    /// # Arguments
257    ///
258    /// * `name` - The name of the model
259    ///
260    /// # Returns
261    ///
262    /// * `Result<()>` - Ok if the model is valid, an error otherwise
263    pub fn validate_model(&self, name: &str) -> Result<()> {
264        // Get the model
265        let model = self.get_model(name).ok_or_else(|| {
266            PoKeysError::ModelLoadError(name.to_string(), "Model not found".to_string())
267        })?;
268
269        // Validate the model
270        model.validate()
271    }
272
273    /// Get the model directory
274    ///
275    /// # Returns
276    ///
277    /// * `&Path` - The model directory
278    pub fn get_model_dir(&self) -> &Path {
279        &self.model_dir
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use tempfile::tempdir;
287
288    #[test]
289    fn test_model_manager() {
290        // Create a temporary directory for the test
291        let dir = tempdir().unwrap();
292
293        // Create a model manager
294        let mut manager = ModelManager::new(Some(dir.path().to_path_buf())).unwrap();
295
296        // Create a model
297        let mut pins = HashMap::new();
298        pins.insert(
299            1,
300            PinModel {
301                capabilities: vec!["DigitalInput".to_string(), "DigitalOutput".to_string()],
302                active: true,
303            },
304        );
305
306        pins.insert(
307            2,
308            PinModel {
309                capabilities: vec!["DigitalInput".to_string(), "AnalogInput".to_string()],
310                active: true,
311            },
312        );
313
314        // Create the model
315        assert!(manager.create_model("TestModel", pins).is_ok());
316
317        // Get the model
318        let model = manager.get_model("TestModel");
319        assert!(model.is_some());
320
321        let model = model.unwrap();
322        assert_eq!(model.name, "TestModel");
323        assert_eq!(model.pins.len(), 2);
324
325        // Copy the model
326        assert!(manager.copy_model("TestModel", "TestModel2").is_ok());
327
328        // Get the copied model
329        let model = manager.get_model("TestModel2");
330        assert!(model.is_some());
331
332        let model = model.unwrap();
333        assert_eq!(model.name, "TestModel2");
334        assert_eq!(model.pins.len(), 2);
335
336        // Delete the model
337        assert!(manager.delete_model("TestModel").is_ok());
338
339        // Check that the model was deleted
340        assert!(manager.get_model("TestModel").is_none());
341
342        // Check that the copied model still exists
343        assert!(manager.get_model("TestModel2").is_some());
344
345        // Reload models
346        assert!(manager.reload_models().is_ok());
347
348        // Check that the copied model was reloaded
349        assert!(manager.get_model("TestModel2").is_some());
350    }
351}
pokeys_lib/model_manager.rs

pokeys_lib/
model_manager.rs
