markbase/filetree-sled/src/lib.rs

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<String>,
    pub sha256: Option<String>,
    pub parent_id: Option<String>,
    pub children: Vec<String>,
    pub node_type: NodeType,
    pub icon: Option<String>,
    pub color: Option<String>,
    pub bg_color: Option<String>,
    pub file_size: Option<i64>,
    pub registered_at: Option<String>,
    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<String, String>,
}

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<String, String> = 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<Self, Self::Err> {
        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<Self> {
        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<Self> {
        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<Option<FileNode>> {
        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<Vec<String>> {
        let key = format!("children:{}", parent_id);
        let value = self.parent_index_tree.get(key.as_bytes())?;

        match value {
            Some(data) => {
                let children: Vec<String> = serde_json::from_slice(&data)?;
                Ok(children)
            }
            None => Ok(Vec::new()),
        }
    }

    pub fn load_all(&self) -> Result<Vec<FileNode>> {
        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<usize> {
        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<i64>,
        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<FileNode> {
    let mut roots = Vec::new();
    let node_map: HashMap<String, &FileNode> =
        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
}