use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use sled::{Db, Tree}; use std::collections::HashMap; use std::str::FromStr; use uuid::Uuid; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FileNode { pub node_id: String, pub label: String, pub aliases: Aliases, pub file_uuid: Option, pub sha256: Option, pub parent_id: Option, pub children: Vec, pub node_type: NodeType, pub icon: Option, pub color: Option, pub bg_color: Option, pub file_size: Option, pub registered_at: Option, pub created_at: String, pub updated_at: String, pub sort_order: i32, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Aliases { #[serde(flatten)] pub map: HashMap, } impl Aliases { pub fn empty() -> Self { Aliases { map: HashMap::new(), } } pub fn to_json(&self) -> String { serde_json::to_string(&self.map).unwrap_or_else(|_| "{}".to_string()) } pub fn from_json(s: &str) -> Self { let map: HashMap = serde_json::from_str(s).unwrap_or_default(); Aliases { map } } } #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] pub enum NodeType { Folder, File, DynamicLayer, } impl NodeType { pub fn as_str(&self) -> &'static str { match self { NodeType::Folder => "folder", NodeType::File => "file", NodeType::DynamicLayer => "dynamic_layer", } } } impl FromStr for NodeType { type Err = String; fn from_str(s: &str) -> Result { match s { "folder" => Ok(NodeType::Folder), "file" => Ok(NodeType::File), "dynamic_layer" => Ok(NodeType::DynamicLayer), _ => Ok(NodeType::Folder), } } } pub struct FileTreeSled { pub user_id: String, pub db: Db, nodes_tree: Tree, registry_tree: Tree, locations_tree: Tree, parent_index_tree: Tree, } impl FileTreeSled { pub fn user_db_path(user_id: &str) -> String { format!("data/users_sled/{}.sled", user_id) } pub fn init_user_db(user_id: &str) -> Result { let db_path = Self::user_db_path(user_id); let parent = std::path::Path::new(&db_path).parent().unwrap(); std::fs::create_dir_all(parent)?; let db = sled::open(&db_path)?; let nodes_tree = db.open_tree("file_nodes")?; let registry_tree = db.open_tree("file_registry")?; let locations_tree = db.open_tree("file_locations")?; let parent_index_tree = db.open_tree("parent_index")?; Ok(FileTreeSled { user_id: user_id.to_string(), db, nodes_tree, registry_tree, locations_tree, parent_index_tree, }) } pub fn open_user_db(user_id: &str) -> Result { let db_path = Self::user_db_path(user_id); let db = sled::open(&db_path).with_context(|| format!("Failed to open {}", db_path))?; let nodes_tree = db.open_tree("file_nodes")?; let registry_tree = db.open_tree("file_registry")?; let locations_tree = db.open_tree("file_locations")?; let parent_index_tree = db.open_tree("parent_index")?; Ok(FileTreeSled { user_id: user_id.to_string(), db, nodes_tree, registry_tree, locations_tree, parent_index_tree, }) } pub fn insert_node(&self, node: &FileNode) -> Result<()> { let node_data = serde_json::to_vec(node)?; self.nodes_tree .insert(node.node_id.as_bytes(), node_data.clone())?; if let Some(parent_id) = &node.parent_id { let mut children = self.get_children(parent_id)?; if !children.contains(&node.node_id) { children.push(node.node_id.clone()); self.parent_index_tree.insert( format!("children:{}", parent_id).as_bytes(), serde_json::to_vec(&children)?, )?; } } Ok(()) } pub fn insert_node_batch(&self, nodes: &[FileNode]) -> Result<()> { let mut batch = sled::Batch::default(); for node in nodes { let node_data = serde_json::to_vec(node)?; batch.insert(node.node_id.as_bytes(), node_data.clone()); } self.nodes_tree.apply_batch(batch)?; for node in nodes { if let Some(parent_id) = &node.parent_id { let mut children = self.get_children(parent_id)?; if !children.contains(&node.node_id) { children.push(node.node_id.clone()); self.parent_index_tree.insert( format!("children:{}", parent_id).as_bytes(), serde_json::to_vec(&children)?, )?; } } } Ok(()) } pub fn get_node(&self, node_id: &str) -> Result> { let value = self.nodes_tree.get(node_id.as_bytes())?; match value { Some(data) => { let node: FileNode = serde_json::from_slice(&data)?; Ok(Some(node)) } None => Ok(None), } } pub fn get_children(&self, parent_id: &str) -> Result> { let key = format!("children:{}", parent_id); let value = self.parent_index_tree.get(key.as_bytes())?; match value { Some(data) => { let children: Vec = serde_json::from_slice(&data)?; Ok(children) } None => Ok(Vec::new()), } } pub fn load_all(&self) -> Result> { let mut nodes = Vec::new(); for item in self.nodes_tree.iter() { let (_, value) = item?; let node: FileNode = serde_json::from_slice(&value)?; nodes.push(node); } nodes.sort_by(|a, b| { a.sort_order .cmp(&b.sort_order) .then_with(|| a.created_at.cmp(&b.created_at)) }); Ok(nodes) } pub fn update_node(&self, node_id: &str, updates: &FileNode) -> Result<()> { let node_data = serde_json::to_vec(updates)?; self.nodes_tree.insert(node_id.as_bytes(), node_data)?; Ok(()) } pub fn delete_node(&self, node_id: &str) -> Result<()> { let node = self.get_node(node_id)?; if let Some(n) = node { if let Some(parent_id) = &n.parent_id { let mut children = self.get_children(parent_id)?; children.retain(|id| id != node_id); self.parent_index_tree.insert( format!("children:{}", parent_id).as_bytes(), serde_json::to_vec(&children)?, )?; } } self.nodes_tree.remove(node_id.as_bytes())?; Ok(()) } pub fn count_nodes(&self) -> Result { Ok(self.nodes_tree.len()) } pub fn delete_all_nodes(&self) -> Result<()> { self.nodes_tree.clear()?; self.parent_index_tree.clear()?; Ok(()) } pub fn flush(&self) -> Result<()> { self.db.flush()?; Ok(()) } pub fn new_folder(label: &str, parent_id: Option<&str>) -> FileNode { FileNode { node_id: Uuid::new_v4().to_string().replace("-", ""), label: label.to_string(), aliases: Aliases::empty(), file_uuid: None, sha256: None, parent_id: parent_id.map(|s| s.to_string()), children: Vec::new(), node_type: NodeType::Folder, icon: None, color: None, bg_color: None, file_size: None, registered_at: None, created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), sort_order: 0, } } pub fn new_file_node( label: &str, file_uuid: &str, sha256: Option<&str>, original_name: &str, file_size: Option, mime_type: Option<&str>, parent_id: Option<&str>, ) -> FileNode { FileNode { node_id: Uuid::new_v4().to_string().replace("-", ""), label: label.to_string(), aliases: Aliases::empty(), file_uuid: Some(file_uuid.to_string()), sha256: sha256.map(|s| s.to_string()), parent_id: parent_id.map(|s| s.to_string()), children: Vec::new(), node_type: NodeType::File, icon: None, color: None, bg_color: None, file_size, registered_at: Some(chrono::Utc::now().to_rfc3339()), created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), sort_order: 0, } } } pub fn build_tree(nodes: &[FileNode]) -> Vec { let mut roots = Vec::new(); let node_map: HashMap = nodes.iter().map(|n| (n.node_id.clone(), n)).collect(); for node in nodes { if node.parent_id.is_none() { roots.push(node.clone()); } } roots }