graph_d 1.3.2

//! Memory-mapped file storage implementation.
//!
//! This module implements persistent storage using memory-mapped files.
//! Note: This is a simplified implementation focused on functionality.
//! Production systems would need more sophisticated record management.

use crate::error::{GraphError, Result};
use crate::graph::{Id, Node, Relationship};
use memmap2::{MmapMut, MmapOptions};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::fs::OpenOptions;
use std::path::Path;
use std::sync::Arc;

// Allow unsafe code only for this module for memory mapping
#[allow(unsafe_code)]
/// Simple file header for the memory-mapped file.
pub struct FileHeader {
    /// Magic bytes for file format identification ("GRPH")
    pub magic: [u8; 4],
    /// File format version number
    pub version: u32,
    /// Number of nodes stored in the file
    pub node_count: u64,
    /// Number of relationships stored in the file
    pub rel_count: u64,
}

impl Default for FileHeader {
    fn default() -> Self {
        Self::new()
    }
}

impl FileHeader {
    const MAGIC: &'static [u8; 4] = b"GRPH";
    const VERSION: u32 = 1;
    const SIZE: usize = 4 + 4 + 8 + 8; // 24 bytes

    /// Create a new file header with default values.
    pub fn new() -> Self {
        FileHeader {
            magic: *Self::MAGIC,
            version: Self::VERSION,
            node_count: 0,
            rel_count: 0,
        }
    }

    /// Check if this header has valid magic bytes and version.
    pub fn is_valid(&self) -> bool {
        self.magic == *Self::MAGIC && self.version == Self::VERSION
    }

    /// Serialize the header to bytes.
    pub fn serialize(&self) -> Vec<u8> {
        let mut data = Vec::with_capacity(Self::SIZE);
        data.extend_from_slice(&self.magic);
        data.extend_from_slice(&self.version.to_le_bytes());
        data.extend_from_slice(&self.node_count.to_le_bytes());
        data.extend_from_slice(&self.rel_count.to_le_bytes());
        data
    }

    /// Deserialize a header from bytes.
    pub fn deserialize(data: &[u8]) -> Result<Self> {
        if data.len() < Self::SIZE {
            return Err(GraphError::Storage("Invalid header size".to_string()));
        }

        let mut magic = [0u8; 4];
        magic.copy_from_slice(&data[0..4]);

        let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
        let node_count = u64::from_le_bytes([
            data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
        ]);
        let rel_count = u64::from_le_bytes([
            data[16], data[17], data[18], data[19], data[20], data[21], data[22], data[23],
        ]);

        Ok(FileHeader {
            magic,
            version,
            node_count,
            rel_count,
        })
    }
}

/// Memory-mapped storage backend.
pub struct MmapStorage {
    /// Memory-mapped file handle
    mmap: Arc<RwLock<MmapMut>>,
    /// Path to the storage file on disk
    file_path: std::path::PathBuf,
    /// In-memory cache for node lookups
    node_cache: Arc<RwLock<HashMap<Id, Node>>>,
    /// In-memory cache for relationship lookups
    relationship_cache: Arc<RwLock<HashMap<Id, Relationship>>>,
    /// Index mapping nodes to their relationship IDs for fast traversal
    node_relationships: Arc<RwLock<HashMap<Id, Vec<Id>>>>,
}

impl MmapStorage {
    /// Create a new memory-mapped storage at the given path.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file_path = path.as_ref().to_path_buf();

        // Create the file with initial size
        let initial_size = 1024 * 1024; // 1MB initial size
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .truncate(true)
            .open(&file_path)?;

        file.set_len(initial_size)?;

        // Create memory map
        #[allow(unsafe_code)]
        let mut mmap = unsafe { MmapOptions::new().map_mut(&file)? };

        // Initialize file header
        let header = FileHeader::new();
        let header_data = header.serialize();
        mmap[0..FileHeader::SIZE].copy_from_slice(&header_data);

        Ok(MmapStorage {
            mmap: Arc::new(RwLock::new(mmap)),
            file_path,
            node_cache: Arc::new(RwLock::new(HashMap::new())),
            relationship_cache: Arc::new(RwLock::new(HashMap::new())),
            node_relationships: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Open existing memory-mapped storage at the given path.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file_path = path.as_ref().to_path_buf();

        if !file_path.exists() {
            return Self::create(path);
        }

        let file = OpenOptions::new().read(true).write(true).open(&file_path)?;

        #[allow(unsafe_code)]
        let mmap = unsafe { MmapOptions::new().map_mut(&file)? };

        // Validate file header
        if mmap.len() < FileHeader::SIZE {
            return Err(GraphError::Storage(
                "File too small to contain header".to_string(),
            ));
        }

        let header = FileHeader::deserialize(&mmap[0..FileHeader::SIZE])?;
        if !header.is_valid() {
            return Err(GraphError::Storage("Invalid file header".to_string()));
        }

        let storage = MmapStorage {
            mmap: Arc::new(RwLock::new(mmap)),
            file_path,
            node_cache: Arc::new(RwLock::new(HashMap::new())),
            relationship_cache: Arc::new(RwLock::new(HashMap::new())),
            node_relationships: Arc::new(RwLock::new(HashMap::new())),
        };

        // Load existing data into caches
        storage.load_caches()?;

        Ok(storage)
    }

    /// Load existing data from the file into memory caches.
    fn load_caches(&self) -> Result<()> {
        let mmap = self.mmap.read();
        let header = FileHeader::deserialize(&mmap[0..FileHeader::SIZE])?;

        let mut offset = FileHeader::SIZE;

        // Load nodes
        for _ in 0..header.node_count {
            if offset + 8 > mmap.len() {
                break; // Corrupted file or incomplete write
            }

            // Read node data length
            let data_len = u64::from_le_bytes([
                mmap[offset],
                mmap[offset + 1],
                mmap[offset + 2],
                mmap[offset + 3],
                mmap[offset + 4],
                mmap[offset + 5],
                mmap[offset + 6],
                mmap[offset + 7],
            ]) as usize;
            offset += 8;

            if offset + data_len > mmap.len() {
                break; // Incomplete record
            }

            // Deserialize node using JSON
            let node_data = &mmap[offset..offset + data_len];
            if let Ok(json_str) = std::str::from_utf8(node_data) {
                if let Ok(node) = serde_json::from_str::<Node>(json_str) {
                    self.node_cache.write().insert(node.id, node.clone());
                    self.node_relationships.write().entry(node.id).or_default();
                }
            }
            offset += data_len;
        }

        // Load relationships
        for _ in 0..header.rel_count {
            if offset + 8 > mmap.len() {
                break;
            }

            let data_len = u64::from_le_bytes([
                mmap[offset],
                mmap[offset + 1],
                mmap[offset + 2],
                mmap[offset + 3],
                mmap[offset + 4],
                mmap[offset + 5],
                mmap[offset + 6],
                mmap[offset + 7],
            ]) as usize;
            offset += 8;

            if offset + data_len > mmap.len() {
                break;
            }

            let rel_data = &mmap[offset..offset + data_len];
            if let Ok(json_str) = std::str::from_utf8(rel_data) {
                if let Ok(rel) = serde_json::from_str::<Relationship>(json_str) {
                    let from_id = rel.from_id;
                    let to_id = rel.to_id;
                    let rel_id = rel.id;

                    self.relationship_cache.write().insert(rel_id, rel);

                    // Update relationship index
                    self.node_relationships
                        .write()
                        .entry(from_id)
                        .or_default()
                        .push(rel_id);

                    if from_id != to_id {
                        self.node_relationships
                            .write()
                            .entry(to_id)
                            .or_default()
                            .push(rel_id);
                    }
                }
            }
            offset += data_len;
        }

        Ok(())
    }

    /// Store a node in the memory-mapped storage.
    pub fn store_node(&self, node: Node) -> Result<()> {
        // Add to cache first
        self.node_cache.write().insert(node.id, node.clone());

        // Initialize empty relationship list
        self.node_relationships.write().entry(node.id).or_default();

        // TODO: Implement actual persistence to memory-mapped file
        // For now, we'll keep it in cache (in a production system,
        // this would write the node record and properties to the file)

        Ok(())
    }

    /// Retrieve a node by its ID.
    pub fn get_node(&self, id: Id) -> Result<Option<Node>> {
        Ok(self.node_cache.read().get(&id).cloned())
    }

    /// Store a relationship in the memory-mapped storage.
    pub fn store_relationship(&self, relationship: Relationship) -> Result<()> {
        let from_id = relationship.from_id;
        let to_id = relationship.to_id;
        let rel_id = relationship.id;

        // Verify nodes exist
        if !self.node_cache.read().contains_key(&from_id) {
            return Err(GraphError::NotFound(format!("Node {from_id} not found")));
        }
        if !self.node_cache.read().contains_key(&to_id) {
            return Err(GraphError::NotFound(format!("Node {to_id} not found")));
        }

        // Add to cache
        self.relationship_cache.write().insert(rel_id, relationship);

        // Update relationship index
        self.node_relationships
            .write()
            .entry(from_id)
            .or_default()
            .push(rel_id);

        if from_id != to_id {
            self.node_relationships
                .write()
                .entry(to_id)
                .or_default()
                .push(rel_id);
        }

        Ok(())
    }

    /// Retrieve a relationship by its ID.
    pub fn get_relationship(&self, id: Id) -> Result<Option<Relationship>> {
        Ok(self.relationship_cache.read().get(&id).cloned())
    }

    /// Get all relationships for a given node.
    pub fn get_relationships_for_node(&self, node_id: Id) -> Result<Vec<Relationship>> {
        let empty_vec = Vec::new();
        let relationship_ids = self
            .node_relationships
            .read()
            .get(&node_id)
            .unwrap_or(&empty_vec)
            .clone();

        let mut relationships = Vec::new();
        let rel_cache = self.relationship_cache.read();

        for rel_id in relationship_ids {
            if let Some(rel) = rel_cache.get(&rel_id) {
                relationships.push(rel.clone());
            }
        }

        Ok(relationships)
    }

    /// Get the number of nodes in storage.
    pub fn node_count(&self) -> usize {
        self.node_cache.read().len()
    }

    /// Get the number of relationships in storage.
    pub fn relationship_count(&self) -> usize {
        self.relationship_cache.read().len()
    }

    /// Delete a node from storage.
    /// Returns the deleted node if it existed.
    pub fn delete_node(&self, id: Id) -> Result<Option<Node>> {
        // Remove from node_relationships tracking
        self.node_relationships.write().remove(&id);
        // Remove the node
        Ok(self.node_cache.write().remove(&id))
    }

    /// Delete a relationship from storage.
    /// Returns the deleted relationship if it existed.
    pub fn delete_relationship(&self, id: Id) -> Result<Option<Relationship>> {
        if let Some(rel) = self.relationship_cache.write().remove(&id) {
            // Remove from node_relationships tracking
            let mut node_rels = self.node_relationships.write();
            if let Some(rels) = node_rels.get_mut(&rel.from_id) {
                rels.retain(|&r| r != id);
            }
            if rel.from_id != rel.to_id {
                if let Some(rels) = node_rels.get_mut(&rel.to_id) {
                    rels.retain(|&r| r != id);
                }
            }
            Ok(Some(rel))
        } else {
            Ok(None)
        }
    }

    /// Check if a node has any relationships.
    pub fn node_has_relationships(&self, node_id: Id) -> bool {
        self.node_relationships
            .read()
            .get(&node_id)
            .map(|rels| !rels.is_empty())
            .unwrap_or(false)
    }

    /// Flush all cached data to the memory-mapped file.
    pub fn flush(&self) -> Result<()> {
        let mut mmap = self.mmap.write();
        let mut offset = FileHeader::SIZE;

        // Serialize all nodes
        let nodes = self.node_cache.read();
        let relationships = self.relationship_cache.read();

        // Calculate total size needed using JSON serialization
        let mut total_size = FileHeader::SIZE;
        for node in nodes.values() {
            if let Ok(json_str) = serde_json::to_string(node) {
                total_size += 8 + json_str.len(); // 8 bytes for length + data
            }
        }
        for rel in relationships.values() {
            if let Ok(json_str) = serde_json::to_string(rel) {
                total_size += 8 + json_str.len(); // 8 bytes for length + data
            }
        }

        // Resize memory map if needed
        if total_size > mmap.len() {
            drop(mmap); // Release the lock before resizing

            // Resize the file
            let file = OpenOptions::new()
                .read(true)
                .write(true)
                .open(&self.file_path)?;
            file.set_len(total_size as u64 * 2)?; // Double for growth room

            // Recreate memory map
            #[allow(unsafe_code)]
            let new_mmap = unsafe { MmapOptions::new().map_mut(&file)? };
            *self.mmap.write() = new_mmap;
            mmap = self.mmap.write();
        }

        // Write nodes using JSON serialization for reliability with serde_json::Value
        for node in nodes.values() {
            if let Ok(json_str) = serde_json::to_string(node) {
                let serialized = json_str.as_bytes();
                if offset + 8 + serialized.len() > mmap.len() {
                    return Err(GraphError::Storage("Not enough space in mmap".to_string()));
                }

                // Write length
                let len_bytes = (serialized.len() as u64).to_le_bytes();
                mmap[offset..offset + 8].copy_from_slice(&len_bytes);
                offset += 8;

                // Write data
                mmap[offset..offset + serialized.len()].copy_from_slice(serialized);
                offset += serialized.len();
            }
        }

        // Write relationships using JSON serialization
        for rel in relationships.values() {
            if let Ok(json_str) = serde_json::to_string(rel) {
                let serialized = json_str.as_bytes();
                if offset + 8 + serialized.len() > mmap.len() {
                    return Err(GraphError::Storage("Not enough space in mmap".to_string()));
                }

                // Write length
                let len_bytes = (serialized.len() as u64).to_le_bytes();
                mmap[offset..offset + 8].copy_from_slice(&len_bytes);
                offset += 8;

                // Write data
                mmap[offset..offset + serialized.len()].copy_from_slice(serialized);
                offset += serialized.len();
            }
        }

        // Update and write header
        let mut header = FileHeader::new();
        header.node_count = nodes.len() as u64;
        header.rel_count = relationships.len() as u64;

        let header_data = header.serialize();
        mmap[0..FileHeader::SIZE].copy_from_slice(&header_data);

        // Force write to disk
        mmap.flush()?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use tempfile::NamedTempFile;

    #[test]
    fn test_mmap_storage_creation() {
        let temp_file = NamedTempFile::new().unwrap();
        let storage = MmapStorage::create(temp_file.path()).unwrap();

        assert_eq!(storage.node_count(), 0);
        assert_eq!(storage.relationship_count(), 0);
    }

    #[test]
    fn test_mmap_node_operations() {
        let temp_file = NamedTempFile::new().unwrap();
        let storage = MmapStorage::create(temp_file.path()).unwrap();

        let node = Node::new(1, HashMap::new());
        storage.store_node(node.clone()).unwrap();

        let retrieved = storage.get_node(1).unwrap();
        assert_eq!(retrieved, Some(node));
        assert_eq!(storage.node_count(), 1);
    }

    #[test]
    fn test_mmap_persistence() {
        let temp_file = NamedTempFile::new().unwrap();
        let path = temp_file.path().to_path_buf();

        // Create storage and add some data
        {
            let storage = MmapStorage::create(&path).unwrap();
            let mut properties = HashMap::new();
            properties.insert("name".to_string(), serde_json::json!("Alice"));
            properties.insert("age".to_string(), serde_json::json!(30));

            let node = Node::new(1, properties);
            storage.store_node(node).unwrap();

            // Verify data is in cache before flush
            assert_eq!(storage.node_count(), 1);

            storage.flush().unwrap();

            // Verify data is still in cache after flush
            let retrieved = storage.get_node(1).unwrap();
            assert!(retrieved.is_some());
            let retrieved_node = retrieved.unwrap();
            assert_eq!(retrieved_node.id, 1);
            assert_eq!(
                retrieved_node.get_property("name"),
                Some(&serde_json::json!("Alice"))
            );
        }

        // Reopen storage and verify persistence
        {
            let storage = MmapStorage::open(&path).unwrap();
            assert_eq!(storage.node_count(), 1, "Data should persist after reopen");

            let retrieved = storage.get_node(1).unwrap();
            assert!(retrieved.is_some());
            let retrieved_node = retrieved.unwrap();
            assert_eq!(retrieved_node.id, 1);
            assert_eq!(
                retrieved_node.get_property("name"),
                Some(&serde_json::json!("Alice"))
            );
            assert_eq!(
                retrieved_node.get_property("age"),
                Some(&serde_json::json!(30))
            );
        }
    }
}