//! Graph Shard for parallel write processing
//!
//! Each shard contains a subset of entities and relations, with its own
//! write queue, WAL, and batch processing for high concurrency.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use parking_lot::{RwLock, Mutex};
use tokio::sync::mpsc;

use crate::types::graph::{Entity, Relation};
use crate::storage::concurrent::WriteMetrics;

use super::config::GraphShardConfig;
use super::types::{
    GraphWriteOp, GraphWriteRequest, GraphWriteResult, GraphWriteBatch, GraphStoreError,
};
use super::wal::{GraphWALWriter, GraphWALConfig, GraphWALEntry, GraphWALReader, GraphRecoveryResult, GraphCompactResult};

/// Indices maintained within a shard for efficient lookups
struct ShardIndices {
    /// entity_id -> outgoing relation_ids
    outgoing: RwLock<HashMap<String, Vec<String>>>,
    
    /// entity_id -> incoming relation_ids
    incoming: RwLock<HashMap<String, Vec<String>>>,
    
    /// entity_type -> entity_ids
    type_index: RwLock<HashMap<String, Vec<String>>>,
}

impl ShardIndices {
    fn new() -> Self {
        Self {
            outgoing: RwLock::new(HashMap::new()),
            incoming: RwLock::new(HashMap::new()),
            type_index: RwLock::new(HashMap::new()),
        }
    }
    
    /// Add entity to type index
    fn add_entity_type(&self, entity_id: &str, entity_type: &str) {
        self.type_index
            .write()
            .entry(entity_type.to_string())
            .or_default()
            .push(entity_id.to_string());
    }
    
    /// Remove entity from type index
    fn remove_entity_type(&self, entity_id: &str, entity_type: &str) {
        if let Some(ids) = self.type_index.write().get_mut(entity_type) {
            ids.retain(|id| id != entity_id);
        }
    }
    
    /// Add relation to indices
    fn add_relation(&self, relation_id: &str, source: &str, target: &str) {
        self.outgoing
            .write()
            .entry(source.to_string())
            .or_default()
            .push(relation_id.to_string());
        
        self.incoming
            .write()
            .entry(target.to_string())
            .or_default()
            .push(relation_id.to_string());
    }
    
    /// Remove relation from indices
    fn remove_relation(&self, relation_id: &str, source: &str, target: &str) {
        if let Some(ids) = self.outgoing.write().get_mut(source) {
            ids.retain(|id| id != relation_id);
        }
        if let Some(ids) = self.incoming.write().get_mut(target) {
            ids.retain(|id| id != relation_id);
        }
    }
    
    /// Get outgoing relation IDs for an entity
    fn get_outgoing(&self, entity_id: &str) -> Vec<String> {
        self.outgoing
            .read()
            .get(entity_id)
            .cloned()
            .unwrap_or_default()
    }
    
    /// Get incoming relation IDs for an entity
    fn get_incoming(&self, entity_id: &str) -> Vec<String> {
        self.incoming
            .read()
            .get(entity_id)
            .cloned()
            .unwrap_or_default()
    }
    
    /// Get entity IDs by type
    fn get_by_type(&self, entity_type: &str) -> Vec<String> {
        self.type_index
            .read()
            .get(entity_type)
            .cloned()
            .unwrap_or_default()
    }
}

/// A single graph shard that processes writes independently
pub struct GraphShard {
    /// Shard ID
    id: usize,
    
    /// Entity storage
    entities: RwLock<HashMap<String, Entity>>,
    
    /// Relation storage
    relations: RwLock<HashMap<String, Relation>>,
    
    /// Shard indices
    indices: ShardIndices,
    
    /// Write queue sender
    queue_sender: mpsc::Sender<GraphWriteRequest>,
    
    /// Write queue receiver (taken by worker)
    queue_receiver: Mutex<Option<mpsc::Receiver<GraphWriteRequest>>>,
    
    /// Queue capacity
    queue_capacity: usize,
    
    /// Sequence counter
    sequence: AtomicU64,
    
    /// Configuration
    config: GraphShardConfig,
    
    /// Metrics reference
    metrics: Arc<WriteMetrics>,
    
    /// Worker handle
    worker_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
    
    /// WAL writer (optional, for durability)
    wal_writer: Mutex<Option<GraphWALWriter>>,
    
    /// WAL configuration
    wal_config: Option<GraphWALConfig>,
}

impl GraphShard {
    /// Create a new graph shard
    pub fn new(config: GraphShardConfig, metrics: Arc<WriteMetrics>) -> Self {
        let capacity = config.queue_config.capacity;
        let (sender, receiver) = mpsc::channel(capacity);
        
        Self {
            id: config.id,
            entities: RwLock::new(HashMap::new()),
            relations: RwLock::new(HashMap::new()),
            indices: ShardIndices::new(),
            queue_sender: sender,
            queue_receiver: Mutex::new(Some(receiver)),
            queue_capacity: capacity,
            sequence: AtomicU64::new(0),
            config,
            metrics,
            worker_handle: Mutex::new(None),
            wal_writer: Mutex::new(None),
            wal_config: None,
        }
    }
    
    /// Create a new graph shard with WAL enabled
    pub fn with_wal(config: GraphShardConfig, metrics: Arc<WriteMetrics>, wal_config: GraphWALConfig) -> std::io::Result<Self> {
        let capacity = config.queue_config.capacity;
        let (sender, receiver) = mpsc::channel(capacity);
        let shard_id = config.id;
        
        // Create WAL writer
        let wal_writer = if wal_config.enabled {
            Some(GraphWALWriter::new(wal_config.clone(), shard_id, metrics.clone())?)
        } else {
            None
        };
        
        Ok(Self {
            id: shard_id,
            entities: RwLock::new(HashMap::new()),
            relations: RwLock::new(HashMap::new()),
            indices: ShardIndices::new(),
            queue_sender: sender,
            queue_receiver: Mutex::new(Some(receiver)),
            queue_capacity: capacity,
            sequence: AtomicU64::new(0),
            config,
            metrics,
            worker_handle: Mutex::new(None),
            wal_writer: Mutex::new(wal_writer),
            wal_config: Some(wal_config),
        })
    }
    
    /// Enable WAL on this shard
    pub fn enable_wal(&self, wal_config: GraphWALConfig) -> std::io::Result<()> {
        let mut wal_guard = self.wal_writer.lock();
        if wal_guard.is_none() && wal_config.enabled {
            *wal_guard = Some(GraphWALWriter::new(wal_config, self.id, self.metrics.clone())?);
        }
        Ok(())
    }
    
    /// Check if WAL is enabled
    pub fn is_wal_enabled(&self) -> bool {
        self.wal_writer.lock().as_ref().map(|w| w.is_enabled()).unwrap_or(false)
    }
    
    /// Recover from WAL
    pub fn recover(&self) -> std::io::Result<GraphRecoveryResult> {
        let wal_config = match &self.wal_config {
            Some(config) => config,
            None => {
                eprintln!("[WAL] Shard {}: No WAL config, skipping recovery", self.id);
                return Ok(GraphRecoveryResult::default());
            }
        };
        
        eprintln!("[WAL] Shard {}: Starting WAL recovery from {:?}", self.id, wal_config.path);
        
        let reader = GraphWALReader::new(&wal_config.path).for_shard(self.id);
        
        // Log the files found
        match reader.list_files() {
            Ok(files) => {
                eprintln!("[WAL] Shard {}: Found {} WAL files: {:?}", 
                    self.id, 
                    files.len(),
                    files.iter().map(|f| f.file_name().unwrap_or_default()).collect::<Vec<_>>()
                );
            }
            Err(e) => {
                eprintln!("[WAL] Shard {}: Failed to list WAL files: {}", self.id, e);
            }
        }
        
        let result = reader.replay(|entry| {
            if let Some(op) = entry.to_write_op() {
                self.apply_op(&op).map(|_| ()).map_err(|e| e.to_string())
            } else {
                Ok(())
            }
        });
        
        // Log recovery result and update shard sequence
        match &result {
            Ok(r) => {
                // CRITICAL: Update shard sequence to continue from where WAL left off
                // This prevents sequence number collision after restart
                if r.last_sequence > 0 {
                    let new_sequence = r.last_sequence + 1;
                    self.sequence.store(new_sequence, Ordering::SeqCst);
                    eprintln!(
                        "[WAL] Shard {}: Updated sequence to {} after recovery",
                        self.id, new_sequence
                    );
                }
                
                eprintln!(
                    "[WAL] Shard {}: Recovery complete - {} entries replayed, {} skipped. Now has {} entities, {} relations",
                    self.id, r.entries_replayed, r.entries_skipped,
                    self.entity_count(), self.relation_count()
                );
            }
            Err(e) => {
                eprintln!("[WAL] Shard {}: Recovery failed: {}", self.id, e);
            }
        }
        
        result
    }
    
    /// Get shard ID
    pub fn id(&self) -> usize {
        self.id
    }
    
    /// Get next sequence number
    fn next_sequence(&self) -> u64 {
        self.sequence.fetch_add(1, Ordering::SeqCst)
    }
    
    /// Get approximate queue length
    pub fn queue_len(&self) -> usize {
        self.queue_capacity - self.queue_sender.capacity()
    }
    
    /// Get queue capacity
    pub fn queue_capacity(&self) -> usize {
        self.queue_capacity
    }
    
    /// Get queue fill ratio
    pub fn queue_fill_ratio(&self) -> f64 {
        self.queue_len() as f64 / self.queue_capacity as f64
    }
    
    /// Check if queue is empty
    pub fn is_queue_empty(&self) -> bool {
        self.queue_len() == 0
    }
    
    /// Get entity count
    pub fn entity_count(&self) -> usize {
        self.entities.read().len()
    }
    
    /// Get relation count
    pub fn relation_count(&self) -> usize {
        self.relations.read().len()
    }
    
    // ========================================================================
    // Async Write API (via queue)
    // ========================================================================
    
    /// Write operation asynchronously (returns immediately after enqueueing)
    pub async fn write(&self, op: GraphWriteOp) -> Result<GraphWriteResult, GraphStoreError> {
        let sequence = self.next_sequence();
        let id = self.get_op_id(&op);
        let request = GraphWriteRequest::new(op, sequence);
        
        self.metrics.record_request();
        
        match tokio::time::timeout(
            Duration::from_millis(self.config.queue_config.send_timeout_ms),
            self.queue_sender.send(request)
        ).await {
            Ok(Ok(())) => {
                self.metrics.update_queue_size(self.queue_len());
                Ok(GraphWriteResult::success(id, sequence, self.id))
            }
            Ok(Err(_)) => {
                self.metrics.record_failure();
                Err(GraphStoreError::QueueClosed)
            }
            Err(_) => {
                self.metrics.record_rejection();
                Err(GraphStoreError::Timeout)
            }
        }
    }
    
    /// Write operation and wait for completion
    pub async fn write_sync(&self, op: GraphWriteOp) -> Result<GraphWriteResult, GraphStoreError> {
        let sequence = self.next_sequence();
        let (request, rx) = GraphWriteRequest::with_completion(op, sequence);
        
        self.metrics.record_request();
        
        // Send to queue
        match tokio::time::timeout(
            Duration::from_millis(self.config.queue_config.send_timeout_ms),
            self.queue_sender.send(request)
        ).await {
            Ok(Ok(())) => {
                self.metrics.update_queue_size(self.queue_len());
            }
            Ok(Err(_)) => {
                self.metrics.record_failure();
                return Err(GraphStoreError::QueueClosed);
            }
            Err(_) => {
                self.metrics.record_rejection();
                return Err(GraphStoreError::Timeout);
            }
        }
        
        // Wait for completion
        rx.await.map_err(|_| GraphStoreError::CompletionFailed)
    }
    
    /// Try to write without blocking
    pub fn try_write(&self, op: GraphWriteOp) -> Result<GraphWriteResult, GraphStoreError> {
        let sequence = self.next_sequence();
        let id = self.get_op_id(&op);
        let request = GraphWriteRequest::new(op, sequence);
        
        self.metrics.record_request();
        
        match self.queue_sender.try_send(request) {
            Ok(()) => {
                self.metrics.update_queue_size(self.queue_len());
                Ok(GraphWriteResult::success(id, sequence, self.id))
            }
            Err(mpsc::error::TrySendError::Full(_)) => {
                self.metrics.record_rejection();
                Err(GraphStoreError::QueueFull)
            }
            Err(mpsc::error::TrySendError::Closed(_)) => {
                self.metrics.record_failure();
                Err(GraphStoreError::QueueClosed)
            }
        }
    }
    
    // ========================================================================
    // Direct Read API (synchronous, no queue)
    // ========================================================================
    
    /// Get an entity by ID
    pub fn get_entity(&self, id: &str) -> Option<Entity> {
        self.entities.read().get(id).cloned()
    }
    
    /// Check if entity exists
    pub fn contains_entity(&self, id: &str) -> bool {
        self.entities.read().contains_key(id)
    }
    
    /// Get a relation by ID
    pub fn get_relation(&self, id: &str) -> Option<Relation> {
        self.relations.read().get(id).cloned()
    }
    
    /// Check if relation exists
    pub fn contains_relation(&self, id: &str) -> bool {
        self.relations.read().contains_key(id)
    }
    
    /// Get outgoing relations from an entity
    pub fn outgoing_from(&self, entity_id: &str) -> Vec<Relation> {
        let relation_ids = self.indices.get_outgoing(entity_id);
        let relations = self.relations.read();
        
        relation_ids
            .iter()
            .filter_map(|id| relations.get(id).cloned())
            .collect()
    }
    
    /// Get incoming relations to an entity
    pub fn incoming_to(&self, entity_id: &str) -> Vec<Relation> {
        let relation_ids = self.indices.get_incoming(entity_id);
        let relations = self.relations.read();
        
        relation_ids
            .iter()
            .filter_map(|id| relations.get(id).cloned())
            .collect()
    }
    
    /// Find entities by type
    pub fn find_entities_by_type(&self, entity_type: &str) -> Vec<Entity> {
        let entity_ids = self.indices.get_by_type(entity_type);
        let entities = self.entities.read();
        
        entity_ids
            .iter()
            .filter_map(|id| entities.get(id).cloned())
            .collect()
    }
    
    /// Get all entities in this shard
    pub fn all_entities(&self) -> Vec<Entity> {
        self.entities.read().values().cloned().collect()
    }
    
    /// Get all relations in this shard
    pub fn all_relations(&self) -> Vec<Relation> {
        self.relations.read().values().cloned().collect()
    }
    
    // ========================================================================
    // Direct Write API (synchronous, for batch processing)
    // ========================================================================
    
    /// Apply a write operation directly (called by batch worker)
    pub fn apply_op(&self, op: &GraphWriteOp) -> Result<String, GraphStoreError> {
        match op {
            GraphWriteOp::PutEntity(entity) => {
                self.put_entity_direct(entity.clone());
                Ok(entity.id.clone())
            }
            GraphWriteOp::DeleteEntity(id) => {
                self.delete_entity_direct(id);
                Ok(id.clone())
            }
            GraphWriteOp::PutRelation(relation) => {
                self.put_relation_direct(relation.clone());
                Ok(relation.id.clone())
            }
            GraphWriteOp::DeleteRelation(id) => {
                self.delete_relation_direct(id);
                Ok(id.clone())
            }
            GraphWriteOp::Batch(ops) => {
                for sub_op in ops {
                    self.apply_op(sub_op)?;
                }
                Ok("batch".to_string())
            }
        }
    }
    
    /// Put entity directly (no queue)
    fn put_entity_direct(&self, entity: Entity) {
        let mut entities = self.entities.write();
        
        // Handle type index updates
        if let Some(old) = entities.get(&entity.id) {
            // Updating existing entity
            if old.entity_type != entity.entity_type {
                // Type changed: remove from old type index, add to new
                self.indices.remove_entity_type(&entity.id, &old.entity_type);
                self.indices.add_entity_type(&entity.id, &entity.entity_type);
            }
            // If type unchanged, no index update needed
        } else {
            // New entity, add to type index
            self.indices.add_entity_type(&entity.id, &entity.entity_type);
        }
        
        entities.insert(entity.id.clone(), entity);
    }
    
    /// Delete entity directly (no queue)
    fn delete_entity_direct(&self, id: &str) {
        let mut entities = self.entities.write();
        
        if let Some(entity) = entities.remove(id) {
            self.indices.remove_entity_type(id, &entity.entity_type);
        }
    }
    
    /// Put relation directly (no queue)
    fn put_relation_direct(&self, relation: Relation) {
        let mut relations = self.relations.write();
        
        // Remove old indices if updating
        if let Some(old) = relations.get(&relation.id) {
            self.indices.remove_relation(&relation.id, &old.source, &old.target);
        }
        
        // Add new indices
        self.indices.add_relation(&relation.id, &relation.source, &relation.target);
        relations.insert(relation.id.clone(), relation);
    }
    
    /// Delete relation directly (no queue)
    fn delete_relation_direct(&self, id: &str) {
        let mut relations = self.relations.write();
        
        if let Some(relation) = relations.remove(id) {
            self.indices.remove_relation(id, &relation.source, &relation.target);
        }
    }
    
    // ========================================================================
    // Worker Management
    // ========================================================================
    
    /// Start the batch worker (must be called on Arc<GraphShard>)
    pub fn start(self: &Arc<Self>) {
        let mut handle_guard = self.worker_handle.lock();
        
        if handle_guard.is_some() {
            return; // Already started
        }
        
        if let Some(receiver) = self.queue_receiver.lock().take() {
            let shard = Arc::clone(self);
            let metrics = self.metrics.clone();
            let batch_config = self.config.batch_config.clone();
            let shard_id = self.id;
            
            let handle = tokio::spawn(async move {
                // Create a worker that processes batches
                let mut batch = GraphWriteBatch::with_capacity(batch_config.max_batch_size);
                let mut receiver = receiver;
                let flush_interval = Duration::from_millis(batch_config.max_wait_ms);
                let mut interval_timer = tokio::time::interval(flush_interval);
                
                loop {
                    tokio::select! {
                        request = receiver.recv() => {
                            match request {
                                Some(req) => {
                                    // Record queue time
                                    let queue_time = req.queue_time();
                                    metrics.record_queue_time(queue_time.as_micros() as u64);
                                    
                                    // Add to batch
                                    batch.push(req);
                                    
                                    // Check if we should flush
                                    if batch.len() >= batch_config.max_batch_size {
                                        Self::flush_batch(&shard, &mut batch, &metrics, shard_id);
                                    }
                                }
                                None => {
                                    // Channel closed, flush remaining and exit
                                    if !batch.is_empty() {
                                        Self::flush_batch(&shard, &mut batch, &metrics, shard_id);
                                    }
                                    break;
                                }
                            }
                        }
                        
                        _ = interval_timer.tick() => {
                            if !batch.is_empty() && batch.age() >= flush_interval {
                                Self::flush_batch(&shard, &mut batch, &metrics, shard_id);
                            }
                        }
                    }
                }
                
                tracing::info!("Graph shard {} worker shutting down", shard_id);
            });
            
            *handle_guard = Some(handle);
            tracing::info!("Started graph shard {} worker", self.id);
        }
    }
    
    /// Flush a batch of requests
    fn flush_batch(
        shard: &Arc<GraphShard>,
        batch: &mut GraphWriteBatch,
        metrics: &WriteMetrics,
        shard_id: usize,
    ) {
        let requests = batch.drain();
        let batch_size = requests.len();
        let timer = Instant::now();
        
        tracing::debug!(
            "Flushing graph batch: shard={}, size={}",
            shard_id,
            batch_size
        );
        
        // Collect WAL entries first
        let mut wal_entries = Vec::new();
        let has_wal = shard.is_wal_enabled();
        
        if has_wal {
            for request in &requests {
                let entries = GraphWALEntry::from_op(request.sequence, shard_id, &request.op);
                wal_entries.extend(entries);
            }
        }
        
        // Write to WAL before applying operations
        if has_wal && !wal_entries.is_empty() {
            let mut wal_guard = shard.wal_writer.lock();
            if let Some(ref mut wal) = *wal_guard {
                if let Err(e) = wal.append_batch(&wal_entries) {
                    tracing::error!("Failed to write to WAL: {}", e);
                    // Continue anyway - WAL failure shouldn't block operations
                    // In production, you might want different behavior
                }
            }
        }
        
        // Apply operations to in-memory storage
        for request in requests {
            let result = match shard.apply_op(&request.op) {
                Ok(id) => {
                    metrics.record_success();
                    GraphWriteResult::success(id, request.sequence, shard_id)
                }
                Err(e) => {
                    metrics.record_failure();
                    GraphWriteResult::failure(
                        shard.get_op_id(&request.op),
                        request.sequence,
                        e.to_string(),
                    )
                }
            };
            
            // Notify completion if requested
            request.notify_completion(result);
        }
        
        metrics.record_batch_flush(batch_size, 0);
        metrics.record_write_time(timer.elapsed().as_micros() as u64);
        
        tracing::debug!(
            "Graph batch flushed: shard={}, size={}, time_ms={}, wal={}",
            shard_id,
            batch_size,
            timer.elapsed().as_millis(),
            has_wal
        );
    }
    
    /// Get the ID from an operation (for result reporting)
    fn get_op_id(&self, op: &GraphWriteOp) -> String {
        match op {
            GraphWriteOp::PutEntity(e) => e.id.clone(),
            GraphWriteOp::DeleteEntity(id) => id.clone(),
            GraphWriteOp::PutRelation(r) => r.id.clone(),
            GraphWriteOp::DeleteRelation(id) => id.clone(),
            GraphWriteOp::Batch(_) => "batch".to_string(),
        }
    }
    
    /// Shutdown the shard worker
    pub async fn shutdown(&self) -> Result<(), GraphStoreError> {
        // Close WAL writer
        {
            let mut wal_guard = self.wal_writer.lock();
            if let Some(ref mut wal) = *wal_guard {
                if let Err(e) = wal.close() {
                    tracing::error!("Failed to close WAL: {}", e);
                }
            }
        }
        
        tracing::info!("Shutdown requested for graph shard {}", self.id);
        Ok(())
    }
    
    /// Sync WAL to disk
    pub fn sync_wal(&self) -> std::io::Result<()> {
        let mut wal_guard = self.wal_writer.lock();
        if let Some(ref mut wal) = *wal_guard {
            wal.sync()?;
        }
        Ok(())
    }
    
    /// Compact WAL files for this shard.
    ///
    /// Merges all .gwal files into a single file. The WAL writer is closed before
    /// compaction and reopened afterward. This is safe to call while the shard is
    /// not actively receiving writes (e.g., after flush/sync).
    pub fn compact_wal(&self) -> std::io::Result<GraphCompactResult> {
        let wal_config = match &self.wal_config {
            Some(config) if config.enabled => config.clone(),
            _ => {
                return Ok(GraphCompactResult::default());
            }
        };
        
        // Close the current WAL writer to release file handles
        {
            let mut wal_guard = self.wal_writer.lock();
            if let Some(ref mut wal) = *wal_guard {
                wal.close()?;
            }
            *wal_guard = None;
        }
        
        // Perform compaction using the reader
        let reader = GraphWALReader::new(&wal_config.path).for_shard(self.id);
        let result = reader.compact()?;
        
        // Reopen WAL writer
        {
            let mut wal_guard = self.wal_writer.lock();
            *wal_guard = Some(GraphWALWriter::new(wal_config, self.id, self.metrics.clone())?);
        }
        
        Ok(result)
    }
    
    /// Get configuration
    pub fn config(&self) -> &GraphShardConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::graph::{entity_types, relation_types};
    use tempfile::tempdir;
    
    fn create_test_shard() -> GraphShard {
        let dir = tempdir().unwrap();
        let config = GraphShardConfig::new(0, dir.path());
        let metrics = Arc::new(WriteMetrics::new());
        GraphShard::new(config, metrics)
    }
    
    #[test]
    fn test_shard_direct_entity_operations() {
        let shard = create_test_shard();
        
        let entity = Entity::new("Alice", entity_types::PERSON);
        let entity_id = entity.id.clone();
        
        // Put entity
        shard.put_entity_direct(entity);
        
        assert!(shard.contains_entity(&entity_id));
        assert_eq!(shard.entity_count(), 1);
        
        let retrieved = shard.get_entity(&entity_id).unwrap();
        assert_eq!(retrieved.name, "Alice");
        
        // Find by type
        let persons = shard.find_entities_by_type(entity_types::PERSON);
        assert_eq!(persons.len(), 1);
        
        // Delete entity
        shard.delete_entity_direct(&entity_id);
        assert!(!shard.contains_entity(&entity_id));
        assert_eq!(shard.entity_count(), 0);
    }
    
    #[test]
    fn test_shard_direct_relation_operations() {
        let shard = create_test_shard();
        
        let alice = Entity::new("Alice", entity_types::PERSON);
        let bob = Entity::new("Bob", entity_types::PERSON);
        let alice_id = alice.id.clone();
        let bob_id = bob.id.clone();
        
        shard.put_entity_direct(alice);
        shard.put_entity_direct(bob);
        
        let relation = Relation::new(&alice_id, &bob_id, relation_types::FRIEND);
        let relation_id = relation.id.clone();
        
        // Put relation
        shard.put_relation_direct(relation);
        
        assert!(shard.contains_relation(&relation_id));
        assert_eq!(shard.relation_count(), 1);
        
        // Check indices
        let outgoing = shard.outgoing_from(&alice_id);
        assert_eq!(outgoing.len(), 1);
        assert_eq!(outgoing[0].target, bob_id);
        
        let incoming = shard.incoming_to(&bob_id);
        assert_eq!(incoming.len(), 1);
        assert_eq!(incoming[0].source, alice_id);
        
        // Delete relation
        shard.delete_relation_direct(&relation_id);
        assert!(!shard.contains_relation(&relation_id));
        assert!(shard.outgoing_from(&alice_id).is_empty());
    }
    
    #[test]
    fn test_shard_apply_op() {
        let shard = create_test_shard();
        
        let entity = Entity::new("Test", entity_types::PERSON);
        let entity_id = entity.id.clone();
        
        // Apply put operation
        let op = GraphWriteOp::PutEntity(entity);
        let result = shard.apply_op(&op);
        assert!(result.is_ok());
        assert!(shard.contains_entity(&entity_id));
        
        // Apply delete operation
        let op = GraphWriteOp::DeleteEntity(entity_id.clone());
        let result = shard.apply_op(&op);
        assert!(result.is_ok());
        assert!(!shard.contains_entity(&entity_id));
    }
    
    fn create_test_shard_arc() -> Arc<GraphShard> {
        let dir = tempfile::tempdir().unwrap();
        let config = GraphShardConfig::new(0, dir.path());
        let metrics = Arc::new(WriteMetrics::new());
        Arc::new(GraphShard::new(config, metrics))
    }
    
    #[tokio::test]
    async fn test_shard_async_write() {
        let shard = create_test_shard_arc();
        shard.start();
        
        let entity = Entity::new("AsyncTest", entity_types::PERSON);
        let entity_id = entity.id.clone();
        
        // Write async
        let op = GraphWriteOp::PutEntity(entity);
        let result = shard.write(op).await.unwrap();
        
        assert!(result.success);
        assert_eq!(result.shard_id, 0);
        
        // Wait for batch to flush
        tokio::time::sleep(Duration::from_millis(100)).await;
        
        // Verify entity was written
        assert!(shard.contains_entity(&entity_id));
    }
    
    #[tokio::test]
    async fn test_shard_sync_write() {
        let shard = create_test_shard_arc();
        shard.start();
        
        let entity = Entity::new("SyncTest", entity_types::PERSON);
        let entity_id = entity.id.clone();
        
        // Write sync (waits for completion)
        let op = GraphWriteOp::PutEntity(entity);
        let result = shard.write_sync(op).await.unwrap();
        
        assert!(result.success);
        
        // Entity should be immediately available after sync write
        assert!(shard.contains_entity(&entity_id));
    }
    
    fn create_test_shard_with_wal() -> Arc<GraphShard> {
        let dir = tempfile::tempdir().unwrap();
        let config = GraphShardConfig::new(0, dir.path());
        let metrics = Arc::new(WriteMetrics::new());
        let wal_config = GraphWALConfig::new(dir.path().join("wal"));
        Arc::new(GraphShard::with_wal(config, metrics, wal_config).unwrap())
    }
    
    #[tokio::test]
    async fn test_shard_with_wal() {
        let shard = create_test_shard_with_wal();
        assert!(shard.is_wal_enabled());
        
        shard.start();
        
        // Write some entities
        for i in 0..5 {
            let entity = Entity::new(&format!("Person {}", i), entity_types::PERSON);
            let op = GraphWriteOp::PutEntity(entity);
            shard.write_sync(op).await.unwrap();
        }
        
        // Verify all written
        assert_eq!(shard.entity_count(), 5);
        
        // Sync WAL
        shard.sync_wal().unwrap();
    }
    
    #[test]
    fn test_shard_wal_recovery() {
        let dir = tempfile::tempdir().unwrap();
        let wal_path = dir.path().join("wal");
        let entity_ids: Vec<String>;
        
        // First, create shard and write some data
        {
            let config = GraphShardConfig::new(0, dir.path());
            let metrics = Arc::new(WriteMetrics::new());
            let wal_config = GraphWALConfig::new(&wal_path);
            let shard = GraphShard::with_wal(config, metrics, wal_config).unwrap();
            
            // Write directly (simulating batch flush)
            entity_ids = (0..3).map(|i| {
                let entity = Entity::new(&format!("Person {}", i), entity_types::PERSON);
                let entity_id = entity.id.clone();
                
                // Write to WAL
                {
                    let mut wal = shard.wal_writer.lock();
                    if let Some(ref mut w) = *wal {
                        let entry = GraphWALEntry::put_entity(i as u64, 0, entity.clone());
                        w.append(&entry).unwrap();
                    }
                }
                
                // Apply to storage
                shard.put_entity_direct(entity);
                entity_id
            }).collect();
            
            shard.sync_wal().unwrap();
        }
        
        // Now create a fresh shard and recover
        {
            let config = GraphShardConfig::new(0, dir.path());
            let metrics = Arc::new(WriteMetrics::new());
            let wal_config = GraphWALConfig::new(&wal_path);
            let shard = GraphShard::with_wal(config, metrics, wal_config).unwrap();
            
            // Recover from WAL
            let result = shard.recover().unwrap();
            
            assert_eq!(result.entries_replayed, 3);
            assert_eq!(result.entries_skipped, 0);
            
            // Verify entities recovered
            for entity_id in &entity_ids {
                assert!(shard.contains_entity(entity_id), "Entity {} not recovered", entity_id);
            }
        }
    }
}
