use super::Binding;
use crate::ir::NodeId;
use indexmap::{IndexMap, IndexSet};
use std::collections::BTreeMap;

/// Tracks which nodes depend on which state paths
pub struct DependencyGraph {
    /// Maps state path → set of NodeIds that depend on it
    dependencies: IndexMap<String, IndexSet<NodeId>>,

    /// Maps NodeId → set of state paths it depends on
    node_bindings: IndexMap<NodeId, IndexSet<String>>,

    /// Prefix index for efficient path lookups
    /// Maps path prefix → set of full paths that start with this prefix
    /// This enables O(log n) lookup instead of O(n) scan
    prefix_index: BTreeMap<String, IndexSet<String>>,
}

impl DependencyGraph {
    pub fn new() -> Self {
        Self {
            dependencies: IndexMap::new(),
            node_bindings: IndexMap::new(),
            prefix_index: BTreeMap::new(),
        }
    }

    /// Record that a node depends on a binding
    pub fn add_dependency(&mut self, node_id: NodeId, binding: &Binding) {
        let path = binding.full_path();

        // Add to dependencies map
        self.dependencies
            .entry(path.clone())
            .or_insert_with(IndexSet::new)
            .insert(node_id);

        // Add to node_bindings map
        self.node_bindings
            .entry(node_id)
            .or_insert_with(IndexSet::new)
            .insert(path.clone());

        // Update prefix index for efficient lookups
        // Add all prefixes of this path to the index
        self.add_path_to_prefix_index(&path);
    }

    /// Add a path and all its prefixes to the prefix index
    fn add_path_to_prefix_index(&mut self, path: &str) {
        // Add the full path to its own prefix
        self.prefix_index
            .entry(path.to_string())
            .or_insert_with(IndexSet::new)
            .insert(path.to_string());

        // Add all parent prefixes
        let mut current = path;
        while let Some(dot_pos) = current.rfind('.') {
            current = &current[..dot_pos];
            self.prefix_index
                .entry(current.to_string())
                .or_insert_with(IndexSet::new)
                .insert(path.to_string());
        }
    }

    /// Remove all dependencies for a node (when unmounting)
    pub fn remove_node(&mut self, node_id: NodeId) {
        if let Some(paths) = self.node_bindings.shift_remove(&node_id) {
            for path in paths {
                if let Some(nodes) = self.dependencies.get_mut(&path) {
                    nodes.shift_remove(&node_id);
                }
            }
        }
    }

    /// Get all nodes that depend on a specific state path
    pub fn get_dependent_nodes(&self, path: &str) -> Option<&IndexSet<NodeId>> {
        self.dependencies.get(path)
    }

    /// Get all nodes affected by a state patch (considers nested paths)
    /// E.g., if "user" changes, both "user" and "user.name" subscribers are affected
    /// Uses prefix index for O(log n) lookup instead of O(n) scan
    pub fn get_affected_nodes(&self, changed_path: &str) -> IndexSet<NodeId> {
        let mut affected = IndexSet::new();
        let mut affected_paths = IndexSet::new();

        // Strategy: Find all paths that could be affected by this change
        // 1. Exact match: changed_path itself
        if self.dependencies.contains_key(changed_path) {
            affected_paths.insert(changed_path.to_string());
        }

        // 2. Parent paths: all prefixes of changed_path (e.g., "user.name" affects "user")
        let mut current = changed_path;
        while let Some(dot_pos) = current.rfind('.') {
            current = &current[..dot_pos];
            if self.dependencies.contains_key(current) {
                affected_paths.insert(current.to_string());
            }
        }

        // 3. Child paths: all paths that start with changed_path (e.g., "user" affects "user.name")
        // Use prefix index for efficient lookup
        if let Some(child_paths) = self.prefix_index.get(changed_path) {
            affected_paths.extend(child_paths.iter().cloned());
        }

        // Collect nodes from all affected paths
        for path in affected_paths {
            if let Some(nodes) = self.dependencies.get(&path) {
                affected.extend(nodes.iter().copied());
            }
        }

        affected
    }

    /// Clear all dependencies
    pub fn clear(&mut self) {
        self.dependencies.clear();
        self.node_bindings.clear();
        self.prefix_index.clear();
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_prefix_index_basic() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();

        // Create real NodeIds using InstanceTree
        let mut tree = InstanceTree::new();
        let node1 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node2 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node3 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));

        // Add dependencies
        let binding1 = Binding::state(vec!["user".to_string()]);
        let binding2 = Binding::state(vec!["user".to_string(), "name".to_string()]);
        let binding3 = Binding::state(vec!["user".to_string(), "email".to_string()]);

        graph.add_dependency(node1, &binding1);
        graph.add_dependency(node2, &binding2);
        graph.add_dependency(node3, &binding3);

        // Verify prefix index was built correctly
        assert!(graph.prefix_index.contains_key("user"));
        assert!(graph.prefix_index.contains_key("user.name"));
        assert!(graph.prefix_index.contains_key("user.email"));

        // user prefix should contain all three paths
        let user_paths = graph.prefix_index.get("user").unwrap();
        assert_eq!(user_paths.len(), 3);
        assert!(user_paths.contains("user"));
        assert!(user_paths.contains("user.name"));
        assert!(user_paths.contains("user.email"));
    }

    #[test]
    fn test_get_affected_nodes_with_prefix_index() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();

        // Create real NodeIds using InstanceTree
        let mut tree = InstanceTree::new();
        let node1 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node2 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node3 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node4 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));

        // Add dependencies
        graph.add_dependency(node1, &Binding::state(vec!["user".to_string()]));
        graph.add_dependency(node2, &Binding::state(vec!["user".to_string(), "name".to_string()]));
        graph.add_dependency(node3, &Binding::state(vec!["user".to_string(), "email".to_string()]));
        graph.add_dependency(node4, &Binding::state(vec!["posts".to_string()]));

        // When "user" changes, should affect node1, node2, and node3 (but not node4)
        let affected = graph.get_affected_nodes("user");
        assert_eq!(affected.len(), 3);
        assert!(affected.contains(&node1));
        assert!(affected.contains(&node2));
        assert!(affected.contains(&node3));
        assert!(!affected.contains(&node4));
    }

    #[test]
    fn test_get_affected_nodes_child_change() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();

        // Create real NodeIds using InstanceTree
        let mut tree = InstanceTree::new();
        let node1 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
        let node2 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));

        graph.add_dependency(node1, &Binding::state(vec!["user".to_string()]));
        graph.add_dependency(node2, &Binding::state(vec!["user".to_string(), "name".to_string()]));

        // When "user.name" changes, should affect both parent "user" and child "user.name"
        let affected = graph.get_affected_nodes("user.name");
        assert_eq!(affected.len(), 2);
        assert!(affected.contains(&node1)); // parent "user" affected
        assert!(affected.contains(&node2)); // exact match "user.name" affected
    }

    #[test]
    fn test_get_affected_nodes_exact_match_only() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();

        // Create real NodeId using InstanceTree
        let mut tree = InstanceTree::new();
        let node1 = tree.create_node(&Element::new("Text"), &serde_json::json!({}));

        graph.add_dependency(node1, &Binding::state(vec!["user".to_string(), "name".to_string()]));

        // When "user.email" changes, should not affect "user.name"
        let affected = graph.get_affected_nodes("user.email");
        assert_eq!(affected.len(), 0);
    }

    #[test]
    fn test_prefix_index_performance() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();
        let mut tree = InstanceTree::new();

        // Add many dependencies with nested paths
        for i in 0..100 {
            let node_id = tree.create_node(&Element::new("Text"), &serde_json::json!({}));
            let path = vec![
                "user".to_string(),
                "profile".to_string(),
                "details".to_string(),
                format!("field{}", i),
            ];
            let binding = Binding::state(path);
            graph.add_dependency(node_id, &binding);
        }

        // Should efficiently find all affected nodes when parent changes
        let affected = graph.get_affected_nodes("user");
        assert_eq!(affected.len(), 100);

        let affected = graph.get_affected_nodes("user.profile");
        assert_eq!(affected.len(), 100);

        let affected = graph.get_affected_nodes("user.profile.details");
        assert_eq!(affected.len(), 100);
    }

    #[test]
    fn test_clear_clears_prefix_index() {
        use super::Binding;
        use crate::reconcile::tree::InstanceTree;
        use crate::ir::Element;

        let mut graph = DependencyGraph::new();

        // Create real NodeId using InstanceTree
        let mut tree = InstanceTree::new();
        let node = tree.create_node(&Element::new("Text"), &serde_json::json!({}));

        graph.add_dependency(node, &Binding::state(vec!["user".to_string(), "name".to_string()]));
        assert!(!graph.prefix_index.is_empty());

        graph.clear();
        assert!(graph.prefix_index.is_empty());
        assert!(graph.dependencies.is_empty());
        assert!(graph.node_bindings.is_empty());
    }
}
