//! Global Graph Index
//!
//! Provides cross-shard query support by maintaining global mappings:
//! - entity_id -> shard_id
//! - relation_id -> shard_id
//! - cross-shard relation indices (for incoming edge queries)
//!
//! The index is updated asynchronously during writes and can be persisted
//! for fast startup recovery.

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::{self, BufReader, BufWriter, Write};
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};

/// Cross-shard relation reference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CrossShardRef {
    /// Shard ID where the relation is stored
    pub shard_id: usize,
    /// Relation ID
    pub relation_id: String,
    /// Source entity ID
    pub source_id: String,
}

/// Global Graph Index for cross-shard queries
pub struct GlobalGraphIndex {
    /// entity_id -> shard_id mapping
    entity_shard_map: RwLock<HashMap<String, usize>>,
    
    /// relation_id -> shard_id mapping
    relation_shard_map: RwLock<HashMap<String, usize>>,
    
    /// target_entity_id -> list of (shard_id, relation_id, source_id)
    /// Used for incoming edge queries when relations are stored in different shards
    cross_shard_incoming: RwLock<HashMap<String, Vec<CrossShardRef>>>,
    
    /// entity_type -> set of entity_ids (for type-based queries)
    type_index: RwLock<HashMap<String, HashSet<String>>>,
    
    /// entity_id -> entity_type reverse mapping (Phase-96: O(1) type cleanup)
    ///
    /// Maintains the current type of each entity so that when the type changes,
    /// we can efficiently remove the entity from the old type set without iterating
    /// all type sets.
    entity_type_map: RwLock<HashMap<String, String>>,
    
    /// Index version for consistency tracking
    version: AtomicU64,
    
    /// Configuration
    config: GlobalIndexConfig,
    
    /// Persistence path
    path: Option<PathBuf>,
}

/// Configuration for the global index
#[derive(Debug, Clone)]
pub struct GlobalIndexConfig {
    /// Enable persistence
    pub persist_enabled: bool,
    
    /// Persistence interval (in number of updates)
    pub persist_interval: usize,
    
    /// Enable type index (for find_by_type queries)
    pub enable_type_index: bool,
    
    /// Maximum entries before compaction (0 = no limit)
    pub max_entries: usize,
}

impl Default for GlobalIndexConfig {
    fn default() -> Self {
        Self {
            persist_enabled: true,
            persist_interval: 1000,
            enable_type_index: true,
            max_entries: 0,
        }
    }
}

impl GlobalIndexConfig {
    /// Create config with persistence disabled
    pub fn no_persistence() -> Self {
        Self {
            persist_enabled: false,
            ..Default::default()
        }
    }
    
    /// Create config for high-frequency updates
    pub fn high_frequency() -> Self {
        Self {
            persist_interval: 100,
            ..Default::default()
        }
    }
}

/// Index integrity information for validation
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct IndexIntegrity {
    /// Expected entity count (from actual shard data when persisted)
    pub expected_entity_count: usize,
    /// Expected relation count (from actual shard data when persisted)
    pub expected_relation_count: usize,
    /// Timestamp when the index was last rebuilt
    pub last_rebuild_timestamp: u64,
    /// Schema version for forward compatibility
    pub schema_version: u32,
}

/// Current schema version for the index format
pub const INDEX_SCHEMA_VERSION: u32 = 2;

/// Snapshot of index data for persistence
#[derive(Debug, Clone, Serialize, Deserialize)]
struct IndexSnapshot {
    /// Version
    version: u64,
    /// Entity shard mappings
    entity_shards: Vec<(String, usize)>,
    /// Relation shard mappings
    relation_shards: Vec<(String, usize)>,
    /// Cross-shard incoming references
    cross_shard_incoming: Vec<(String, Vec<CrossShardRef>)>,
    /// Type index
    type_index: Vec<(String, Vec<String>)>,
    /// Entity → type reverse mapping (Phase-96, added in schema v3)
    #[serde(default)]
    entity_type_map: Vec<(String, String)>,
    /// Integrity information for validation (added in schema v2)
    #[serde(default)]
    integrity: IndexIntegrity,
}

impl GlobalGraphIndex {
    /// Create a new global index
    pub fn new(config: GlobalIndexConfig) -> Self {
        Self {
            entity_shard_map: RwLock::new(HashMap::new()),
            relation_shard_map: RwLock::new(HashMap::new()),
            cross_shard_incoming: RwLock::new(HashMap::new()),
            type_index: RwLock::new(HashMap::new()),
            entity_type_map: RwLock::new(HashMap::new()),
            version: AtomicU64::new(0),
            config,
            path: None,
        }
    }
    
    /// Create a new global index with persistence path
    pub fn with_path(config: GlobalIndexConfig, path: impl Into<PathBuf>) -> Self {
        Self {
            entity_shard_map: RwLock::new(HashMap::new()),
            relation_shard_map: RwLock::new(HashMap::new()),
            cross_shard_incoming: RwLock::new(HashMap::new()),
            type_index: RwLock::new(HashMap::new()),
            entity_type_map: RwLock::new(HashMap::new()),
            version: AtomicU64::new(0),
            config,
            path: Some(path.into()),
        }
    }
    
    /// Create with default config
    pub fn default_config() -> Self {
        Self::new(GlobalIndexConfig::default())
    }
    
    // ========================================================================
    // Entity Index Operations
    // ========================================================================
    
    /// Register an entity with its shard
    pub fn register_entity(&self, entity_id: &str, shard_id: usize) {
        self.entity_shard_map.write().insert(entity_id.to_string(), shard_id);
        self.increment_version();
    }
    
    /// Register an entity with its shard and type
    ///
    /// Phase-96 fix: When an entity's type changes, the old type mapping is
    /// automatically removed before the new one is added. This prevents stale
    /// type entries from accumulating (e.g., entity appearing in both
    /// "tree_section" and "tree_section_shadow" type sets).
    pub fn register_entity_with_type(&self, entity_id: &str, entity_type: &str, shard_id: usize) {
        self.entity_shard_map.write().insert(entity_id.to_string(), shard_id);
        
        if self.config.enable_type_index {
            let mut type_index = self.type_index.write();
            let mut entity_type_map = self.entity_type_map.write();
            
            // Phase-96: Remove from old type set if entity was previously registered
            // with a different type (O(1) via reverse mapping instead of O(types))
            if let Some(old_type) = entity_type_map.get(entity_id) {
                if old_type != entity_type {
                    if let Some(old_set) = type_index.get_mut(old_type.as_str()) {
                        old_set.remove(entity_id);
                        // Clean up empty type sets to avoid memory leaks
                        if old_set.is_empty() {
                            type_index.remove(old_type.as_str());
                        }
                    }
                }
            }
            
            // Add to new type set
            type_index
                .entry(entity_type.to_string())
                .or_insert_with(HashSet::new)
                .insert(entity_id.to_string());
            
            // Update reverse mapping
            entity_type_map.insert(entity_id.to_string(), entity_type.to_string());
        }
        
        self.increment_version();
    }
    
    /// Remove an entity from the index
    ///
    /// Phase-96 fix: Uses the entity_type_map reverse index for O(1) type
    /// cleanup instead of iterating all type sets. Falls back to full scan
    /// if the reverse mapping is missing (defensive).
    pub fn remove_entity(&self, entity_id: &str) {
        self.entity_shard_map.write().remove(entity_id);
        
        // Remove from type index using reverse mapping (O(1))
        if self.config.enable_type_index {
            let mut entity_type_map = self.entity_type_map.write();
            let mut type_index = self.type_index.write();
            
            if let Some(old_type) = entity_type_map.remove(entity_id) {
                // O(1) removal via reverse mapping
                if let Some(entities) = type_index.get_mut(&old_type) {
                    entities.remove(entity_id);
                    if entities.is_empty() {
                        type_index.remove(&old_type);
                    }
                }
            } else {
                // Defensive fallback: iterate all types (shouldn't happen normally)
                let mut empty_types = Vec::new();
                for (type_name, entities) in type_index.iter_mut() {
                    entities.remove(entity_id);
                    if entities.is_empty() {
                        empty_types.push(type_name.clone());
                    }
                }
                for t in empty_types {
                    type_index.remove(&t);
                }
            }
        }
        
        // Remove incoming references pointing to this entity
        self.cross_shard_incoming.write().remove(entity_id);
        
        self.increment_version();
    }
    
    /// Get the shard ID for an entity
    pub fn get_entity_shard(&self, entity_id: &str) -> Option<usize> {
        self.entity_shard_map.read().get(entity_id).copied()
    }
    
    /// Check if an entity is registered
    pub fn contains_entity(&self, entity_id: &str) -> bool {
        self.entity_shard_map.read().contains_key(entity_id)
    }
    
    /// Get all entity IDs of a specific type
    pub fn get_entities_by_type(&self, entity_type: &str) -> Vec<String> {
        if !self.config.enable_type_index {
            return Vec::new();
        }
        
        self.type_index
            .read()
            .get(entity_type)
            .map(|set| set.iter().cloned().collect())
            .unwrap_or_default()
    }
    
    /// Get a paginated slice of entity IDs of a specific type (Phase-97)
    ///
    /// Returns `(entity_ids, total_count)` where `entity_ids` is the requested
    /// page of IDs and `total_count` is the total number of entities of this type.
    ///
    /// # Arguments
    /// * `entity_type` - The entity type to query
    /// * `offset` - Number of items to skip
    /// * `limit` - Maximum number of items to return
    ///
    /// # Example
    /// ```
    /// let (ids, total) = index.get_entities_by_type_paginated("person", 0, 100);
    /// assert!(ids.len() <= 100);
    /// ```
    pub fn get_entities_by_type_paginated(
        &self,
        entity_type: &str,
        offset: usize,
        limit: usize,
    ) -> (Vec<String>, usize) {
        if !self.config.enable_type_index {
            return (Vec::new(), 0);
        }
        
        let type_index = self.type_index.read();
        match type_index.get(entity_type) {
            Some(set) => {
                let total = set.len();
                let ids: Vec<String> = set
                    .iter()
                    .skip(offset)
                    .take(limit)
                    .cloned()
                    .collect();
                (ids, total)
            }
            None => (Vec::new(), 0),
        }
    }
    
    /// Count entities of a specific type without loading all IDs (Phase-97)
    ///
    /// More efficient than `get_entities_by_type().len()` as it avoids
    /// cloning all entity ID strings.
    pub fn count_entities_by_type(&self, entity_type: &str) -> usize {
        if !self.config.enable_type_index {
            return 0;
        }
        
        self.type_index
            .read()
            .get(entity_type)
            .map(|set| set.len())
            .unwrap_or(0)
    }
    
    /// Get all registered entity types
    ///
    /// Returns a list of all entity_type strings that have at least one entity
    /// in the type index. Includes both built-in types (person, tag, etc.)
    /// and custom types (article, category, etc.).
    pub fn get_all_entity_types(&self) -> Vec<String> {
        if !self.config.enable_type_index {
            return Vec::new();
        }
        
        self.type_index
            .read()
            .keys()
            .cloned()
            .collect()
    }
    
    /// Get the current type of an entity (Phase-96)
    ///
    /// Returns the entity_type string if the entity is registered, None otherwise.
    pub fn get_entity_type(&self, entity_id: &str) -> Option<String> {
        self.entity_type_map.read().get(entity_id).cloned()
    }
    
    /// Get entity count
    pub fn entity_count(&self) -> usize {
        self.entity_shard_map.read().len()
    }
    
    // ========================================================================
    // Relation Index Operations
    // ========================================================================
    
    /// Register a relation with its shard
    pub fn register_relation(&self, relation_id: &str, shard_id: usize) {
        self.relation_shard_map.write().insert(relation_id.to_string(), shard_id);
        self.increment_version();
    }
    
    /// Register a relation with cross-shard reference
    /// Call this when the relation's target entity is in a different shard
    pub fn register_relation_with_target(
        &self,
        relation_id: &str,
        source_id: &str,
        target_id: &str,
        relation_shard: usize,
        target_shard: usize,
    ) {
        // Register relation -> shard mapping
        self.relation_shard_map.write().insert(relation_id.to_string(), relation_shard);
        
        // If cross-shard, add to incoming index for the target
        if relation_shard != target_shard {
            let cross_ref = CrossShardRef {
                shard_id: relation_shard,
                relation_id: relation_id.to_string(),
                source_id: source_id.to_string(),
            };
            
            self.cross_shard_incoming
                .write()
                .entry(target_id.to_string())
                .or_insert_with(Vec::new)
                .push(cross_ref);
        }
        
        self.increment_version();
    }
    
    /// Remove a relation from the index
    pub fn remove_relation(&self, relation_id: &str, target_id: Option<&str>) {
        self.relation_shard_map.write().remove(relation_id);
        
        // Remove from cross-shard incoming index if target is provided
        if let Some(target) = target_id {
            let mut incoming = self.cross_shard_incoming.write();
            if let Some(refs) = incoming.get_mut(target) {
                refs.retain(|r| r.relation_id != relation_id);
                if refs.is_empty() {
                    incoming.remove(target);
                }
            }
        }
        
        self.increment_version();
    }
    
    /// Get the shard ID for a relation
    pub fn get_relation_shard(&self, relation_id: &str) -> Option<usize> {
        self.relation_shard_map.read().get(relation_id).copied()
    }
    
    /// Check if a relation is registered
    pub fn contains_relation(&self, relation_id: &str) -> bool {
        self.relation_shard_map.read().contains_key(relation_id)
    }
    
    /// Get relation count
    pub fn relation_count(&self) -> usize {
        self.relation_shard_map.read().len()
    }
    
    // ========================================================================
    // Cross-Shard Query Support
    // ========================================================================
    
    /// Get incoming relation references from other shards
    /// Returns list of (shard_id, relation_id, source_id) tuples
    pub fn get_incoming_refs(&self, target_id: &str) -> Vec<CrossShardRef> {
        self.cross_shard_incoming
            .read()
            .get(target_id)
            .cloned()
            .unwrap_or_default()
    }
    
    /// Get all shards that have relations pointing to a target
    pub fn get_incoming_shards(&self, target_id: &str) -> Vec<usize> {
        self.cross_shard_incoming
            .read()
            .get(target_id)
            .map(|refs| {
                refs.iter()
                    .map(|r| r.shard_id)
                    .collect::<HashSet<_>>()
                    .into_iter()
                    .collect()
            })
            .unwrap_or_default()
    }
    
    /// Get cross-shard relation count for a target
    pub fn cross_shard_count(&self, target_id: &str) -> usize {
        self.cross_shard_incoming
            .read()
            .get(target_id)
            .map(|refs| refs.len())
            .unwrap_or(0)
    }
    
    /// Get total cross-shard relation count
    pub fn total_cross_shard_count(&self) -> usize {
        self.cross_shard_incoming
            .read()
            .values()
            .map(|refs| refs.len())
            .sum()
    }
    
    // ========================================================================
    // Batch Operations (for efficiency)
    // ========================================================================
    
    /// Register multiple entities at once
    pub fn register_entities_batch(&self, entities: &[(String, usize)]) {
        let mut map = self.entity_shard_map.write();
        for (entity_id, shard_id) in entities {
            map.insert(entity_id.clone(), *shard_id);
        }
        drop(map);
        self.increment_version();
    }
    
    /// Register multiple entities with types at once
    ///
    /// Phase-96: Also maintains entity_type_map and cleans up old type mappings.
    pub fn register_entities_with_types_batch(&self, entities: &[(String, String, usize)]) {
        let mut entity_map = self.entity_shard_map.write();
        let mut type_index = self.type_index.write();
        let mut entity_type_map = self.entity_type_map.write();
        
        for (entity_id, entity_type, shard_id) in entities {
            entity_map.insert(entity_id.clone(), *shard_id);
            
            if self.config.enable_type_index {
                // Phase-96: Remove from old type set if type changed
                if let Some(old_type) = entity_type_map.get(entity_id.as_str()) {
                    if old_type != entity_type {
                        if let Some(old_set) = type_index.get_mut(old_type.as_str()) {
                            old_set.remove(entity_id.as_str());
                            if old_set.is_empty() {
                                type_index.remove(old_type.as_str());
                            }
                        }
                    }
                }
                
                type_index
                    .entry(entity_type.clone())
                    .or_insert_with(HashSet::new)
                    .insert(entity_id.clone());
                
                entity_type_map.insert(entity_id.clone(), entity_type.clone());
            }
        }
        
        drop(entity_map);
        drop(type_index);
        drop(entity_type_map);
        self.increment_version();
    }
    
    /// Register multiple relations at once
    pub fn register_relations_batch(&self, relations: &[(String, usize)]) {
        let mut map = self.relation_shard_map.write();
        for (relation_id, shard_id) in relations {
            map.insert(relation_id.clone(), *shard_id);
        }
        drop(map);
        self.increment_version();
    }
    
    // ========================================================================
    // Persistence
    // ========================================================================
    
    /// Save index to disk
    pub fn persist(&self) -> io::Result<()> {
        let path = match &self.path {
            Some(p) => p.clone(),
            None => return Ok(()), // No path configured
        };
        
        if !self.config.persist_enabled {
            return Ok(());
        }
        
        // Create snapshot
        let snapshot = self.create_snapshot();
        
        // Ensure directory exists
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        
        // Write to temp file first
        let temp_path = path.with_extension("tmp");
        let file = File::create(&temp_path)?;
        let mut writer = BufWriter::new(file);
        
        serde_json::to_writer(&mut writer, &snapshot)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        writer.flush()?;
        drop(writer);
        
        // Atomic rename
        std::fs::rename(&temp_path, &path)?;
        
        tracing::debug!("Global index persisted: version={}, entities={}, relations={}", 
            snapshot.version, snapshot.entity_shards.len(), snapshot.relation_shards.len());
        
        Ok(())
    }
    
    /// Load index from disk
    pub fn load(&self) -> io::Result<bool> {
        let path = match &self.path {
            Some(p) => p.clone(),
            None => return Ok(false),
        };
        
        if !path.exists() {
            return Ok(false);
        }
        
        let file = File::open(&path)?;
        let reader = BufReader::new(file);
        
        let snapshot: IndexSnapshot = serde_json::from_reader(reader)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        self.apply_snapshot(snapshot);
        
        tracing::info!("Global index loaded: version={}, entities={}, relations={}",
            self.version.load(Ordering::SeqCst),
            self.entity_count(),
            self.relation_count());
        
        Ok(true)
    }
    
    /// Load index from disk with integrity validation
    /// Returns (loaded, integrity_valid)
    pub fn load_with_validation(&self) -> io::Result<(bool, bool)> {
        let path = match &self.path {
            Some(p) => p.clone(),
            None => return Ok((false, false)),
        };
        
        if !path.exists() {
            return Ok((false, false));
        }
        
        let file = File::open(&path)?;
        let reader = BufReader::new(file);
        
        let snapshot: IndexSnapshot = serde_json::from_reader(reader)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        // Check schema version
        if snapshot.integrity.schema_version != INDEX_SCHEMA_VERSION && snapshot.integrity.schema_version != 0 {
            tracing::warn!(
                "Index schema version mismatch: expected {}, got {}",
                INDEX_SCHEMA_VERSION,
                snapshot.integrity.schema_version
            );
        }
        
        // Check integrity
        let entity_count = snapshot.entity_shards.len();
        let relation_count = snapshot.relation_shards.len();
        
        let integrity_valid = snapshot.integrity.expected_entity_count == entity_count
            && snapshot.integrity.expected_relation_count == relation_count;
        
        if !integrity_valid {
            tracing::warn!(
                "Index integrity check failed: expected ({} entities, {} relations), actual ({} entities, {} relations)",
                snapshot.integrity.expected_entity_count,
                snapshot.integrity.expected_relation_count,
                entity_count,
                relation_count
            );
        }
        
        self.apply_snapshot(snapshot);
        
        tracing::info!(
            "Global index loaded: version={}, entities={}, relations={}, integrity_valid={}",
            self.version.load(Ordering::SeqCst),
            self.entity_count(),
            self.relation_count(),
            integrity_valid
        );
        
        Ok((true, integrity_valid))
    }
    
    /// Verify index integrity against actual shard data
    /// Returns (is_valid, indexed_count, actual_count)
    pub fn verify_integrity(&self, actual_entity_count: usize, actual_relation_count: usize) -> IndexIntegrityResult {
        let indexed_entities = self.entity_count();
        let indexed_relations = self.relation_count();
        
        let entity_match = indexed_entities == actual_entity_count;
        let relation_match = indexed_relations == actual_relation_count;
        let is_valid = entity_match && relation_match;
        
        if !is_valid {
            tracing::warn!(
                "Index integrity mismatch: indexed ({} entities, {} relations) vs actual ({} entities, {} relations)",
                indexed_entities,
                indexed_relations,
                actual_entity_count,
                actual_relation_count
            );
        }
        
        IndexIntegrityResult {
            is_valid,
            indexed_entity_count: indexed_entities,
            indexed_relation_count: indexed_relations,
            actual_entity_count,
            actual_relation_count,
        }
    }
    
    /// Create a snapshot of the current state
    fn create_snapshot(&self) -> IndexSnapshot {
        let entity_shards: Vec<_> = self.entity_shard_map
            .read()
            .iter()
            .map(|(k, v)| (k.clone(), *v))
            .collect();
        
        let relation_shards: Vec<_> = self.relation_shard_map
            .read()
            .iter()
            .map(|(k, v)| (k.clone(), *v))
            .collect();
        
        let cross_shard_incoming: Vec<_> = self.cross_shard_incoming
            .read()
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        
        let type_index: Vec<_> = self.type_index
            .read()
            .iter()
            .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
            .collect();
        
        // Phase-96: Persist entity_type_map
        let entity_type_map: Vec<_> = self.entity_type_map
            .read()
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        
        // Build integrity info
        let integrity = IndexIntegrity {
            expected_entity_count: entity_shards.len(),
            expected_relation_count: relation_shards.len(),
            last_rebuild_timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_millis() as u64,
            schema_version: INDEX_SCHEMA_VERSION,
        };
        
        IndexSnapshot {
            version: self.version.load(Ordering::SeqCst),
            entity_shards,
            relation_shards,
            cross_shard_incoming,
            type_index,
            entity_type_map,
            integrity,
        }
    }
    
    /// Apply a snapshot to restore state
    fn apply_snapshot(&self, snapshot: IndexSnapshot) {
        // Clear existing data
        self.entity_shard_map.write().clear();
        self.relation_shard_map.write().clear();
        self.cross_shard_incoming.write().clear();
        self.type_index.write().clear();
        self.entity_type_map.write().clear();
        
        // Apply entity mappings
        {
            let mut map = self.entity_shard_map.write();
            for (entity_id, shard_id) in snapshot.entity_shards {
                map.insert(entity_id, shard_id);
            }
        }
        
        // Apply relation mappings
        {
            let mut map = self.relation_shard_map.write();
            for (relation_id, shard_id) in snapshot.relation_shards {
                map.insert(relation_id, shard_id);
            }
        }
        
        // Apply cross-shard incoming
        {
            let mut incoming = self.cross_shard_incoming.write();
            for (target_id, refs) in snapshot.cross_shard_incoming {
                incoming.insert(target_id, refs);
            }
        }
        
        // Apply type index
        {
            let mut type_idx = self.type_index.write();
            for (entity_type, entity_ids) in snapshot.type_index {
                type_idx.insert(entity_type, entity_ids.into_iter().collect());
            }
        }
        
        // Phase-96: Apply entity_type_map (backward compatible — empty for old snapshots)
        {
            let mut etm = self.entity_type_map.write();
            for (entity_id, entity_type) in snapshot.entity_type_map {
                etm.insert(entity_id, entity_type);
            }
            
            // If entity_type_map was empty (loading from pre-Phase-96 snapshot),
            // rebuild it from the type_index
            if etm.is_empty() {
                let type_idx = self.type_index.read();
                for (entity_type, entity_ids) in type_idx.iter() {
                    for entity_id in entity_ids {
                        etm.insert(entity_id.clone(), entity_type.clone());
                    }
                }
            }
        }
        
        self.version.store(snapshot.version, Ordering::SeqCst);
    }
    
    // ========================================================================
    // Utility Methods
    // ========================================================================
    
    /// Get current version
    pub fn version(&self) -> u64 {
        self.version.load(Ordering::SeqCst)
    }
    
    /// Increment version
    fn increment_version(&self) -> u64 {
        self.version.fetch_add(1, Ordering::SeqCst)
    }
    
    /// Clear all index data
    pub fn clear(&self) {
        self.entity_shard_map.write().clear();
        self.relation_shard_map.write().clear();
        self.cross_shard_incoming.write().clear();
        self.type_index.write().clear();
        self.entity_type_map.write().clear();
        self.version.store(0, Ordering::SeqCst);
    }
    
    /// Get index statistics
    pub fn stats(&self) -> GlobalIndexStats {
        GlobalIndexStats {
            version: self.version(),
            entity_count: self.entity_count(),
            relation_count: self.relation_count(),
            cross_shard_count: self.total_cross_shard_count(),
            type_count: self.type_index.read().len(),
        }
    }
}

/// Index statistics
#[derive(Debug, Clone)]
pub struct GlobalIndexStats {
    /// Current version
    pub version: u64,
    /// Number of entities indexed
    pub entity_count: usize,
    /// Number of relations indexed
    pub relation_count: usize,
    /// Number of cross-shard relations
    pub cross_shard_count: usize,
    /// Number of unique entity types
    pub type_count: usize,
}

/// Result of index integrity verification
#[derive(Debug, Clone)]
pub struct IndexIntegrityResult {
    /// Whether the index is valid (matches actual data)
    pub is_valid: bool,
    /// Number of entities in the index
    pub indexed_entity_count: usize,
    /// Number of relations in the index
    pub indexed_relation_count: usize,
    /// Actual entity count from shards
    pub actual_entity_count: usize,
    /// Actual relation count from shards
    pub actual_relation_count: usize,
}

impl IndexIntegrityResult {
    /// Get a summary message
    pub fn summary(&self) -> String {
        if self.is_valid {
            format!(
                "Index integrity OK: {} entities, {} relations",
                self.indexed_entity_count,
                self.indexed_relation_count
            )
        } else {
            format!(
                "Index integrity MISMATCH: indexed ({} entities, {} relations) vs actual ({} entities, {} relations)",
                self.indexed_entity_count,
                self.indexed_relation_count,
                self.actual_entity_count,
                self.actual_relation_count
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    
    #[test]
    fn test_entity_registration() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register entities
        index.register_entity("entity_1", 0);
        index.register_entity("entity_2", 1);
        index.register_entity("entity_3", 0);
        
        assert_eq!(index.get_entity_shard("entity_1"), Some(0));
        assert_eq!(index.get_entity_shard("entity_2"), Some(1));
        assert_eq!(index.get_entity_shard("entity_3"), Some(0));
        assert_eq!(index.get_entity_shard("nonexistent"), None);
        
        assert_eq!(index.entity_count(), 3);
    }
    
    #[test]
    fn test_entity_with_type() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register entities with types
        index.register_entity_with_type("person_1", "Person", 0);
        index.register_entity_with_type("person_2", "Person", 1);
        index.register_entity_with_type("company_1", "Company", 0);
        
        // Query by type
        let persons = index.get_entities_by_type("Person");
        assert_eq!(persons.len(), 2);
        assert!(persons.contains(&"person_1".to_string()));
        assert!(persons.contains(&"person_2".to_string()));
        
        let companies = index.get_entities_by_type("Company");
        assert_eq!(companies.len(), 1);
        assert!(companies.contains(&"company_1".to_string()));
    }
    
    #[test]
    fn test_relation_registration() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register relations
        index.register_relation("rel_1", 0);
        index.register_relation("rel_2", 1);
        
        assert_eq!(index.get_relation_shard("rel_1"), Some(0));
        assert_eq!(index.get_relation_shard("rel_2"), Some(1));
        
        assert_eq!(index.relation_count(), 2);
    }
    
    #[test]
    fn test_cross_shard_relations() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register cross-shard relation
        // Relation "rel_1" is in shard 0, but target "entity_2" is in shard 1
        index.register_relation_with_target(
            "rel_1",
            "entity_1", // source
            "entity_2", // target
            0,          // relation_shard
            1,          // target_shard
        );
        
        // Another relation to the same target
        index.register_relation_with_target(
            "rel_2",
            "entity_3",
            "entity_2",
            2,
            1,
        );
        
        // Query incoming refs for entity_2
        let refs = index.get_incoming_refs("entity_2");
        assert_eq!(refs.len(), 2);
        
        // Check shards with incoming relations
        let incoming_shards = index.get_incoming_shards("entity_2");
        assert!(incoming_shards.contains(&0));
        assert!(incoming_shards.contains(&2));
        
        // Cross-shard count
        assert_eq!(index.cross_shard_count("entity_2"), 2);
        assert_eq!(index.total_cross_shard_count(), 2);
    }
    
    #[test]
    fn test_remove_operations() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Setup
        index.register_entity_with_type("entity_1", "Person", 0);
        index.register_relation_with_target("rel_1", "entity_1", "entity_2", 0, 1);
        
        assert!(index.contains_entity("entity_1"));
        assert!(index.contains_relation("rel_1"));
        
        // Remove entity
        index.remove_entity("entity_1");
        assert!(!index.contains_entity("entity_1"));
        
        // Remove relation with target
        index.remove_relation("rel_1", Some("entity_2"));
        assert!(!index.contains_relation("rel_1"));
        assert_eq!(index.get_incoming_refs("entity_2").len(), 0);
    }
    
    #[test]
    fn test_batch_operations() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Batch register entities
        let entities = vec![
            ("e1".to_string(), 0),
            ("e2".to_string(), 1),
            ("e3".to_string(), 2),
        ];
        index.register_entities_batch(&entities);
        
        assert_eq!(index.entity_count(), 3);
        assert_eq!(index.get_entity_shard("e1"), Some(0));
        assert_eq!(index.get_entity_shard("e2"), Some(1));
        assert_eq!(index.get_entity_shard("e3"), Some(2));
        
        // Batch register relations
        let relations = vec![
            ("r1".to_string(), 0),
            ("r2".to_string(), 1),
        ];
        index.register_relations_batch(&relations);
        
        assert_eq!(index.relation_count(), 2);
    }
    
    #[test]
    fn test_persistence() {
        let dir = tempdir().unwrap();
        let index_path = dir.path().join("global_index.json");
        
        // Create and populate index
        {
            let config = GlobalIndexConfig::default();
            let index = GlobalGraphIndex::with_path(config, &index_path);
            
            index.register_entity_with_type("person_1", "Person", 0);
            index.register_entity_with_type("person_2", "Person", 1);
            index.register_relation_with_target("rel_1", "person_1", "person_2", 0, 1);
            
            index.persist().unwrap();
        }
        
        // Load into new index
        {
            let config = GlobalIndexConfig::default();
            let index = GlobalGraphIndex::with_path(config, &index_path);
            
            let loaded = index.load().unwrap();
            assert!(loaded);
            
            assert_eq!(index.entity_count(), 2);
            assert_eq!(index.relation_count(), 1);
            assert_eq!(index.get_entity_shard("person_1"), Some(0));
            assert_eq!(index.get_entity_shard("person_2"), Some(1));
            assert_eq!(index.get_entities_by_type("Person").len(), 2);
            assert_eq!(index.get_incoming_refs("person_2").len(), 1);
        }
    }
    
    #[test]
    fn test_stats() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        index.register_entity_with_type("e1", "TypeA", 0);
        index.register_entity_with_type("e2", "TypeB", 1);
        index.register_relation_with_target("r1", "e1", "e2", 0, 1);
        
        let stats = index.stats();
        
        assert!(stats.version > 0);
        assert_eq!(stats.entity_count, 2);
        assert_eq!(stats.relation_count, 1);
        assert_eq!(stats.cross_shard_count, 1);
        assert_eq!(stats.type_count, 2);
    }
    
    // ========================================================================
    // Phase-96: Type cleanup regression tests
    // ========================================================================
    
    #[test]
    fn test_type_change_cleans_old_mapping() {
        // Phase-96: When entity type changes via register_entity_with_type,
        // the old type mapping must be removed.
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register entity with initial type
        index.register_entity_with_type("node-1", "tree_section", 0);
        assert_eq!(index.get_entities_by_type("tree_section").len(), 1);
        assert_eq!(index.get_entity_type("node-1"), Some("tree_section".to_string()));
        
        // Change type (e.g., demote to shadow)
        index.register_entity_with_type("node-1", "tree_section_shadow", 0);
        
        // Old type mapping must be removed
        assert_eq!(
            index.get_entities_by_type("tree_section").len(), 0,
            "Phase-96 bug: entity still in old type set after type change"
        );
        
        // New type mapping must be present
        assert_eq!(index.get_entities_by_type("tree_section_shadow").len(), 1);
        assert_eq!(index.get_entity_type("node-1"), Some("tree_section_shadow".to_string()));
        
        // Total entity count should remain 1
        assert_eq!(index.entity_count(), 1);
    }
    
    #[test]
    fn test_type_change_then_delete_no_resurrection() {
        // Phase-96: Reproduce the exact race condition scenario:
        // put(type=A) → put(type=B) → delete → verify not in any type set
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Step 1: Initial registration
        index.register_entity_with_type("node-1", "tree_section", 0);
        
        // Step 2: Type change (demote)
        index.register_entity_with_type("node-1", "tree_section_shadow", 0);
        
        // Step 3: Delete
        index.remove_entity("node-1");
        
        // Step 4: Verify complete cleanup
        assert_eq!(index.get_entities_by_type("tree_section").len(), 0,
            "Entity still in original type set after delete");
        assert_eq!(index.get_entities_by_type("tree_section_shadow").len(), 0,
            "Entity still in new type set after delete");
        assert_eq!(index.get_entity_type("node-1"), None,
            "Entity type reverse mapping not cleaned after delete");
        assert!(!index.contains_entity("node-1"),
            "Entity still in shard map after delete");
        assert_eq!(index.entity_count(), 0);
    }
    
    #[test]
    fn test_multiple_type_changes_cleanup() {
        // Phase-96: Entity type changes multiple times — only latest should remain
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        index.register_entity_with_type("e1", "TypeA", 0);
        index.register_entity_with_type("e1", "TypeB", 0);
        index.register_entity_with_type("e1", "TypeC", 0);
        index.register_entity_with_type("e1", "TypeD", 0);
        
        // Only TypeD should contain e1
        assert_eq!(index.get_entities_by_type("TypeA").len(), 0);
        assert_eq!(index.get_entities_by_type("TypeB").len(), 0);
        assert_eq!(index.get_entities_by_type("TypeC").len(), 0);
        assert_eq!(index.get_entities_by_type("TypeD").len(), 1);
        assert_eq!(index.get_entity_type("e1"), Some("TypeD".to_string()));
        
        // Type count should be 1 (only TypeD has entries, others were cleaned)
        assert_eq!(index.stats().type_count, 1);
    }
    
    #[test]
    fn test_same_type_registration_idempotent() {
        // Registering the same type again should not cause issues
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        index.register_entity_with_type("e1", "Person", 0);
        index.register_entity_with_type("e1", "Person", 0); // same type again
        
        assert_eq!(index.get_entities_by_type("Person").len(), 1);
        assert_eq!(index.entity_count(), 1);
    }
    
    #[test]
    fn test_batch_type_change_cleanup() {
        // Phase-96: Batch registration should also clean up old types
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Initial registration
        index.register_entity_with_type("e1", "TypeA", 0);
        index.register_entity_with_type("e2", "TypeA", 1);
        
        // Batch change type of e1
        let batch = vec![
            ("e1".to_string(), "TypeB".to_string(), 0),
        ];
        index.register_entities_with_types_batch(&batch);
        
        // e1 should be in TypeB, not TypeA
        let type_a = index.get_entities_by_type("TypeA");
        assert!(!type_a.contains(&"e1".to_string()), "e1 still in TypeA after batch type change");
        assert!(type_a.contains(&"e2".to_string()), "e2 should still be in TypeA");
        
        let type_b = index.get_entities_by_type("TypeB");
        assert!(type_b.contains(&"e1".to_string()), "e1 should be in TypeB");
    }
    
    #[test]
    fn test_persistence_with_entity_type_map() {
        // Phase-96: entity_type_map should survive persistence roundtrip
        let dir = tempdir().unwrap();
        let index_path = dir.path().join("global_index.json");
        
        // Create, populate and persist
        {
            let config = GlobalIndexConfig::default();
            let index = GlobalGraphIndex::with_path(config, &index_path);
            
            index.register_entity_with_type("e1", "Person", 0);
            index.register_entity_with_type("e2", "Company", 1);
            // Type change — only "Manager" should be persisted for e1
            index.register_entity_with_type("e1", "Manager", 0);
            
            index.persist().unwrap();
        }
        
        // Load and verify
        {
            let config = GlobalIndexConfig::default();
            let index = GlobalGraphIndex::with_path(config, &index_path);
            
            let loaded = index.load().unwrap();
            assert!(loaded);
            
            // Verify entity_type_map was restored
            assert_eq!(index.get_entity_type("e1"), Some("Manager".to_string()));
            assert_eq!(index.get_entity_type("e2"), Some("Company".to_string()));
            
            // Verify type_index is consistent
            assert_eq!(index.get_entities_by_type("Person").len(), 0);
            assert_eq!(index.get_entities_by_type("Manager").len(), 1);
            assert_eq!(index.get_entities_by_type("Company").len(), 1);
        }
    }
    
    #[test]
    fn test_backward_compatible_snapshot_loading() {
        // Phase-96: Loading a pre-Phase-96 snapshot (no entity_type_map)
        // should automatically rebuild the reverse map from type_index
        let dir = tempdir().unwrap();
        let index_path = dir.path().join("global_index.json");
        
        // Create a pre-Phase-96 style snapshot (without entity_type_map)
        let old_snapshot = serde_json::json!({
            "version": 42,
            "entity_shards": [["e1", 0], ["e2", 1]],
            "relation_shards": [],
            "cross_shard_incoming": [],
            "type_index": [
                ["Person", ["e1"]],
                ["Company", ["e2"]]
            ],
            "integrity": {
                "expected_entity_count": 2,
                "expected_relation_count": 0,
                "last_rebuild_timestamp": 0,
                "schema_version": 2
            }
        });
        
        std::fs::write(&index_path, serde_json::to_string(&old_snapshot).unwrap()).unwrap();
        
        // Load the old format
        let config = GlobalIndexConfig::default();
        let index = GlobalGraphIndex::with_path(config, &index_path);
        let loaded = index.load().unwrap();
        assert!(loaded);
        
        // entity_type_map should be rebuilt from type_index
        assert_eq!(index.get_entity_type("e1"), Some("Person".to_string()));
        assert_eq!(index.get_entity_type("e2"), Some("Company".to_string()));
    }
    
    // ========================================================================
    // Phase-97: Pagination Tests
    // ========================================================================
    
    #[test]
    fn test_get_entities_by_type_paginated_basic() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        // Register 10 persons
        for i in 0..10 {
            index.register_entity_with_type(&format!("person_{}", i), "Person", i % 3);
        }
        
        // First page
        let (ids, total) = index.get_entities_by_type_paginated("Person", 0, 3);
        assert_eq!(total, 10);
        assert_eq!(ids.len(), 3);
        
        // Middle page
        let (ids, total) = index.get_entities_by_type_paginated("Person", 3, 3);
        assert_eq!(total, 10);
        assert_eq!(ids.len(), 3);
        
        // Last page (partial)
        let (ids, total) = index.get_entities_by_type_paginated("Person", 9, 3);
        assert_eq!(total, 10);
        assert_eq!(ids.len(), 1);
        
        // Full page
        let (ids, total) = index.get_entities_by_type_paginated("Person", 0, 100);
        assert_eq!(total, 10);
        assert_eq!(ids.len(), 10);
    }
    
    #[test]
    fn test_get_entities_by_type_paginated_empty() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        let (ids, total) = index.get_entities_by_type_paginated("NonExistent", 0, 10);
        assert_eq!(total, 0);
        assert_eq!(ids.len(), 0);
    }
    
    #[test]
    fn test_get_entities_by_type_paginated_offset_beyond() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        for i in 0..5 {
            index.register_entity_with_type(&format!("e_{}", i), "Tag", 0);
        }
        
        // offset beyond total
        let (ids, total) = index.get_entities_by_type_paginated("Tag", 100, 10);
        assert_eq!(total, 5);
        assert_eq!(ids.len(), 0);
    }
    
    #[test]
    fn test_count_entities_by_type() {
        let index = GlobalGraphIndex::new(GlobalIndexConfig::default());
        
        for i in 0..7 {
            index.register_entity_with_type(&format!("p_{}", i), "Person", 0);
        }
        for i in 0..3 {
            index.register_entity_with_type(&format!("c_{}", i), "Company", 0);
        }
        
        assert_eq!(index.count_entities_by_type("Person"), 7);
        assert_eq!(index.count_entities_by_type("Company"), 3);
        assert_eq!(index.count_entities_by_type("NonExistent"), 0);
    }
    
    #[test]
    fn test_get_entities_by_type_paginated_type_index_disabled() {
        let config = GlobalIndexConfig {
            enable_type_index: false,
            ..Default::default()
        };
        let index = GlobalGraphIndex::new(config);
        
        index.register_entity_with_type("e1", "Person", 0);
        
        let (ids, total) = index.get_entities_by_type_paginated("Person", 0, 10);
        assert_eq!(total, 0);
        assert_eq!(ids.len(), 0);
        assert_eq!(index.count_entities_by_type("Person"), 0);
    }
}
