//! HNSW Node and storage structures
//!
//! Contains HNSWNode and FineGrainedNodes for concurrent access.

use std::sync::Arc;

use parking_lot::{RwLock, RwLockReadGuard};
use serde::{Deserialize, Serialize};

/// Custom serde for Arc<Vec<f32>>
pub(crate) mod arc_vec_serde {
    use super::*;
    use serde::{Deserializer, Serializer};
    
    pub fn serialize<S>(data: &Arc<Vec<f32>>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        data.as_ref().serialize(serializer)
    }
    
    pub fn deserialize<'de, D>(deserializer: D) -> Result<Arc<Vec<f32>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let vec = Vec::<f32>::deserialize(deserializer)?;
        Ok(Arc::new(vec))
    }
}

/// Internal node structure for HNSW (optimized with Arc for vector sharing)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HNSWNode {
    /// The vector data (shared via Arc to avoid cloning)
    #[serde(with = "arc_vec_serde")]
    pub(crate) vector: Arc<Vec<f32>>,
    
    /// Neighbors at each level: level -> [neighbor_ids]
    pub(crate) neighbors: Vec<Vec<usize>>,
    
    /// Maximum level this node appears in
    pub(crate) max_level: usize,
}

impl HNSWNode {
    pub(crate) fn new(vector: Vec<f32>, max_level: usize, m_max0: usize, m: usize) -> Self {
        let mut neighbors = Vec::with_capacity(max_level + 1);
        for level in 0..=max_level {
            let capacity = if level == 0 { m_max0 } else { m };
            neighbors.push(Vec::with_capacity(capacity));
        }
        Self {
            vector: Arc::new(vector),
            neighbors,
            max_level,
        }
    }
    
    /// Create node with shared vector (avoids cloning)
    pub(crate) fn with_arc_vector(vector: Arc<Vec<f32>>, max_level: usize, m_max0: usize, m: usize) -> Self {
        let mut neighbors = Vec::with_capacity(max_level + 1);
        for level in 0..=max_level {
            let capacity = if level == 0 { m_max0 } else { m };
            neighbors.push(Vec::with_capacity(capacity));
        }
        Self {
            vector,
            neighbors,
            max_level,
        }
    }
    
    /// Get reference to vector (no cloning)
    #[inline]
    pub(crate) fn vector(&self) -> &[f32] {
        &self.vector
    }
    
    /// Get Arc reference to vector (for sharing)
    #[inline]
    pub(crate) fn vector_arc(&self) -> Arc<Vec<f32>> {
        Arc::clone(&self.vector)
    }
}

/// Fine-grained node storage with per-node locking
/// 
/// This structure allows concurrent access to different nodes,
/// significantly improving parallel insert/search performance.
#[allow(dead_code)]
pub struct FineGrainedNodes {
    /// Node data with per-node locks for neighbors
    nodes: RwLock<Vec<HNSWNode>>,
    /// Per-node locks for neighbor modifications (indexed by node_id)
    neighbor_locks: RwLock<Vec<RwLock<()>>>,
}

#[allow(dead_code)]
impl FineGrainedNodes {
    pub(crate) fn new(capacity: usize) -> Self {
        Self {
            nodes: RwLock::new(Vec::with_capacity(capacity)),
            neighbor_locks: RwLock::new(Vec::with_capacity(capacity)),
        }
    }
    
    /// Add a new node
    pub(crate) fn push(&self, node: HNSWNode) -> usize {
        let mut nodes = self.nodes.write();
        let id = nodes.len();
        nodes.push(node);
        
        let mut locks = self.neighbor_locks.write();
        locks.push(RwLock::new(()));
        
        id
    }
    
    /// Get node count
    pub(crate) fn len(&self) -> usize {
        self.nodes.read().len()
    }
    
    /// Check if empty
    pub(crate) fn is_empty(&self) -> bool {
        self.nodes.read().is_empty()
    }
    
    /// Read nodes (shared access)
    pub(crate) fn read(&self) -> RwLockReadGuard<'_, Vec<HNSWNode>> {
        self.nodes.read()
    }
    
    /// Get vector reference without full lock
    pub(crate) fn get_vector(&self, id: usize) -> Option<Arc<Vec<f32>>> {
        let nodes = self.nodes.read();
        nodes.get(id).map(|n| n.vector_arc())
    }
    
    /// Update neighbors for a specific node with fine-grained lock
    pub(crate) fn update_neighbors<F, R>(&self, node_id: usize, level: usize, f: F) -> Option<R>
    where
        F: FnOnce(&mut Vec<usize>) -> R,
    {
        // Acquire per-node lock first
        let locks = self.neighbor_locks.read();
        if node_id >= locks.len() {
            return None;
        }
        let _node_lock = locks[node_id].write();
        
        // Then modify neighbors
        let mut nodes = self.nodes.write();
        if node_id >= nodes.len() || level >= nodes[node_id].neighbors.len() {
            return None;
        }
        
        Some(f(&mut nodes[node_id].neighbors[level]))
    }
    
    /// Set neighbors for a node at a specific level
    pub(crate) fn set_neighbors(&self, node_id: usize, level: usize, neighbors: Vec<usize>) {
        let locks = self.neighbor_locks.read();
        if node_id >= locks.len() {
            return;
        }
        let _node_lock = locks[node_id].write();
        
        let mut nodes = self.nodes.write();
        if node_id < nodes.len() && level < nodes[node_id].neighbors.len() {
            nodes[node_id].neighbors[level] = neighbors;
        }
    }
    
    /// Get neighbors for a node (read-only, no per-node lock needed)
    pub(crate) fn get_neighbors(&self, node_id: usize, level: usize) -> Vec<usize> {
        let nodes = self.nodes.read();
        if node_id < nodes.len() && level < nodes[node_id].neighbors.len() {
            nodes[node_id].neighbors[level].clone()
        } else {
            Vec::new()
        }
    }
    
    /// Get max level for a node
    pub(crate) fn get_max_level(&self, node_id: usize) -> Option<usize> {
        let nodes = self.nodes.read();
        nodes.get(node_id).map(|n| n.max_level)
    }
}
