//! Dynamic Tree Manager — Phase-95 (T95.3–T95.6)
//!
//! Implements access-driven, reversible promotion/demotion of tree nodes.
//!
//! # Design Inspiration
//!
//! Fuses two ideas:
//! - **BriLLM Three-Tier Weight Model**: access-driven promotion (shared → independent)
//!   and demotion (independent → shared) with periodic scanning.
//! - **PageIndex Tree Thinning**: merging small nodes into parents to reduce tree depth.
//!
//! # Lifecycle
//!
//! ```text
//!         record_access()
//!              │
//!     ┌────────▼─────────┐
//!     │   Active Node    │◄───── promote_node()
//!     │ (independent)    │         (Shadow → Active)
//!     └────────┬─────────┘
//!              │
//!     demote_subtree()
//!     (Active → Shadow)
//!              │
//!     ┌────────▼─────────┐
//!     │   Shadow Node    │───── preserved for future restore
//!     │  (merged into    │
//!     │   parent)        │
//!     └──────────────────┘
//! ```
//!
//! # Key Properties
//!
//! - **Reversible**: demote saves `original_content` + `pre_merge_content` on shadow/parent.
//! - **O(1) restore**: promote just flips entity types and relation types.
//! - **Access-driven**: nodes are promoted when parent hit count ≥ threshold.
//! - **Idle-driven**: nodes are demoted when idle for longer than timeout.

use std::sync::Arc;
use std::time::Instant;

use crate::social::SocialStoreV2;
use crate::types::graph::{Entity, Relation};

use super::types::{
    NodeAccessStats, DynamicTreeConfig, ScanReport, DemoteReport, PromoteReport,
    TreeIndexError, TreeIndexResult,
    tree_entity_types, tree_relation_types, tree_property_keys,
    estimate_tokens,
};
use super::store::TreeIndexStore;

/// Dynamic Tree Manager — orchestrates access-driven promotion/demotion.
///
/// Holds a reference to the `TreeIndexStore` and a `DynamicTreeConfig`.
/// Provides methods to record accesses, demote cold subtrees, promote hot
/// nodes, and run periodic scans.
pub struct DynamicTreeManager {
    store: Arc<TreeIndexStore>,
    config: DynamicTreeConfig,
}

impl DynamicTreeManager {
    /// Create a new DynamicTreeManager.
    pub fn new(store: Arc<TreeIndexStore>, config: DynamicTreeConfig) -> Self {
        Self { store, config }
    }

    /// Get a reference to the configuration.
    pub fn config(&self) -> &DynamicTreeConfig {
        &self.config
    }

    /// Get a reference to the underlying store.
    pub fn store(&self) -> &Arc<TreeIndexStore> {
        &self.store
    }

    /// Get the underlying graph store for direct access.
    fn graph(&self) -> &Arc<SocialStoreV2> {
        self.store.graph()
    }

    // ========================================================================
    // T95.3: record_access — Node Access Recording
    // ========================================================================

    /// Record a retrieval hit on a node.
    ///
    /// Increments `hit_count` and updates `last_accessed_at`.
    /// Called by `TreeRetriever` when a node is matched during retrieval.
    ///
    /// # Arguments
    /// * `node_entity_id` - The entity ID of the accessed node.
    ///
    /// # Returns
    /// Updated `NodeAccessStats`, or `None` if the entity was not found.
    pub fn record_access(&self, node_entity_id: &str) -> Option<NodeAccessStats> {
        let entity = self.graph().get_entity_by_id(node_entity_id)?;

        let now = chrono_now();
        let mut stats = Self::extract_access_stats(&entity);
        stats.record_hit(now);

        // Write back the updated entity with new access stats
        let mut updated = entity;
        updated.properties.insert(
            tree_property_keys::ACCESS_STATS.to_string(),
            serde_json::to_value(&stats).unwrap_or_default(),
        );
        self.graph().put_entity(updated);

        Some(stats)
    }

    /// Record access on multiple nodes at once (batch).
    pub fn record_access_batch(&self, node_entity_ids: &[&str]) {
        for id in node_entity_ids {
            self.record_access(id);
        }
    }

    /// Get the access stats for a node.
    pub fn get_access_stats(&self, node_entity_id: &str) -> Option<NodeAccessStats> {
        let entity = self.graph().get_entity_by_id(node_entity_id)?;
        Some(Self::extract_access_stats(&entity))
    }

    /// Extract `NodeAccessStats` from an entity's properties.
    fn extract_access_stats(entity: &Entity) -> NodeAccessStats {
        entity
            .properties
            .get(tree_property_keys::ACCESS_STATS)
            .and_then(|v| serde_json::from_value::<NodeAccessStats>(v.clone()).ok())
            .unwrap_or_default()
    }

    // ========================================================================
    // T95.4: demote_subtree — Cold Subtree Demotion
    // ========================================================================

    /// Demote (merge) all children of a parent node into the parent.
    ///
    /// This is the "demotion" operation, analogous to BriLLM's
    /// "independent → shared" edge weight demotion.
    ///
    /// # Process
    /// 1. Save parent's current content as `pre_merge_content`.
    /// 2. For each active child:
    ///    a. Save child's content as `original_content`.
    ///    b. Change child's entity_type to `tree_section_shadow`.
    ///    c. Change relation from `HAS_CHILD` to `HAS_SHADOW`.
    ///    d. Recursively shadow any grandchildren.
    /// 3. Merge all children's text into parent.
    /// 4. Mark parent with `has_shadows = true`.
    ///
    /// # Guards
    /// - Nodes at level ≤ `preserve_top_levels` are protected.
    /// - Nodes that already have shadows are skipped (already demoted).
    ///
    /// # Returns
    /// `DemoteReport` with the number of nodes shadowed.
    pub fn demote_subtree(&self, parent_entity_id: &str) -> TreeIndexResult<DemoteReport> {
        let parent = self.graph().get_entity_by_id(parent_entity_id)
            .ok_or_else(|| TreeIndexError::NodeNotFound(parent_entity_id.to_string()))?;

        // Guard: must be a tree_section
        if parent.entity_type != tree_entity_types::TREE_SECTION {
            return Err(TreeIndexError::InvalidStructure(format!(
                "Cannot demote entity of type '{}', expected '{}'",
                parent.entity_type, tree_entity_types::TREE_SECTION
            )));
        }

        // Guard: protected level
        let level = parent.properties.get("level")
            .and_then(|v| v.as_u64()).unwrap_or(0) as u8;
        if level <= self.config.preserve_top_levels {
            return Err(TreeIndexError::ConfigError(format!(
                "Cannot demote node at level {} (preserve_top_levels = {})",
                level, self.config.preserve_top_levels
            )));
        }

        // Guard: already has shadows
        let already_has_shadows = parent.properties.get(tree_property_keys::HAS_SHADOWS)
            .and_then(|v| v.as_bool()).unwrap_or(false);
        if already_has_shadows {
            return Err(TreeIndexError::InvalidStructure(
                "Node already has shadow children; promote first before re-demoting".to_string()
            ));
        }

        // Find all active child relations
        let all_relations = self.graph().get_entity_relations(parent_entity_id);
        let child_relations: Vec<Relation> = all_relations.into_iter()
            .filter(|r| r.source == parent_entity_id && r.relation_type == tree_relation_types::HAS_CHILD)
            .collect();

        if child_relations.is_empty() {
            return Ok(DemoteReport { shadow_count: 0 }); // Nothing to demote
        }

        // 1. Save parent's original content
        let parent_content = parent.properties.get("content")
            .and_then(|v| v.as_str()).unwrap_or("").to_string();
        let parent_token_count = parent.properties.get("token_count")
            .and_then(|v| v.as_u64()).unwrap_or(0);

        let mut updated_parent = parent;
        updated_parent.properties.insert(
            tree_property_keys::PRE_MERGE_CONTENT.to_string(),
            serde_json::json!(parent_content),
        );
        updated_parent.properties.insert(
            tree_property_keys::PRE_MERGE_TOKEN_COUNT.to_string(),
            serde_json::json!(parent_token_count),
        );

        // 2. Collect text from children and shadow them
        let mut merged_text_parts = vec![parent_content.clone()];
        let mut total_shadowed = 0;

        for rel in &child_relations {
            // Collect text recursively from this child subtree
            self.collect_text_from_subtree(&rel.target, &mut merged_text_parts);

            // Shadow the child subtree recursively
            total_shadowed += self.shadow_subtree_recursive(&rel.target)?;

            // Change relation type: HAS_CHILD → HAS_SHADOW
            let shadow_rel = Relation::new(
                &rel.source, &rel.target, tree_relation_types::HAS_SHADOW,
            ).with_metadata("order",
                rel.metadata.get("order").cloned().unwrap_or(serde_json::json!(0))
            );

            self.graph().delete_relation(&rel.id);
            self.graph().put_relation(shadow_rel);
        }

        // 3. Merge text into parent
        let merged_text = merged_text_parts
            .into_iter()
            .filter(|s| !s.is_empty())
            .collect::<Vec<_>>()
            .join("\n\n");

        updated_parent.properties.insert(
            "content".to_string(), serde_json::json!(merged_text),
        );
        updated_parent.properties.insert(
            "token_count".to_string(), serde_json::json!(estimate_tokens(&merged_text)),
        );

        // 4. Mark parent as having shadows
        updated_parent.properties.insert(
            tree_property_keys::HAS_SHADOWS.to_string(), serde_json::json!(true),
        );
        updated_parent.properties.insert(
            tree_property_keys::DEMOTED_AT.to_string(), serde_json::json!(chrono_now()),
        );

        self.graph().put_entity(updated_parent);

        Ok(DemoteReport { shadow_count: total_shadowed })
    }

    /// Recursively convert an active subtree to shadow.
    ///
    /// Returns the number of nodes shadowed (including the root).
    fn shadow_subtree_recursive(&self, entity_id: &str) -> TreeIndexResult<usize> {
        let entity = self.graph().get_entity_by_id(entity_id)
            .ok_or_else(|| TreeIndexError::NodeNotFound(entity_id.to_string()))?;

        // Only shadow active sections
        if entity.entity_type != tree_entity_types::TREE_SECTION {
            return Ok(0);
        }

        let mut count = 1;

        // Save original content before shadowing
        let content = entity.properties.get("content")
            .and_then(|v| v.as_str()).unwrap_or("").to_string();
        let token_count = entity.properties.get("token_count")
            .and_then(|v| v.as_u64()).unwrap_or(0);

        let mut shadowed = entity;
        shadowed.entity_type = tree_entity_types::TREE_SECTION_SHADOW.to_string();
        shadowed.properties.insert(
            tree_property_keys::ORIGINAL_CONTENT.to_string(),
            serde_json::json!(content),
        );
        shadowed.properties.insert(
            tree_property_keys::ORIGINAL_TOKEN_COUNT.to_string(),
            serde_json::json!(token_count),
        );
        shadowed.properties.insert(
            tree_property_keys::DEMOTED_AT.to_string(),
            serde_json::json!(chrono_now()),
        );

        self.graph().put_entity(shadowed);

        // Recursively shadow children
        let all_relations = self.graph().get_entity_relations(entity_id);
        let child_relations: Vec<Relation> = all_relations.into_iter()
            .filter(|r| r.source == entity_id && r.relation_type == tree_relation_types::HAS_CHILD)
            .collect();

        for rel in &child_relations {
            count += self.shadow_subtree_recursive(&rel.target)?;

            // Change relation: HAS_CHILD → HAS_SHADOW
            let shadow_rel = Relation::new(
                &rel.source, &rel.target, tree_relation_types::HAS_SHADOW,
            ).with_metadata("order",
                rel.metadata.get("order").cloned().unwrap_or(serde_json::json!(0))
            );
            self.graph().delete_relation(&rel.id);
            self.graph().put_relation(shadow_rel);
        }

        Ok(count)
    }

    /// Collect text from a subtree (for merging into parent).
    fn collect_text_from_subtree(&self, entity_id: &str, parts: &mut Vec<String>) {
        if let Some(entity) = self.graph().get_entity_by_id(entity_id) {
            if let Some(text) = entity.properties.get("content").and_then(|v| v.as_str()) {
                if !text.is_empty() {
                    parts.push(text.to_string());
                }
            }

            // Recurse into children
            let all_relations = self.graph().get_entity_relations(entity_id);
            let mut children: Vec<(usize, String)> = all_relations.into_iter()
                .filter(|r| r.source == entity_id && r.relation_type == tree_relation_types::HAS_CHILD)
                .map(|r| {
                    let order = r.metadata.get("order")
                        .and_then(|v| v.as_u64()).unwrap_or(0) as usize;
                    (order, r.target.clone())
                })
                .collect();
            children.sort_by_key(|(order, _)| *order);

            for (_, child_id) in children {
                self.collect_text_from_subtree(&child_id, parts);
            }
        }
    }

    // ========================================================================
    // T95.5: promote_node — Hot Node Promotion
    // ========================================================================

    /// Promote (expand) shadow children of a parent node back to active.
    ///
    /// This is the "promotion" operation, analogous to BriLLM's
    /// "shared → independent" edge weight promotion.
    ///
    /// # Process
    /// 1. Restore parent's content from `pre_merge_content`.
    /// 2. For each shadow child:
    ///    a. Change entity_type from `tree_section_shadow` to `tree_section`.
    ///    b. Restore content from `original_content`.
    ///    c. Change relation from `HAS_SHADOW` to `HAS_CHILD`.
    ///    d. Recursively restore any shadow grandchildren.
    /// 3. Clear parent's `has_shadows`, `pre_merge_content`.
    /// 4. Reset parent's hit count to avoid immediate re-promotion.
    ///
    /// # Returns
    /// `PromoteReport` with the number of nodes restored.
    pub fn promote_node(&self, parent_entity_id: &str) -> TreeIndexResult<PromoteReport> {
        let parent = self.graph().get_entity_by_id(parent_entity_id)
            .ok_or_else(|| TreeIndexError::NodeNotFound(parent_entity_id.to_string()))?;

        // Guard: must have shadows
        let has_shadows = parent.properties.get(tree_property_keys::HAS_SHADOWS)
            .and_then(|v| v.as_bool()).unwrap_or(false);
        if !has_shadows {
            return Err(TreeIndexError::InvalidStructure(
                "Node has no shadow children to promote".to_string()
            ));
        }

        // 1. Restore parent's original content
        let mut updated_parent = parent;

        if let Some(pre_merge) = updated_parent.properties.get(tree_property_keys::PRE_MERGE_CONTENT).cloned() {
            updated_parent.properties.insert("content".to_string(), pre_merge);
        }
        if let Some(pre_merge_tc) = updated_parent.properties.get(tree_property_keys::PRE_MERGE_TOKEN_COUNT).cloned() {
            updated_parent.properties.insert("token_count".to_string(), pre_merge_tc);
        }

        // Clean up merge metadata
        updated_parent.properties.remove(tree_property_keys::PRE_MERGE_CONTENT);
        updated_parent.properties.remove(tree_property_keys::PRE_MERGE_TOKEN_COUNT);
        updated_parent.properties.remove(tree_property_keys::HAS_SHADOWS);
        updated_parent.properties.remove(tree_property_keys::DEMOTED_AT);

        // Reset hit count to avoid immediate re-promotion cycle
        let mut stats = Self::extract_access_stats(&updated_parent);
        stats.hit_count = 0;
        updated_parent.properties.insert(
            tree_property_keys::ACCESS_STATS.to_string(),
            serde_json::to_value(&stats).unwrap_or_default(),
        );

        self.graph().put_entity(updated_parent);

        // 2. Restore all shadow children
        let all_relations = self.graph().get_entity_relations(parent_entity_id);
        let shadow_relations: Vec<Relation> = all_relations.into_iter()
            .filter(|r| r.source == parent_entity_id && r.relation_type == tree_relation_types::HAS_SHADOW)
            .collect();

        let mut total_restored = 0;
        for rel in &shadow_relations {
            total_restored += self.restore_subtree_recursive(&rel.target)?;

            // Change relation: HAS_SHADOW → HAS_CHILD
            let active_rel = Relation::new(
                &rel.source, &rel.target, tree_relation_types::HAS_CHILD,
            ).with_metadata("order",
                rel.metadata.get("order").cloned().unwrap_or(serde_json::json!(0))
            );
            self.graph().delete_relation(&rel.id);
            self.graph().put_relation(active_rel);
        }

        Ok(PromoteReport { restored_count: total_restored })
    }

    /// Recursively restore a shadow subtree to active.
    ///
    /// Returns the number of nodes restored (including the root).
    fn restore_subtree_recursive(&self, entity_id: &str) -> TreeIndexResult<usize> {
        let entity = self.graph().get_entity_by_id(entity_id)
            .ok_or_else(|| TreeIndexError::NodeNotFound(entity_id.to_string()))?;

        // Only restore shadows
        if entity.entity_type != tree_entity_types::TREE_SECTION_SHADOW {
            return Ok(0);
        }

        let mut count = 1;

        // Restore original content
        let mut restored = entity;
        restored.entity_type = tree_entity_types::TREE_SECTION.to_string();

        if let Some(original) = restored.properties.get(tree_property_keys::ORIGINAL_CONTENT).cloned() {
            restored.properties.insert("content".to_string(), original);
        }
        if let Some(original_tc) = restored.properties.get(tree_property_keys::ORIGINAL_TOKEN_COUNT).cloned() {
            restored.properties.insert("token_count".to_string(), original_tc);
        }

        // Clean up shadow metadata
        restored.properties.remove(tree_property_keys::ORIGINAL_CONTENT);
        restored.properties.remove(tree_property_keys::ORIGINAL_TOKEN_COUNT);
        restored.properties.remove(tree_property_keys::DEMOTED_AT);

        // Set last_accessed_at to now so restored nodes don't get immediately re-demoted
        let mut stats = Self::extract_access_stats(&restored);
        stats.last_accessed_at = chrono_now();
        if stats.first_accessed_at == 0 {
            stats.first_accessed_at = stats.last_accessed_at;
        }
        restored.properties.insert(
            tree_property_keys::ACCESS_STATS.to_string(),
            serde_json::to_value(&stats).unwrap_or_default(),
        );

        self.graph().put_entity(restored);

        // Recursively restore shadow children
        let all_relations = self.graph().get_entity_relations(entity_id);
        let shadow_relations: Vec<Relation> = all_relations.into_iter()
            .filter(|r| r.source == entity_id && r.relation_type == tree_relation_types::HAS_SHADOW)
            .collect();

        for rel in &shadow_relations {
            count += self.restore_subtree_recursive(&rel.target)?;

            // Change relation: HAS_SHADOW → HAS_CHILD
            let active_rel = Relation::new(
                &rel.source, &rel.target, tree_relation_types::HAS_CHILD,
            ).with_metadata("order",
                rel.metadata.get("order").cloned().unwrap_or(serde_json::json!(0))
            );
            self.graph().delete_relation(&rel.id);
            self.graph().put_relation(active_rel);
        }

        Ok(count)
    }

    // ========================================================================
    // T95.6: DynamicTreeManager — Scan Engine
    // ========================================================================

    /// Run a full promotion/demotion scan across all documents.
    ///
    /// This is the periodic scan, analogous to BriLLM's
    /// `run_promotion_scan()` + `run_demotion_scan()`.
    ///
    /// # Process
    /// 1. For each document:
    ///    a. **Promotion scan**: find parents with shadows whose hit count ≥ threshold → promote.
    ///    b. **Demotion scan**: find parents whose children are all idle → demote.
    ///
    /// # Returns
    /// `ScanReport` summarizing all actions taken.
    pub fn run_scan(&self) -> ScanReport {
        let start = Instant::now();
        let mut report = ScanReport::default();

        let documents = self.store.list_documents();
        report.scanned_docs = documents.len();

        for (doc_id, _metadata) in &documents {
            // --- Phase 1: Promotion scan ---
            let promotion_candidates = self.find_promotion_candidates(doc_id);
            for candidate_id in promotion_candidates {
                if let Ok(r) = self.promote_node(&candidate_id) {
                    report.promoted += r.restored_count;
                }
            }

            // --- Phase 2: Demotion scan ---
            let demotion_candidates = self.find_demotion_candidates(doc_id);
            for candidate_id in demotion_candidates {
                if let Ok(r) = self.demote_subtree(&candidate_id) {
                    report.demoted += r.shadow_count;
                }
            }
        }

        report.scan_duration_ms = start.elapsed().as_millis() as u64;
        report
    }

    /// Find parent nodes that qualify for promotion.
    ///
    /// Criteria:
    /// - Entity type = `tree_section`
    /// - `has_shadows` = true
    /// - `hit_count` ≥ `promote_hit_threshold`
    fn find_promotion_candidates(&self, doc_id: &str) -> Vec<String> {
        let mut candidates = Vec::new();

        // Get all section nodes for this document
        let sections = self.get_doc_sections(doc_id);

        for entity in sections {
            // Must have shadows
            let has_shadows = entity.properties.get(tree_property_keys::HAS_SHADOWS)
                .and_then(|v| v.as_bool()).unwrap_or(false);
            if !has_shadows {
                continue;
            }

            // Check hit count
            let stats = Self::extract_access_stats(&entity);
            if stats.hit_count >= self.config.promote_hit_threshold {
                candidates.push(entity.id.clone());
            }
        }

        candidates
    }

    /// Find parent nodes whose children qualify for demotion.
    ///
    /// Criteria:
    /// - Entity type = `tree_section` with active children
    /// - Node level > `preserve_top_levels`
    /// - All children idle > `demote_idle_timeout_secs`
    /// - Subtree total tokens < `demote_min_tokens`
    /// - No existing shadows (already demoted)
    fn find_demotion_candidates(&self, doc_id: &str) -> Vec<String> {
        let mut candidates = Vec::new();
        let now = chrono_now();

        let sections = self.get_doc_sections(doc_id);

        for entity in sections {
            // Must be active section (not shadow, not doc)
            if entity.entity_type != tree_entity_types::TREE_SECTION {
                continue;
            }

            // Guard: protected level
            let level = entity.properties.get("level")
                .and_then(|v| v.as_u64()).unwrap_or(0) as u8;
            if level <= self.config.preserve_top_levels {
                continue;
            }

            // Skip if already has shadows
            let has_shadows = entity.properties.get(tree_property_keys::HAS_SHADOWS)
                .and_then(|v| v.as_bool()).unwrap_or(false);
            if has_shadows {
                continue;
            }

            // Must have children
            let all_relations = self.graph().get_entity_relations(&entity.id);
            let child_relations: Vec<&Relation> = all_relations.iter()
                .filter(|r| r.source == entity.id && r.relation_type == tree_relation_types::HAS_CHILD)
                .collect();
            if child_relations.is_empty() {
                continue; // Leaf nodes can't demote children
            }

            // Check: all children idle?
            let all_children_idle = child_relations.iter().all(|rel| {
                if let Some(child_entity) = self.graph().get_entity_by_id(&rel.target) {
                    let child_stats = Self::extract_access_stats(&child_entity);
                    child_stats.is_idle(now, self.config.demote_idle_timeout_secs)
                } else {
                    true // Missing entity → treat as idle
                }
            });

            if !all_children_idle {
                continue;
            }

            // Check: subtree token count below threshold?
            let subtree_tokens = self.compute_subtree_tokens(&entity.id);
            if subtree_tokens >= self.config.demote_min_tokens {
                continue;
            }

            candidates.push(entity.id.clone());
        }

        candidates
    }

    // ========================================================================
    // Query Helpers
    // ========================================================================

    /// List all shadow nodes for a given document.
    pub fn list_shadow_nodes(&self, doc_id: &str) -> Vec<Entity> {
        self.get_doc_sections(doc_id)
            .into_iter()
            .filter(|e| e.entity_type == tree_entity_types::TREE_SECTION_SHADOW)
            .collect()
    }

    /// Count shadow nodes for a given document.
    pub fn shadow_node_count(&self, doc_id: &str) -> usize {
        self.list_shadow_nodes(doc_id).len()
    }

    /// Check if the dynamic tree feature is enabled.
    pub fn is_enabled(&self) -> bool {
        self.config.enabled
    }

    // ========================================================================
    // Internal Helpers
    // ========================================================================

    /// Get all section entities belonging to a document.
    ///
    /// Traverses the tree starting from root nodes via HAS_ROOT relations.
    fn get_doc_sections(&self, doc_id: &str) -> Vec<Entity> {
        let mut sections = Vec::new();

        // Get root relations
        let all_relations = self.graph().get_entity_relations(doc_id);
        let root_relations: Vec<&Relation> = all_relations.iter()
            .filter(|r| r.source == doc_id && r.relation_type == tree_relation_types::HAS_ROOT)
            .collect();

        for rel in root_relations {
            self.collect_all_sections(&rel.target, &mut sections);
        }

        sections
    }

    /// Recursively collect all section entities (active and shadow) from a subtree.
    fn collect_all_sections(&self, entity_id: &str, sections: &mut Vec<Entity>) {
        if let Some(entity) = self.graph().get_entity_by_id(entity_id) {
            let is_section = entity.entity_type == tree_entity_types::TREE_SECTION
                || entity.entity_type == tree_entity_types::TREE_SECTION_SHADOW;

            if is_section {
                sections.push(entity);

                // Follow both active and shadow child relations
                let all_relations = self.graph().get_entity_relations(entity_id);
                for rel in all_relations.iter() {
                    if rel.source == entity_id
                        && (rel.relation_type == tree_relation_types::HAS_CHILD
                            || rel.relation_type == tree_relation_types::HAS_SHADOW)
                    {
                        self.collect_all_sections(&rel.target, sections);
                    }
                }
            }
        }
    }

    /// Compute total token count for a subtree rooted at the given entity.
    fn compute_subtree_tokens(&self, entity_id: &str) -> usize {
        let entity = match self.graph().get_entity_by_id(entity_id) {
            Some(e) => e,
            None => return 0,
        };

        let own_tokens = entity.properties.get("token_count")
            .and_then(|v| v.as_u64()).unwrap_or(0) as usize;

        let all_relations = self.graph().get_entity_relations(entity_id);
        let child_tokens: usize = all_relations.iter()
            .filter(|r| r.source == entity_id && r.relation_type == tree_relation_types::HAS_CHILD)
            .map(|r| self.compute_subtree_tokens(&r.target))
            .sum();

        own_tokens + child_tokens
    }
}

/// Get current Unix timestamp in milliseconds.
fn chrono_now() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64
}

// ─── Tests ───────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::social::SocialStoreV2Config;
    use crate::tree_index::types::{TreeDocument, TreeMetadata, TreeNode};

    fn create_test_env() -> (DynamicTreeManager, TreeIndexStore, tempfile::TempDir) {
        let tmp = tempfile::tempdir().unwrap();
        let config = SocialStoreV2Config::new(tmp.path());
        let graph = Arc::new(SocialStoreV2::new(config));
        let store = TreeIndexStore::new(graph);
        let store_arc = Arc::new(store);

        let dynamic_config = DynamicTreeConfig {
            enabled: true,
            promote_hit_threshold: 5, // Low threshold for testing
            demote_idle_timeout_secs: 1, // 1 second for fast testing
            demote_min_tokens: 10000, // High so small nodes always qualify
            preserve_top_levels: 1,   // Only protect H1
            scan_interval_secs: 60,
            max_shadow_nodes_per_doc: 100,
        };

        let manager = DynamicTreeManager::new(store_arc.clone(), dynamic_config);

        // Need a plain store reference for import operations
        let store_ref = TreeIndexStore::new(manager.graph().clone());

        (manager, store_ref, tmp)
    }

    fn create_sample_doc() -> TreeDocument {
        let meta = TreeMetadata::new("test-dynamic")
            .with_description("Test doc for dynamic tree");

        // Build: H1 > H2 > H3
        let mut root = TreeNode::new("0001", "Chapter 1", 1)
            .with_content("Chapter 1 introduction text.");

        let mut section = TreeNode::new("0002", "Section 1.1", 2)
            .with_content("Section 1.1 content.");

        let subsection_a = TreeNode::new("0003", "Sub A", 3)
            .with_content("Subsection A details.");
        let subsection_b = TreeNode::new("0004", "Sub B", 3)
            .with_content("Subsection B details.");

        section.add_child(subsection_a);
        section.add_child(subsection_b);

        let section2 = TreeNode::new("0005", "Section 1.2", 2)
            .with_content("Section 1.2 content.");

        root.add_child(section);
        root.add_child(section2);

        TreeDocument::new(meta, vec![root])
    }

    // ─── T95.3: record_access tests ─────────────────────────────────

    #[test]
    fn test_record_access_basic() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        // Find the root section entity
        let _tree = store.get_tree(&doc_id).unwrap();
        let root_relations = manager.graph().get_entity_relations(&doc_id);
        let root_entity_id = root_relations.iter()
            .find(|r| r.relation_type == tree_relation_types::HAS_ROOT)
            .map(|r| r.target.clone())
            .unwrap();

        // Record access
        let stats = manager.record_access(&root_entity_id).unwrap();
        assert_eq!(stats.hit_count, 1);
        assert!(stats.last_accessed_at > 0);
        assert!(stats.first_accessed_at > 0);

        // Record again
        let stats2 = manager.record_access(&root_entity_id).unwrap();
        assert_eq!(stats2.hit_count, 2);
        assert!(stats2.last_accessed_at >= stats.last_accessed_at);
        assert_eq!(stats2.first_accessed_at, stats.first_accessed_at);
    }

    #[test]
    fn test_record_access_nonexistent() {
        let (manager, _store, _tmp) = create_test_env();
        let result = manager.record_access("nonexistent_id");
        assert!(result.is_none());
    }

    #[test]
    fn test_get_access_stats() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let root_relations = manager.graph().get_entity_relations(&doc_id);
        let root_entity_id = root_relations.iter()
            .find(|r| r.relation_type == tree_relation_types::HAS_ROOT)
            .map(|r| r.target.clone())
            .unwrap();

        // Initially zero
        let stats = manager.get_access_stats(&root_entity_id).unwrap();
        assert_eq!(stats.hit_count, 0);

        // After recording
        manager.record_access(&root_entity_id);
        let stats = manager.get_access_stats(&root_entity_id).unwrap();
        assert_eq!(stats.hit_count, 1);
    }

    // ─── T95.4: demote_subtree tests ────────────────────────────────

    #[test]
    fn test_demote_subtree_basic() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        // Find "Section 1.1" entity (level 2, has children Sub A and Sub B)
        let section_entity_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");

        // Demote: merge Sub A + Sub B into Section 1.1
        let report = manager.demote_subtree(&section_entity_id).unwrap();
        assert_eq!(report.shadow_count, 2); // Sub A + Sub B

        // Verify parent now has merged content
        let parent = manager.graph().get_entity_by_id(&section_entity_id).unwrap();
        let content = parent.properties.get("content").and_then(|v| v.as_str()).unwrap();
        assert!(content.contains("Section 1.1 content"));
        assert!(content.contains("Subsection A details"));
        assert!(content.contains("Subsection B details"));

        // Verify has_shadows flag
        let has_shadows = parent.properties.get(tree_property_keys::HAS_SHADOWS)
            .and_then(|v| v.as_bool()).unwrap();
        assert!(has_shadows);

        // Verify pre_merge_content saved
        let pre_merge = parent.properties.get(tree_property_keys::PRE_MERGE_CONTENT)
            .and_then(|v| v.as_str()).unwrap();
        assert_eq!(pre_merge, "Section 1.1 content.");

        // Verify children are now shadows
        let shadow_count = manager.shadow_node_count(&doc_id);
        assert_eq!(shadow_count, 2);
    }

    #[test]
    fn test_demote_protected_level() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        // Try to demote root (level 1) — should fail (preserve_top_levels = 1)
        let root_entity_id = find_section_by_title(manager.graph(), &doc_id, "Chapter 1");
        let result = manager.demote_subtree(&root_entity_id);
        assert!(result.is_err());
    }

    #[test]
    fn test_demote_no_children() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        // "Section 1.2" is a leaf with level 2 — should succeed with 0 shadows
        let leaf_entity_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.2");
        let report = manager.demote_subtree(&leaf_entity_id).unwrap();
        assert_eq!(report.shadow_count, 0);
    }

    #[test]
    fn test_demote_double_demote_fails() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");
        manager.demote_subtree(&section_id).unwrap();

        // Second demote should fail (already has shadows)
        let result = manager.demote_subtree(&section_id);
        assert!(result.is_err());
    }

    // ─── T95.5: promote_node tests ──────────────────────────────────

    #[test]
    fn test_promote_restores_structure() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");

        // Get original content before demote
        let original_parent_content = manager.graph().get_entity_by_id(&section_id).unwrap()
            .properties.get("content").and_then(|v| v.as_str()).unwrap().to_string();

        // Demote
        manager.demote_subtree(&section_id).unwrap();
        assert_eq!(manager.shadow_node_count(&doc_id), 2);

        // Promote (restore)
        let report = manager.promote_node(&section_id).unwrap();
        assert_eq!(report.restored_count, 2);
        assert_eq!(manager.shadow_node_count(&doc_id), 0);

        // Verify parent content restored
        let parent = manager.graph().get_entity_by_id(&section_id).unwrap();
        let restored_content = parent.properties.get("content").and_then(|v| v.as_str()).unwrap();
        assert_eq!(restored_content, original_parent_content);

        // Verify has_shadows removed
        assert!(parent.properties.get(tree_property_keys::HAS_SHADOWS).is_none());
        assert!(parent.properties.get(tree_property_keys::PRE_MERGE_CONTENT).is_none());

        // Verify children are active again
        let all_rels = manager.graph().get_entity_relations(&section_id);
        let active_children: Vec<_> = all_rels.iter()
            .filter(|r| r.source == section_id && r.relation_type == tree_relation_types::HAS_CHILD)
            .collect();
        assert_eq!(active_children.len(), 2);

        // Verify no shadow relations remain
        let shadow_rels: Vec<_> = all_rels.iter()
            .filter(|r| r.source == section_id && r.relation_type == tree_relation_types::HAS_SHADOW)
            .collect();
        assert_eq!(shadow_rels.len(), 0);
    }

    #[test]
    fn test_promote_without_shadows_fails() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");
        let result = manager.promote_node(&section_id);
        assert!(result.is_err()); // No shadows to promote
    }

    #[test]
    fn test_demote_promote_roundtrip_preserves_content() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");

        // Capture original tree structure
        let original_tree = store.get_tree(&doc_id).unwrap();
        let original_section = original_tree.find_node("0002").unwrap();
        let original_child_a_content = original_section.children[0].content.text.clone();
        let original_child_b_content = original_section.children[1].content.text.clone();

        // Demote → Promote roundtrip
        manager.demote_subtree(&section_id).unwrap();
        manager.promote_node(&section_id).unwrap();

        // Verify child content is exactly restored
        let restored_tree = store.get_tree(&doc_id).unwrap();
        let restored_section = restored_tree.find_node("0002").unwrap();
        assert_eq!(restored_section.children.len(), 2);
        assert_eq!(restored_section.children[0].content.text, original_child_a_content);
        assert_eq!(restored_section.children[1].content.text, original_child_b_content);
    }

    #[test]
    fn test_promote_resets_hit_count() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");

        // Record hits
        for _ in 0..10 {
            manager.record_access(&section_id);
        }
        let stats = manager.get_access_stats(&section_id).unwrap();
        assert_eq!(stats.hit_count, 10);

        // Demote and promote
        manager.demote_subtree(&section_id).unwrap();
        manager.promote_node(&section_id).unwrap();

        // Hit count should be reset
        let stats = manager.get_access_stats(&section_id).unwrap();
        assert_eq!(stats.hit_count, 0);
    }

    // ─── T95.6: run_scan tests ──────────────────────────────────────

    #[test]
    fn test_scan_promotes_hot_nodes() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let doc_id = store.import_tree(&doc).unwrap();

        let section_id = find_section_by_title(manager.graph(), &doc_id, "Section 1.1");

        // Demote first
        manager.demote_subtree(&section_id).unwrap();
        assert_eq!(manager.shadow_node_count(&doc_id), 2);

        // Simulate hot access (threshold = 5)
        for _ in 0..5 {
            manager.record_access(&section_id);
        }

        // Run scan — should promote
        let report = manager.run_scan();
        assert!(report.promoted > 0);
        assert_eq!(manager.shadow_node_count(&doc_id), 0);
    }

    #[test]
    fn test_scan_demotes_cold_nodes() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        let _doc_id = store.import_tree(&doc).unwrap();

        // Wait for idle timeout (1 second in test config)
        std::thread::sleep(std::time::Duration::from_millis(1100));

        // Run scan — should demote idle children
        let report = manager.run_scan();
        // "Section 1.1" has children (Sub A, Sub B) that are idle
        // and level 2 > preserve_top_levels(1), so should be demoted
        assert!(report.demoted > 0);
    }

    #[test]
    fn test_scan_report_fields() {
        let (manager, store, _tmp) = create_test_env();
        let doc = create_sample_doc();
        store.import_tree(&doc).unwrap();

        let report = manager.run_scan();
        assert_eq!(report.scanned_docs, 1);
        assert!(report.scan_duration_ms < 10000); // Should be fast
    }

    // ─── Helper ─────────────────────────────────────────────────────

    /// Find a section entity by title within a document.
    fn find_section_by_title(graph: &Arc<SocialStoreV2>, doc_id: &str, title: &str) -> String {
        let all_rels = graph.get_entity_relations(doc_id);
        for rel in &all_rels {
            if rel.source == doc_id && rel.relation_type == tree_relation_types::HAS_ROOT {
                if let Some(found) = find_entity_recursive(graph, &rel.target, title) {
                    return found;
                }
            }
        }
        panic!("Section '{}' not found in doc '{}'", title, doc_id);
    }

    fn find_entity_recursive(graph: &Arc<SocialStoreV2>, entity_id: &str, title: &str) -> Option<String> {
        let entity = graph.get_entity_by_id(entity_id)?;
        if entity.name == title {
            return Some(entity_id.to_string());
        }

        let all_rels = graph.get_entity_relations(entity_id);
        for rel in &all_rels {
            if rel.source == entity_id
                && (rel.relation_type == tree_relation_types::HAS_CHILD
                    || rel.relation_type == tree_relation_types::HAS_SHADOW)
            {
                if let Some(found) = find_entity_recursive(graph, &rel.target, title) {
                    return Some(found);
                }
            }
        }
        None
    }
}
