1 files changed, 133 insertions, 0 deletions
diff --git a/src/state.rs b/src/state.rs
new file mode 100644
index 0000000..b9a6618
--- /dev/null
+++ b/src/state.rs
@@ -0,0 +1,133 @@
+use std::{collections::HashSet, fs};
+
+use derive_more::Display;
+use expanduser::expanduser;
+use serde::{Deserialize, Serialize};
+use walkdir::WalkDir;
+
+use crate::{
+    config::{Config, default_state_path},
+    learner::Learner,
+    trajectory::Trajectory,
+};
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct State {
+    pub config: Config,
+    trajectory: Trajectory,
+    learner: Learner,
+    albums: HashSet<String>,
+}
+
+impl State {
+    pub fn next(&mut self) -> Option<(String, f32)> {
+        let learning = self.trajectory.step(self.config.skip_window_secs);
+        if let Some(ref learning) = learning {
+            println!("Obtained learning: {learning:?}");
+            self.learner.learn(&self.config, learning);
+        }
+        if let Some((new_album, prob)) = self.learner.sample(
+            self.trajectory.slice(),
+            &learning.as_ref().into(),
+            &self.albums,
+            self.config.temperature,
+        ) {
+            self.trajectory.log(&new_album);
+            Some((new_album, prob))
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Display)]
+#[display("+{added}/-{removed}, total {new_total}")]
+pub struct UpdateResult {
+    added: usize,
+    removed: usize,
+    new_total: usize,
+}
+
+impl State {
+    pub fn load(path: &str) -> Result<State, String> {
+        expanduser(path)
+            .map_err(|e| e.to_string())
+            .and_then(|p| fs::read_to_string(p).map_err(|e| e.to_string()))
+            .and_then(|contents| ron::from_str(&contents).map_err(|e| e.to_string()))
+    }
+
+    pub fn try_save(&self) -> Result<(), String> {
+        let str_path = default_state_path();
+        let path = expanduser(&str_path).map_err(|e| e.to_string())?;
+        _ = path.parent().map(|parent| std::fs::create_dir_all(parent));
+        let serialised = ron::ser::to_string_pretty(self, ron::ser::PrettyConfig::default())
+            .map_err(|e| format!("Failed to serialize state: {}", e))?;
+        fs::write(&path, serialised)
+            .map_err(|e| format!("Failed to write state file at {}: {}", str_path, e))?;
+        Ok(())
+    }
+
+    pub fn update_albums(&mut self) -> UpdateResult {
+        let roots: Vec<_> = self
+            .config
+            .music_roots
+            .iter()
+            .filter_map(|x| expanduser(x).ok().filter(|p| p.is_dir()))
+            .collect();
+
+        let mut new_albums: HashSet<String> = HashSet::new();
+
+        for root_path in roots {
+            for entry in WalkDir::new(&root_path) {
+                let entry = match entry {
+                    Ok(e) if e.path().is_dir() => e,
+                    _ => continue,
+                };
+
+                if !self.contains_audio(&entry) {
+                    continue;
+                }
+
+                if let Some(rel_str) = entry
+                    .path()
+                    .strip_prefix(&root_path)
+                    .ok()
+                    .and_then(|p| p.to_str())
+                {
+                    new_albums.insert(rel_str.to_string());
+                }
+            }
+        }
+
+        let added = new_albums.difference(&self.albums).count();
+        let removed = self.albums.difference(&new_albums).count();
+        self.albums.clone_from(&new_albums);
+
+        self.trajectory.prune(&self.albums);
+        self.learner.prune(&self.albums);
+
+        UpdateResult {
+            added,
+            removed,
+            new_total: self.albums.len(),
+        }
+    }
+
+    fn contains_audio(&self, entry: &walkdir::DirEntry) -> bool {
+        fs::read_dir(entry.path()).ok().is_some_and(|mut entries| {
+            entries.any(|e| {
+                e.as_ref()
+                    .is_ok_and(|entry| self.is_audio_file(&entry.path()))
+            })
+        })
+    }
+
+    fn is_audio_file(&self, path: &std::path::Path) -> bool {
+        !path.is_dir()
+            && path
+                .extension()
+                .and_then(|ext| ext.to_str())
+                .map(|s| s.to_lowercase())
+                .is_some_and(|ext| self.config.audio_exts.contains(&format!(".{ext}")))
+    }
+}