markbase/filetree-rocksdb/src/lib.rs

use anyhow::{Context, Result};
use rocksdb::{ColumnFamilyDescriptor, Options, WriteBatch, DB};
use serde::{Deserialize, Serialize};
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 FileTreeRocksDB {
    pub user_id: String,
    pub db: DB,
}

impl FileTreeRocksDB {
    pub fn user_db_path(user_id: &str) -> String {
        format!("data/users_rocksdb/{}.rocksdb", 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 mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![
            ColumnFamilyDescriptor::new("file_nodes", Options::default()),
            ColumnFamilyDescriptor::new("file_registry", Options::default()),
            ColumnFamilyDescriptor::new("file_locations", Options::default()),
            ColumnFamilyDescriptor::new("parent_index", Options::default()),
        ];

        let db = DB::open_cf_descriptors(&opts, &db_path, cfs)
            .with_context(|| format!("Failed to open RocksDB at {}", db_path))?;

        Ok(FileTreeRocksDB {
            user_id: user_id.to_string(),
            db,
        })
    }

    pub fn open_user_db(user_id: &str) -> Result<Self> {
        let db_path = Self::user_db_path(user_id);

        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![
            ColumnFamilyDescriptor::new("file_nodes", Options::default()),
            ColumnFamilyDescriptor::new("file_registry", Options::default()),
            ColumnFamilyDescriptor::new("file_locations", Options::default()),
            ColumnFamilyDescriptor::new("parent_index", Options::default()),
        ];

        let db = DB::open_cf_descriptors(&opts, &db_path, cfs)
            .with_context(|| format!("Failed to open RocksDB at {}", db_path))?;

        Ok(FileTreeRocksDB {
            user_id: user_id.to_string(),
            db,
        })
    }

    pub fn insert_node(&self, node: &FileNode) -> Result<()> {
        let cf = self.db.cf_handle("file_nodes").unwrap();
        let node_data = serde_json::to_vec(node)?;
        self.db.put_cf(cf, node.node_id.as_bytes(), node_data)?;

        if let Some(parent_id) = &node.parent_id {
            let cf_parent = self.db.cf_handle("parent_index").unwrap();
            let key = format!("children:{}", parent_id);
            let existing = self.db.get_cf(cf_parent, key.as_bytes())?;

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

            if !children.contains(&node.node_id) {
                children.push(node.node_id.clone());
                self.db
                    .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?;
            }
        }

        Ok(())
    }

    pub fn insert_node_batch(&self, nodes: &[FileNode]) -> Result<()> {
        let mut batch = WriteBatch::default();
        let cf = self.db.cf_handle("file_nodes").unwrap();
        let cf_parent = self.db.cf_handle("parent_index").unwrap();

        for node in nodes {
            let node_data = serde_json::to_vec(node)?;
            batch.put_cf(cf, node.node_id.as_bytes(), node_data);
        }

        self.db.write(batch)?;

        for node in nodes {
            if let Some(parent_id) = &node.parent_id {
                let key = format!("children:{}", parent_id);
                let existing = self.db.get_cf(cf_parent, key.as_bytes())?;

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

                if !children.contains(&node.node_id) {
                    children.push(node.node_id.clone());
                    self.db
                        .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?;
                }
            }
        }

        Ok(())
    }

    pub fn get_node(&self, node_id: &str) -> Result<Option<FileNode>> {
        let cf = self.db.cf_handle("file_nodes").unwrap();
        let value = self.db.get_cf(cf, 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 cf = self.db.cf_handle("parent_index").unwrap();
        let key = format!("children:{}", parent_id);
        let value = self.db.get_cf(cf, 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 cf = self.db.cf_handle("file_nodes").unwrap();
        let mut nodes = Vec::new();

        let iter = self.db.iterator_cf(cf, rocksdb::IteratorMode::Start);
        for item in 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 cf = self.db.cf_handle("file_nodes").unwrap();
        let node_data = serde_json::to_vec(updates)?;
        self.db.put_cf(cf, 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 cf_parent = self.db.cf_handle("parent_index").unwrap();
                let key = format!("children:{}", parent_id);
                let existing = self.db.get_cf(cf_parent, key.as_bytes())?;

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

                children.retain(|id| id != node_id);
                self.db
                    .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?;
            }
        }

        let cf = self.db.cf_handle("file_nodes").unwrap();
        self.db.delete_cf(cf, node_id.as_bytes())?;
        Ok(())
    }

    pub fn count_nodes(&self) -> Result<usize> {
        let cf = self.db.cf_handle("file_nodes").unwrap();
        let mut count = 0;

        let iter = self.db.iterator_cf(cf, rocksdb::IteratorMode::Start);
        for item in iter {
            let _ = item?;
            count += 1;
        }

        Ok(count)
    }

    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,
        }
    }
}