//! Graph Memory System - Inspired by Graphiti
//!
//! Temporal knowledge graph for tracking entities, relationships, and episodic memories.
//! Implements bi-temporal tracking and hybrid retrieval (semantic + graph traversal).

use anyhow::Result;
use chrono::{DateTime, Utc};
use rocksdb::{Options, DB};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use uuid::Uuid;

/// Entity node in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntityNode {
    /// Unique identifier
    pub uuid: Uuid,

    /// Entity name (e.g., "John", "Paris", "Rust programming")
    pub name: String,

    /// Entity labels/types (e.g., ["Person"], ["Location", "City"])
    pub labels: Vec<EntityLabel>,

    /// When this entity was first created in the graph
    pub created_at: DateTime<Utc>,

    /// When this entity was last observed
    pub last_seen_at: DateTime<Utc>,

    /// How many times this entity has been mentioned
    pub mention_count: usize,

    /// Summary of this entity's context (built from surrounding edges)
    pub summary: String,

    /// Additional attributes based on entity type
    pub attributes: HashMap<String, String>,

    /// Semantic embedding of the entity name (for similarity search)
    pub name_embedding: Option<Vec<f32>>,

    /// Salience score (0.0 - 1.0): How important is this entity?
    /// Higher salience = larger gravitational well in the memory universe
    /// Factors: proper noun status, mention frequency, recency, user-defined importance
    #[serde(default = "default_salience")]
    pub salience: f32,

    /// Whether this is a proper noun (names, places, products)
    /// Proper nouns have higher base salience than common nouns
    #[serde(default)]
    pub is_proper_noun: bool,
}

fn default_salience() -> f32 {
    0.5 // Default middle salience
}

/// Entity labels following Graphiti's categorization
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum EntityLabel {
    Person,
    Organization,
    Location,
    Technology,
    Concept,
    Event,
    Date,
    Product,
    Skill,
    Other(String),
}

impl EntityLabel {
    /// Get string representation of the entity label
    #[allow(unused)] // Public API for serialization/display
    pub fn as_str(&self) -> &str {
        match self {
            Self::Person => "Person",
            Self::Organization => "Organization",
            Self::Location => "Location",
            Self::Technology => "Technology",
            Self::Concept => "Concept",
            Self::Event => "Event",
            Self::Date => "Date",
            Self::Product => "Product",
            Self::Skill => "Skill",
            Self::Other(s) => s.as_str(),
        }
    }
}

/// Relationship edge between entities
///
/// Implements Hebbian synaptic plasticity: "Neurons that fire together, wire together"
/// - Strength increases with co-activation (strengthen method)
/// - Strength decays over time without use (decay method)
/// - Long-Term Potentiation (LTP): After threshold activations, becomes permanent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelationshipEdge {
    /// Unique identifier for this edge
    pub uuid: Uuid,

    /// Source entity UUID
    pub from_entity: Uuid,

    /// Target entity UUID
    pub to_entity: Uuid,

    /// Type of relationship
    pub relation_type: RelationType,

    /// Confidence/strength of this relationship (0.0 to 1.0)
    /// Dynamic: increases with co-activation, decays without use
    pub strength: f32,

    /// When this relationship was created
    pub created_at: DateTime<Utc>,

    /// When this relationship was last observed (temporal tracking)
    pub valid_at: DateTime<Utc>,

    /// Whether this relationship has been invalidated (temporal edge invalidation)
    pub invalidated_at: Option<DateTime<Utc>>,

    /// Source episode that created this relationship
    pub source_episode_id: Option<Uuid>,

    /// Additional context about the relationship
    pub context: String,

    // === Hebbian Synaptic Plasticity Fields ===
    /// When this synapse was last activated (used in retrieval/traversal)
    /// Used to calculate time-based decay
    #[serde(default = "default_last_activated")]
    pub last_activated: DateTime<Utc>,

    /// Number of times both entities were co-accessed (Hebbian co-activation)
    /// Higher count = stronger learned association
    #[serde(default)]
    pub activation_count: u32,

    /// Long-Term Potentiation flag: synapse becomes permanent after threshold
    /// Once potentiated, decay is dramatically reduced (like biological LTP)
    #[serde(default)]
    pub potentiated: bool,
}

fn default_last_activated() -> DateTime<Utc> {
    Utc::now()
}

/// Hebbian learning constants
const LEARNING_RATE: f32 = 0.1; // η: How much strength increases per co-activation
const DECAY_HALF_LIFE_DAYS: f64 = 14.0; // λ: Strength halves every 14 days without use
const LTP_THRESHOLD: u32 = 10; // Activations needed for Long-Term Potentiation
const LTP_DECAY_FACTOR: f32 = 0.1; // Potentiated synapses decay 10x slower
const MIN_STRENGTH: f32 = 0.01; // Floor to prevent complete forgetting

impl RelationshipEdge {
    /// Strengthen this synapse (Hebbian learning)
    ///
    /// Called when both connected entities are accessed together.
    /// Formula: w_new = w_old + η × (1 - w_old) × co_activation_boost
    ///
    /// The (1 - w_old) term ensures asymptotic approach to 1.0,
    /// preventing unbounded growth while allowing strong associations.
    pub fn strengthen(&mut self) {
        self.activation_count += 1;
        self.last_activated = Utc::now();

        // Hebbian strengthening: diminishing returns as strength approaches 1.0
        let boost = LEARNING_RATE * (1.0 - self.strength);
        self.strength = (self.strength + boost).min(1.0);

        // Check for Long-Term Potentiation threshold
        if !self.potentiated && self.activation_count >= LTP_THRESHOLD {
            self.potentiated = true;
            // LTP bonus: immediate strength boost
            self.strength = (self.strength + 0.2).min(1.0);
        }
    }

    /// Apply time-based decay to this synapse
    ///
    /// Formula: w(t) = w₀ × e^(-λt) where λ = ln(2) / half_life
    /// Potentiated synapses decay 10x slower (LTP protection)
    ///
    /// Returns true if synapse should be pruned (strength below threshold)
    pub fn decay(&mut self) -> bool {
        let now = Utc::now();
        let elapsed = now.signed_duration_since(self.last_activated);
        let days_elapsed = elapsed.num_seconds() as f64 / 86400.0;

        if days_elapsed <= 0.0 {
            return false;
        }

        // Calculate decay rate (λ = ln(2) / half_life)
        let lambda = std::f64::consts::LN_2 / DECAY_HALF_LIFE_DAYS;

        // Potentiated synapses decay much slower (biological LTP)
        let effective_lambda = if self.potentiated {
            lambda * LTP_DECAY_FACTOR as f64
        } else {
            lambda
        };

        // Exponential decay: w(t) = w₀ × e^(-λt)
        let decay_factor = (-effective_lambda * days_elapsed).exp() as f32;
        self.strength *= decay_factor;

        // Apply floor to prevent complete forgetting
        if self.strength < MIN_STRENGTH {
            self.strength = MIN_STRENGTH;
        }

        // Return whether this synapse should be pruned
        // Non-potentiated synapses with minimal strength can be removed
        !self.potentiated && self.strength <= MIN_STRENGTH
    }

    /// Get the effective strength considering recency
    ///
    /// This is a read-only version that calculates what the strength
    /// would be after decay, without modifying the edge.
    pub fn effective_strength(&self) -> f32 {
        let now = Utc::now();
        let elapsed = now.signed_duration_since(self.last_activated);
        let days_elapsed = elapsed.num_seconds() as f64 / 86400.0;

        if days_elapsed <= 0.0 {
            return self.strength;
        }

        let lambda = std::f64::consts::LN_2 / DECAY_HALF_LIFE_DAYS;
        let effective_lambda = if self.potentiated {
            lambda * LTP_DECAY_FACTOR as f64
        } else {
            lambda
        };

        let decay_factor = (-effective_lambda * days_elapsed).exp() as f32;
        (self.strength * decay_factor).max(MIN_STRENGTH)
    }
}

/// Relationship types following Graphiti's semantic model
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum RelationType {
    /// Work relationships
    WorksWith,
    WorksAt,
    EmployedBy,

    /// Structural relationships
    PartOf,
    Contains,
    OwnedBy,

    /// Location relationships
    LocatedIn,
    LocatedAt,

    /// Usage relationships
    Uses,
    CreatedBy,
    DevelopedBy,

    /// Causal relationships
    Causes,
    ResultsIn,

    /// Learning relationships
    Learned,
    Knows,
    Teaches,

    /// Generic relationships
    RelatedTo,
    AssociatedWith,

    /// Custom relationship
    Custom(String),
}

impl RelationType {
    /// Get string representation of the relation type
    #[allow(unused)] // Public API for serialization/display
    pub fn as_str(&self) -> &str {
        match self {
            Self::WorksWith => "WorksWith",
            Self::WorksAt => "WorksAt",
            Self::EmployedBy => "EmployedBy",
            Self::PartOf => "PartOf",
            Self::Contains => "Contains",
            Self::OwnedBy => "OwnedBy",
            Self::LocatedIn => "LocatedIn",
            Self::LocatedAt => "LocatedAt",
            Self::Uses => "Uses",
            Self::CreatedBy => "CreatedBy",
            Self::DevelopedBy => "DevelopedBy",
            Self::Causes => "Causes",
            Self::ResultsIn => "ResultsIn",
            Self::Learned => "Learned",
            Self::Knows => "Knows",
            Self::Teaches => "Teaches",
            Self::RelatedTo => "RelatedTo",
            Self::AssociatedWith => "AssociatedWith",
            Self::Custom(s) => s.as_str(),
        }
    }
}

/// Episodic node representing a discrete experience/memory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpisodicNode {
    /// Unique identifier
    pub uuid: Uuid,

    /// Human-readable name/title
    pub name: String,

    /// Episode content (the actual experience data)
    pub content: String,

    /// When the original event occurred (event time)
    pub valid_at: DateTime<Utc>,

    /// When this was ingested into the system (ingestion time)
    pub created_at: DateTime<Utc>,

    /// Entities extracted from this episode
    pub entity_refs: Vec<Uuid>,

    /// Source type
    pub source: EpisodeSource,

    /// Additional metadata
    pub metadata: HashMap<String, String>,
}

/// Episode source types
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum EpisodeSource {
    Message,
    Document,
    Event,
    Observation,
}

/// Graph memory storage and operations
pub struct GraphMemory {
    /// RocksDB storage for entities
    entities_db: Arc<DB>,

    /// RocksDB storage for relationships
    relationships_db: Arc<DB>,

    /// RocksDB storage for episodes
    episodes_db: Arc<DB>,

    /// RocksDB storage for entity -> relationships index
    entity_edges_db: Arc<DB>,

    /// RocksDB storage for entity -> episodes index (inverted index for fast lookup)
    entity_episodes_db: Arc<DB>,

    /// RocksDB storage for entity name -> UUID index (persisted, O(1) startup)
    entity_name_index_db: Arc<DB>,

    /// In-memory entity name index for fast lookups (loaded from entity_name_index_db)
    entity_name_index: Arc<parking_lot::RwLock<HashMap<String, Uuid>>>,

    // === Atomic counters for O(1) stats (P1 fix) ===
    /// Entity count - initialized from entity_name_index.len(), updated on add
    entity_count: Arc<AtomicUsize>,

    /// Relationship count - initialized on startup, updated on add
    relationship_count: Arc<AtomicUsize>,

    /// Episode count - initialized on startup, updated on add
    episode_count: Arc<AtomicUsize>,
}

impl GraphMemory {
    /// Create a new graph memory system
    pub fn new(path: &Path) -> Result<Self> {
        std::fs::create_dir_all(path)?;

        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_compression_type(rocksdb::DBCompressionType::Lz4);

        let entities_db = Arc::new(DB::open(&opts, path.join("graph_entities"))?);
        let relationships_db = Arc::new(DB::open(&opts, path.join("graph_relationships"))?);
        let episodes_db = Arc::new(DB::open(&opts, path.join("graph_episodes"))?);
        let entity_edges_db = Arc::new(DB::open(&opts, path.join("graph_entity_edges"))?);
        let entity_episodes_db = Arc::new(DB::open(&opts, path.join("graph_entity_episodes"))?);
        let entity_name_index_db = Arc::new(DB::open(&opts, path.join("graph_entity_name_index"))?);

        // Load entity name index from persisted DB (O(n) but faster than deserializing entities)
        // If empty, migrate from entities_db (one-time migration for existing data)
        let entity_name_index =
            Self::load_or_migrate_name_index(&entity_name_index_db, &entities_db)?;
        let entity_count = entity_name_index.len();

        // Count relationships and episodes during startup (one-time cost)
        // This is O(n) at startup, but get_stats() will be O(1) at runtime
        let relationship_count = Self::count_db_entries(&relationships_db);
        let episode_count = Self::count_db_entries(&episodes_db);

        let graph = Self {
            entities_db,
            relationships_db,
            episodes_db,
            entity_edges_db,
            entity_episodes_db,
            entity_name_index_db,
            entity_name_index: Arc::new(parking_lot::RwLock::new(entity_name_index)),
            entity_count: Arc::new(AtomicUsize::new(entity_count)),
            relationship_count: Arc::new(AtomicUsize::new(relationship_count)),
            episode_count: Arc::new(AtomicUsize::new(episode_count)),
        };

        if entity_count > 0 || relationship_count > 0 || episode_count > 0 {
            tracing::info!(
                "Loaded graph with {} entities, {} relationships, {} episodes",
                entity_count,
                relationship_count,
                episode_count
            );
        }

        Ok(graph)
    }

    /// Load entity name->UUID index from persisted DB, or migrate from entities_db if empty
    fn load_or_migrate_name_index(
        index_db: &DB,
        entities_db: &DB,
    ) -> Result<HashMap<String, Uuid>> {
        let mut index = HashMap::new();

        // Try to load from dedicated index DB first
        let iter = index_db.iterator(rocksdb::IteratorMode::Start);
        for (key, value) in iter.flatten() {
            if let (Ok(name), Ok(uuid_bytes)) = (
                std::str::from_utf8(&key),
                <[u8; 16]>::try_from(value.as_ref()),
            ) {
                index.insert(name.to_string(), Uuid::from_bytes(uuid_bytes));
            }
        }

        // If index DB is empty but entities exist, migrate (one-time operation)
        if index.is_empty() {
            let entity_iter = entities_db.iterator(rocksdb::IteratorMode::Start);
            let mut migrated_count = 0;
            for (_, value) in entity_iter.flatten() {
                if let Ok(entity) = bincode::deserialize::<EntityNode>(&value) {
                    // Store in index DB: name -> UUID bytes
                    index_db.put(entity.name.as_bytes(), entity.uuid.as_bytes())?;
                    index.insert(entity.name.clone(), entity.uuid);
                    migrated_count += 1;
                }
            }
            if migrated_count > 0 {
                tracing::info!("Migrated {} entities to name index DB", migrated_count);
            }
        }

        Ok(index)
    }

    /// Count entries in a RocksDB database (one-time startup cost)
    fn count_db_entries(db: &DB) -> usize {
        db.iterator(rocksdb::IteratorMode::Start).count()
    }

    /// Add or update an entity node
    /// Salience is updated using the formula: salience = base_salience * (1 + 0.1 * ln(mention_count))
    /// This means frequently mentioned entities grow in salience (gravitational wells get heavier)
    pub fn add_entity(&self, mut entity: EntityNode) -> Result<Uuid> {
        // Check if entity already exists by name
        let existing_uuid = {
            let index = self.entity_name_index.read();
            index.get(&entity.name).cloned()
        };

        if let Some(uuid) = existing_uuid {
            // Update existing entity
            entity.uuid = uuid;
            if let Some(existing) = self.get_entity(&uuid)? {
                entity.mention_count = existing.mention_count + 1;
                entity.last_seen_at = Utc::now();
                entity.created_at = existing.created_at; // Preserve original creation time
                entity.is_proper_noun = existing.is_proper_noun || entity.is_proper_noun;

                // Update salience with frequency boost
                // Formula: salience = base_salience * (1 + 0.1 * ln(mention_count))
                // This caps at about 1.3x boost at 20 mentions
                let frequency_boost = 1.0 + 0.1 * (entity.mention_count as f32).ln();
                entity.salience = (existing.salience * frequency_boost).min(1.0);
            }
        } else {
            // New entity
            entity.uuid = Uuid::new_v4();
            entity.created_at = Utc::now();
            entity.last_seen_at = entity.created_at;
            entity.mention_count = 1;
            // Salience stays at base_salience for new entities

            // Increment entity counter for new entities only
            self.entity_count.fetch_add(1, Ordering::Relaxed);
        }

        // Store in database
        let key = entity.uuid.as_bytes();
        let value = bincode::serialize(&entity)?;
        self.entities_db.put(key, value)?;

        // Update both in-memory index and persisted index DB
        {
            let mut index = self.entity_name_index.write();
            index.insert(entity.name.clone(), entity.uuid);
        }
        // Persist name->UUID mapping for O(1) startup
        self.entity_name_index_db
            .put(entity.name.as_bytes(), entity.uuid.as_bytes())?;

        Ok(entity.uuid)
    }

    /// Get entity by UUID
    pub fn get_entity(&self, uuid: &Uuid) -> Result<Option<EntityNode>> {
        let key = uuid.as_bytes();
        match self.entities_db.get(key)? {
            Some(value) => {
                let entity: EntityNode = bincode::deserialize(&value)?;
                Ok(Some(entity))
            }
            None => Ok(None),
        }
    }

    /// Find entity by name
    pub fn find_entity_by_name(&self, name: &str) -> Result<Option<EntityNode>> {
        let uuid = {
            let index = self.entity_name_index.read();
            index.get(name).cloned()
        };

        match uuid {
            Some(uuid) => self.get_entity(&uuid),
            None => Ok(None),
        }
    }

    /// Add a relationship edge
    pub fn add_relationship(&self, mut edge: RelationshipEdge) -> Result<Uuid> {
        edge.uuid = Uuid::new_v4();
        edge.created_at = Utc::now();

        // Store relationship
        let key = edge.uuid.as_bytes();
        let value = bincode::serialize(&edge)?;
        self.relationships_db.put(key, value)?;

        // Increment relationship counter
        self.relationship_count.fetch_add(1, Ordering::Relaxed);

        // Update entity->edges index for both entities
        self.index_entity_edge(&edge.from_entity, &edge.uuid)?;
        self.index_entity_edge(&edge.to_entity, &edge.uuid)?;

        Ok(edge.uuid)
    }

    /// Index an edge for an entity
    fn index_entity_edge(&self, entity_uuid: &Uuid, edge_uuid: &Uuid) -> Result<()> {
        let key = format!("{entity_uuid}:{edge_uuid}");
        self.entity_edges_db.put(key.as_bytes(), b"1")?;
        Ok(())
    }

    /// Get all relationships for an entity
    pub fn get_entity_relationships(&self, entity_uuid: &Uuid) -> Result<Vec<RelationshipEdge>> {
        let mut edges = Vec::new();
        let prefix = format!("{entity_uuid}:");

        let iter = self.entity_edges_db.prefix_iterator(prefix.as_bytes());
        for (key, _) in iter.flatten() {
            if let Ok(key_str) = std::str::from_utf8(&key) {
                if !key_str.starts_with(&prefix) {
                    break;
                }

                // Extract edge UUID
                if let Some(edge_uuid_str) = key_str.split(':').nth(1) {
                    if let Ok(edge_uuid) = Uuid::parse_str(edge_uuid_str) {
                        if let Some(edge) = self.get_relationship(&edge_uuid)? {
                            edges.push(edge);
                        }
                    }
                }
            }
        }

        Ok(edges)
    }

    /// Get relationship by UUID (raw, without decay applied)
    pub fn get_relationship(&self, uuid: &Uuid) -> Result<Option<RelationshipEdge>> {
        let key = uuid.as_bytes();
        match self.relationships_db.get(key)? {
            Some(value) => {
                let edge: RelationshipEdge = bincode::deserialize(&value)?;
                Ok(Some(edge))
            }
            None => Ok(None),
        }
    }

    /// Get relationship by UUID with effective strength (lazy decay calculation)
    ///
    /// Returns the edge with strength reflecting time-based decay.
    /// This doesn't persist the decay - just calculates what the strength would be.
    /// Use this for API responses to show accurate current strength.
    pub fn get_relationship_with_effective_strength(
        &self,
        uuid: &Uuid,
    ) -> Result<Option<RelationshipEdge>> {
        let key = uuid.as_bytes();
        match self.relationships_db.get(key)? {
            Some(value) => {
                let mut edge: RelationshipEdge = bincode::deserialize(&value)?;
                // Apply effective strength calculation (doesn't persist)
                edge.strength = edge.effective_strength();
                Ok(Some(edge))
            }
            None => Ok(None),
        }
    }

    /// Add an episodic node
    pub fn add_episode(&self, episode: EpisodicNode) -> Result<Uuid> {
        let key = episode.uuid.as_bytes();
        let value = bincode::serialize(&episode)?;
        self.episodes_db.put(key, value)?;

        // Increment episode counter
        self.episode_count.fetch_add(1, Ordering::Relaxed);

        // Update inverted index: entity_uuid -> episode_uuid
        for entity_uuid in &episode.entity_refs {
            self.index_entity_episode(entity_uuid, &episode.uuid)?;
        }

        Ok(episode.uuid)
    }

    /// Index an episode for an entity (inverted index)
    fn index_entity_episode(&self, entity_uuid: &Uuid, episode_uuid: &Uuid) -> Result<()> {
        let key = format!("{entity_uuid}:{episode_uuid}");
        self.entity_episodes_db.put(key.as_bytes(), b"1")?;
        Ok(())
    }

    /// Get episode by UUID
    pub fn get_episode(&self, uuid: &Uuid) -> Result<Option<EpisodicNode>> {
        let key = uuid.as_bytes();
        match self.episodes_db.get(key)? {
            Some(value) => {
                let episode: EpisodicNode = bincode::deserialize(&value)?;
                Ok(Some(episode))
            }
            None => Ok(None),
        }
    }

    /// Get all episodes that contain a specific entity
    ///
    /// Uses inverted index for O(k) lookup instead of O(n) full scan.
    /// Crucial for spreading activation algorithm.
    pub fn get_episodes_by_entity(&self, entity_uuid: &Uuid) -> Result<Vec<EpisodicNode>> {
        let mut episodes = Vec::new();
        let prefix = format!("{entity_uuid}:");

        // Use inverted index: entity_uuid -> episode_uuids
        let iter = self.entity_episodes_db.prefix_iterator(prefix.as_bytes());
        for (key, _) in iter.flatten() {
            if let Ok(key_str) = std::str::from_utf8(&key) {
                if !key_str.starts_with(&prefix) {
                    break; // Prefix iterator exhausted
                }

                // Extract episode UUID from key
                if let Some(episode_uuid_str) = key_str.split(':').nth(1) {
                    if let Ok(episode_uuid) = Uuid::parse_str(episode_uuid_str) {
                        if let Some(episode) = self.get_episode(&episode_uuid)? {
                            episodes.push(episode);
                        }
                    }
                }
            }
        }

        Ok(episodes)
    }

    /// Traverse graph starting from an entity (breadth-first)
    ///
    /// Implements Hebbian learning: edges traversed during retrieval are strengthened.
    /// This means frequently accessed pathways become stronger over time.
    pub fn traverse_from_entity(
        &self,
        start_uuid: &Uuid,
        max_depth: usize,
    ) -> Result<GraphTraversal> {
        let mut visited_entities = HashSet::new();
        let mut visited_edges = HashSet::new();
        let mut current_level = vec![*start_uuid];
        let mut all_entities = Vec::new();
        let mut all_edges = Vec::new();
        let mut edges_to_strengthen = Vec::new();

        visited_entities.insert(*start_uuid);
        if let Some(entity) = self.get_entity(start_uuid)? {
            all_entities.push(entity);
        }

        for _ in 0..max_depth {
            let mut next_level = Vec::new();

            for entity_uuid in current_level {
                let edges = self.get_entity_relationships(&entity_uuid)?;

                for edge in edges {
                    if visited_edges.contains(&edge.uuid) {
                        continue;
                    }

                    visited_edges.insert(edge.uuid);

                    // Only traverse non-invalidated edges
                    if edge.invalidated_at.is_some() {
                        continue;
                    }

                    // Collect edge UUID for Hebbian strengthening
                    edges_to_strengthen.push(edge.uuid);

                    // Return edge with effective strength (lazy decay calculation)
                    let mut edge_with_decay = edge.clone();
                    edge_with_decay.strength = edge_with_decay.effective_strength();
                    all_edges.push(edge_with_decay);

                    // Add connected entity
                    let connected_uuid = if edge.from_entity == entity_uuid {
                        edge.to_entity
                    } else {
                        edge.from_entity
                    };

                    if !visited_entities.contains(&connected_uuid) {
                        visited_entities.insert(connected_uuid);
                        if let Some(entity) = self.get_entity(&connected_uuid)? {
                            all_entities.push(entity);
                        }
                        next_level.push(connected_uuid);
                    }
                }
            }

            if next_level.is_empty() {
                break;
            }

            current_level = next_level;
        }

        // Apply Hebbian strengthening to all traversed edges
        // "Neurons that fire together, wire together"
        for edge_uuid in edges_to_strengthen {
            if let Err(e) = self.strengthen_synapse(&edge_uuid) {
                tracing::debug!("Failed to strengthen synapse {}: {}", edge_uuid, e);
            }
        }

        Ok(GraphTraversal {
            entities: all_entities,
            relationships: all_edges,
        })
    }

    /// Invalidate a relationship (temporal edge invalidation)
    pub fn invalidate_relationship(&self, edge_uuid: &Uuid) -> Result<()> {
        if let Some(mut edge) = self.get_relationship(edge_uuid)? {
            edge.invalidated_at = Some(Utc::now());

            let key = edge.uuid.as_bytes();
            let value = bincode::serialize(&edge)?;
            self.relationships_db.put(key, value)?;
        }

        Ok(())
    }

    /// Strengthen a synapse (Hebbian learning)
    ///
    /// Called when an edge is traversed during memory retrieval.
    /// Implements "neurons that fire together, wire together".
    pub fn strengthen_synapse(&self, edge_uuid: &Uuid) -> Result<()> {
        if let Some(mut edge) = self.get_relationship(edge_uuid)? {
            edge.strengthen();

            let key = edge.uuid.as_bytes();
            let value = bincode::serialize(&edge)?;
            self.relationships_db.put(key, value)?;
        }

        Ok(())
    }

    /// Apply decay to a synapse
    ///
    /// Returns true if the synapse should be pruned (non-potentiated and below threshold)
    pub fn decay_synapse(&self, edge_uuid: &Uuid) -> Result<bool> {
        if let Some(mut edge) = self.get_relationship(edge_uuid)? {
            let should_prune = edge.decay();

            let key = edge.uuid.as_bytes();
            let value = bincode::serialize(&edge)?;
            self.relationships_db.put(key, value)?;

            return Ok(should_prune);
        }

        Ok(false)
    }

    /// Get graph statistics - O(1) using atomic counters
    pub fn get_stats(&self) -> Result<GraphStats> {
        Ok(GraphStats {
            entity_count: self.entity_count.load(Ordering::Relaxed),
            relationship_count: self.relationship_count.load(Ordering::Relaxed),
            episode_count: self.episode_count.load(Ordering::Relaxed),
        })
    }

    /// Get all entities in the graph
    pub fn get_all_entities(&self) -> Result<Vec<EntityNode>> {
        let mut entities = Vec::new();

        let iter = self.entities_db.iterator(rocksdb::IteratorMode::Start);
        for (_, value) in iter.flatten() {
            if let Ok(entity) = bincode::deserialize::<EntityNode>(&value) {
                entities.push(entity);
            }
        }

        // Sort by mention count (most mentioned first)
        entities.sort_by(|a, b| b.mention_count.cmp(&a.mention_count));

        Ok(entities)
    }

    /// Get all relationships in the graph
    pub fn get_all_relationships(&self) -> Result<Vec<RelationshipEdge>> {
        let mut relationships = Vec::new();

        let iter = self.relationships_db.iterator(rocksdb::IteratorMode::Start);
        for (_, value) in iter.flatten() {
            if let Ok(edge) = bincode::deserialize::<RelationshipEdge>(&value) {
                // Only include non-invalidated relationships
                if edge.invalidated_at.is_none() {
                    relationships.push(edge);
                }
            }
        }

        // Sort by strength (strongest first)
        relationships.sort_by(|a, b| {
            b.strength
                .partial_cmp(&a.strength)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        Ok(relationships)
    }

    /// Get the Memory Universe visualization data
    /// Returns entities as "stars" with positions based on their relationships,
    /// sized by salience, and colored by entity type.
    pub fn get_universe(&self) -> Result<MemoryUniverse> {
        let entities = self.get_all_entities()?;
        let relationships = self.get_all_relationships()?;

        // Create entity UUID to index mapping for position calculation
        let entity_indices: HashMap<Uuid, usize> = entities
            .iter()
            .enumerate()
            .map(|(i, e)| (e.uuid, i))
            .collect();

        // Calculate 3D positions using a force-directed layout approximation
        // High-salience entities are positioned more centrally
        let mut stars: Vec<UniverseStar> = entities
            .iter()
            .enumerate()
            .map(|(i, entity)| {
                // Use a spiral galaxy layout with salience affecting radius
                // Higher salience = closer to center
                let angle = (i as f32) * 2.4; // Golden angle for even distribution
                let base_radius = 1.0 - entity.salience; // High salience = small radius
                let radius = base_radius * 100.0 + 10.0; // 10-110 range

                let x = radius * angle.cos();
                let y = radius * angle.sin();
                let z = ((i as f32) * 0.1).sin() * 20.0; // Slight z variation

                UniverseStar {
                    id: entity.uuid.to_string(),
                    name: entity.name.clone(),
                    entity_type: entity.labels.first().map(|l| l.as_str().to_string()),
                    salience: entity.salience,
                    mention_count: entity.mention_count,
                    is_proper_noun: entity.is_proper_noun,
                    position: Position3D { x, y, z },
                    color: entity_type_color(entity.labels.first()),
                    size: 5.0 + entity.salience * 20.0, // Size 5-25 based on salience
                }
            })
            .collect();

        // Apply gravitational forces FIRST, before creating connections
        // This ensures connection positions match final star positions
        for rel in &relationships {
            if let (Some(from_idx), Some(to_idx)) = (
                entity_indices.get(&rel.from_entity),
                entity_indices.get(&rel.to_entity),
            ) {
                // Apply small gravitational pull based on connection strength
                let pull_factor = rel.strength * 0.05;

                let from_pos = stars[*from_idx].position.clone();
                let to_pos = stars[*to_idx].position.clone();

                let dx = (to_pos.x - from_pos.x) * pull_factor;
                let dy = (to_pos.y - from_pos.y) * pull_factor;
                let dz = (to_pos.z - from_pos.z) * pull_factor;

                stars[*from_idx].position.x += dx;
                stars[*from_idx].position.y += dy;
                stars[*from_idx].position.z += dz;

                stars[*to_idx].position.x -= dx;
                stars[*to_idx].position.y -= dy;
                stars[*to_idx].position.z -= dz;
            }
        }

        // Create gravitational connections AFTER star positions are finalized
        // This ensures from_position/to_position match current star positions
        let connections: Vec<GravitationalConnection> = relationships
            .iter()
            .filter_map(|rel| {
                let from_idx = entity_indices.get(&rel.from_entity)?;
                let to_idx = entity_indices.get(&rel.to_entity)?;

                Some(GravitationalConnection {
                    id: rel.uuid.to_string(),
                    from_id: rel.from_entity.to_string(),
                    to_id: rel.to_entity.to_string(),
                    relation_type: rel.relation_type.as_str().to_string(),
                    strength: rel.strength,
                    from_position: stars[*from_idx].position.clone(),
                    to_position: stars[*to_idx].position.clone(),
                })
            })
            .collect();

        // Calculate universe bounds
        let (min_x, max_x, min_y, max_y, min_z, max_z) = stars.iter().fold(
            (f32::MAX, f32::MIN, f32::MAX, f32::MIN, f32::MAX, f32::MIN),
            |(min_x, max_x, min_y, max_y, min_z, max_z), star| {
                (
                    min_x.min(star.position.x),
                    max_x.max(star.position.x),
                    min_y.min(star.position.y),
                    max_y.max(star.position.y),
                    min_z.min(star.position.z),
                    max_z.max(star.position.z),
                )
            },
        );

        Ok(MemoryUniverse {
            stars,
            connections,
            total_entities: entities.len(),
            total_connections: relationships.len(),
            bounds: UniverseBounds {
                min: Position3D {
                    x: min_x,
                    y: min_y,
                    z: min_z,
                },
                max: Position3D {
                    x: max_x,
                    y: max_y,
                    z: max_z,
                },
            },
        })
    }
}

/// Helper function to get color for entity type
fn entity_type_color(label: Option<&EntityLabel>) -> String {
    match label {
        Some(EntityLabel::Person) => "#FF6B6B".to_string(), // Coral red
        Some(EntityLabel::Organization) => "#4ECDC4".to_string(), // Teal
        Some(EntityLabel::Location) => "#45B7D1".to_string(), // Sky blue
        Some(EntityLabel::Technology) => "#96CEB4".to_string(), // Sage green
        Some(EntityLabel::Product) => "#FFEAA7".to_string(), // Soft yellow
        Some(EntityLabel::Event) => "#DDA0DD".to_string(),  // Plum
        Some(EntityLabel::Skill) => "#98D8C8".to_string(),  // Mint
        Some(EntityLabel::Concept) => "#F7DC6F".to_string(), // Gold
        Some(EntityLabel::Date) => "#BB8FCE".to_string(),   // Light purple
        Some(EntityLabel::Other(_)) => "#AEB6BF".to_string(), // Gray
        None => "#AEB6BF".to_string(),                      // Gray default
    }
}

/// 3D position in the memory universe
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Position3D {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

/// A star in the memory universe (represents an entity)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UniverseStar {
    pub id: String,
    pub name: String,
    pub entity_type: Option<String>,
    pub salience: f32,
    pub mention_count: usize,
    pub is_proper_noun: bool,
    pub position: Position3D,
    pub color: String,
    pub size: f32,
}

/// A gravitational connection between stars (represents a relationship)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GravitationalConnection {
    pub id: String,
    pub from_id: String,
    pub to_id: String,
    pub relation_type: String,
    pub strength: f32,
    pub from_position: Position3D,
    pub to_position: Position3D,
}

/// Bounds of the memory universe
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UniverseBounds {
    pub min: Position3D,
    pub max: Position3D,
}

/// The complete memory universe visualization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryUniverse {
    pub stars: Vec<UniverseStar>,
    pub connections: Vec<GravitationalConnection>,
    pub total_entities: usize,
    pub total_connections: usize,
    pub bounds: UniverseBounds,
}

/// Result of graph traversal
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphTraversal {
    pub entities: Vec<EntityNode>,
    pub relationships: Vec<RelationshipEdge>,
}

/// Graph statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphStats {
    pub entity_count: usize,
    pub relationship_count: usize,
    pub episode_count: usize,
}

/// Extracted entity with salience information
#[derive(Debug, Clone)]
pub struct ExtractedEntity {
    pub name: String,
    pub label: EntityLabel,
    pub is_proper_noun: bool,
    pub base_salience: f32,
}

/// Simple entity extraction (rule-based NER) with salience detection
pub struct EntityExtractor {
    /// Common person name indicators
    person_indicators: HashSet<String>,

    /// Common organization indicators
    org_indicators: HashSet<String>,

    /// Common technology keywords
    tech_keywords: HashSet<String>,

    /// Common words that should NOT be extracted as entities
    /// (stop words that start with capitals at sentence beginning)
    stop_words: HashSet<String>,
}

impl EntityExtractor {
    pub fn new() -> Self {
        let person_indicators: HashSet<String> =
            vec!["mr", "mrs", "ms", "dr", "prof", "sir", "madam"]
                .into_iter()
                .map(String::from)
                .collect();

        let org_indicators: HashSet<String> = vec![
            "inc",
            "corp",
            "ltd",
            "llc",
            "company",
            "corporation",
            "university",
            "institute",
            "foundation",
        ]
        .into_iter()
        .map(String::from)
        .collect();

        let tech_keywords: HashSet<String> = vec![
            "rust",
            "python",
            "java",
            "javascript",
            "typescript",
            "react",
            "vue",
            "angular",
            "docker",
            "kubernetes",
            "aws",
            "azure",
            "gcp",
            "sql",
            "nosql",
            "mongodb",
            "postgresql",
            "redis",
            "kafka",
            "api",
            "rest",
            "graphql",
        ]
        .into_iter()
        .map(String::from)
        .collect();

        // Common words that are capitalized at sentence start but aren't entities
        let stop_words: HashSet<String> = vec![
            "the",
            "a",
            "an",
            "this",
            "that",
            "these",
            "those",
            "i",
            "we",
            "you",
            "he",
            "she",
            "it",
            "they",
            "is",
            "are",
            "was",
            "were",
            "been",
            "being",
            "have",
            "has",
            "had",
            "do",
            "does",
            "did",
            "will",
            "would",
            "could",
            "should",
            "may",
            "might",
            "if",
            "when",
            "where",
            "what",
            "why",
            "how",
            "user",
            "error",
            "task",
            "code",
            "pattern",
            "search",
            "discovered",
            "completed",
            "performed",
            "mentioned",
        ]
        .into_iter()
        .map(String::from)
        .collect();

        Self {
            person_indicators,
            org_indicators,
            tech_keywords,
            stop_words,
        }
    }

    /// Calculate base salience for an entity based on its type and detection confidence
    fn calculate_base_salience(label: &EntityLabel, is_proper_noun: bool) -> f32 {
        let type_salience = match label {
            EntityLabel::Person => 0.8,       // People are highly salient
            EntityLabel::Organization => 0.7, // Organizations are important
            EntityLabel::Location => 0.6,     // Locations matter for context
            EntityLabel::Technology => 0.6,   // Tech keywords matter for dev context
            EntityLabel::Product => 0.7,      // Products are specific entities
            EntityLabel::Event => 0.6,        // Events are temporal anchors
            EntityLabel::Skill => 0.5,        // Skills are somewhat important
            EntityLabel::Concept => 0.4,      // Concepts are more generic
            EntityLabel::Date => 0.3,         // Dates are structural, not salient
            EntityLabel::Other(_) => 0.3,     // Unknown types get low salience
        };

        // Proper nouns get a 20% boost
        if is_proper_noun {
            (type_salience * 1.2_f32).min(1.0_f32)
        } else {
            type_salience
        }
    }

    /// Check if a word is likely a proper noun (not just capitalized at sentence start)
    fn is_likely_proper_noun(&self, word: &str, position: usize, prev_char: Option<char>) -> bool {
        // If it's not at position 0 and is capitalized, it's likely a proper noun
        if position > 0 {
            return true;
        }

        // At position 0, check if previous character was punctuation (sentence start)
        // If previous char was '.', '!', '?' then this might just be sentence capitalization
        if let Some(c) = prev_char {
            if c == '.' || c == '!' || c == '?' {
                // It's at sentence start - could be either
                // Check if it's a common word
                let lower = word.to_lowercase();
                return !self.stop_words.contains(&lower);
            }
        }

        // Default to proper noun for capitalized words
        true
    }

    /// Extract entities from text with salience information
    pub fn extract_with_salience(&self, text: &str) -> Vec<ExtractedEntity> {
        let mut entities = Vec::new();
        let mut seen = HashSet::new();

        // Split into words and detect capitalized sequences
        let words: Vec<&str> = text.split_whitespace().collect();

        for (i, word) in words.iter().enumerate() {
            let clean_word = word.trim_matches(|c: char| !c.is_alphanumeric());

            if clean_word.is_empty() || clean_word.len() < 2 {
                continue;
            }

            let lower = clean_word.to_lowercase();

            // Skip common stop words
            if self.stop_words.contains(&lower) {
                continue;
            }

            // Check for technology keywords (always proper nouns in tech context)
            if self.tech_keywords.contains(&lower) && !seen.contains(&lower) {
                let entity = ExtractedEntity {
                    name: clean_word.to_string(),
                    label: EntityLabel::Technology,
                    is_proper_noun: true, // Tech keywords are always "proper" in context
                    base_salience: Self::calculate_base_salience(&EntityLabel::Technology, true),
                };
                entities.push(entity);
                seen.insert(lower.clone());
                continue;
            }

            // Check for capitalized words (potential entities)
            if clean_word
                .chars()
                .next()
                .map(|c| c.is_uppercase())
                .unwrap_or(false)
            {
                let mut entity_name = clean_word.to_string();
                let mut entity_label = EntityLabel::Other("Unknown".to_string());

                // Determine previous character for proper noun detection
                let prev_char = if i > 0 {
                    words[i - 1].chars().last()
                } else {
                    None
                };

                let is_proper = self.is_likely_proper_noun(clean_word, i, prev_char);

                // Check for person indicators
                if i > 0
                    && self
                        .person_indicators
                        .contains(&words[i - 1].to_lowercase())
                {
                    entity_label = EntityLabel::Person;
                }

                // Check for multi-word capitalized sequences
                let mut j = i + 1;
                while j < words.len()
                    && words[j]
                        .chars()
                        .next()
                        .map(|c| c.is_uppercase())
                        .unwrap_or(false)
                {
                    let next_word = words[j].trim_matches(|c: char| !c.is_alphanumeric());
                    // Skip stop words in multi-word sequences
                    if !self.stop_words.contains(&next_word.to_lowercase()) {
                        entity_name.push(' ');
                        entity_name.push_str(next_word);
                    }
                    j += 1;
                }

                // Check for organization indicators
                for word in entity_name.split_whitespace() {
                    if self.org_indicators.contains(&word.to_lowercase()) {
                        entity_label = EntityLabel::Organization;
                        break;
                    }
                }

                // Default to Person for single capitalized words without other indicators
                if matches!(entity_label, EntityLabel::Other(_)) && !entity_name.contains(' ') {
                    entity_label = EntityLabel::Person;
                } else if matches!(entity_label, EntityLabel::Other(_)) {
                    entity_label = EntityLabel::Concept;
                }

                let entity_key = entity_name.to_lowercase();
                if !seen.contains(&entity_key) {
                    let base_salience = Self::calculate_base_salience(&entity_label, is_proper);
                    let entity = ExtractedEntity {
                        name: entity_name,
                        label: entity_label,
                        is_proper_noun: is_proper,
                        base_salience,
                    };
                    entities.push(entity);
                    seen.insert(entity_key);
                }
            }
        }

        entities
    }

    /// Legacy method for backward compatibility - returns (name, label) tuples
    pub fn extract(&self, text: &str) -> Vec<(String, EntityLabel)> {
        self.extract_with_salience(text)
            .into_iter()
            .map(|e| (e.name, e.label))
            .collect()
    }
}

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