//! Query Planner - Type-aware query optimization for heterogeneous stores
//!
//! Analyzes queries and generates optimized execution plans for cross-store
//! and cross-type queries.
//!
//! # Key capabilities
//!
//! - **Smart filter routing**: Parses filter field prefixes (e.g. `person.name`)
//!   and pushes them down to the correct target store.
//! - **JOIN condition parsing**: Splits `on` expressions like
//!   `"article.author_id = person.id"` into left/right key pairs.
//! - **Strategy selection**: Uses store types, sizes, and index availability to
//!   pick among DirectJoin / IndexLookup / HashJoin / NestedLoop.
//! - **Post-join operations**: Appends SELECT (project), ORDER BY, LIMIT/OFFSET
//!   nodes after the join tree.

use std::collections::HashMap;

use super::error::{FederationError, Result};
use super::types::{
    StoreType, UniversalQuery, QueryTarget, JoinClause, JoinType, JoinStrategy, CrossDbQuery, Filter,
};

/// Query plan node representing a single operation
#[derive(Debug, Clone)]
pub enum QueryPlanNode {
    /// Single store scan
    StoreScan {
        store_alias: String,
        store_type: StoreType,
        query: UniversalQuery,
        estimated_rows: usize,
    },
    /// Graph traversal operation
    GraphTraverse {
        store_alias: String,
        start_ids: Vec<String>,
        max_depth: u32,
        relation_types: Vec<String>,
    },
    /// Index lookup
    IndexLookup {
        index_name: String,
        keys: Vec<String>,
    },
    /// Hash join between two results
    HashJoin {
        left: Box<QueryPlanNode>,
        right: Box<QueryPlanNode>,
        /// Left-side key field (e.g. "author_id")
        left_key: String,
        /// Right-side key field (e.g. "id")
        right_key: String,
        /// Original raw expression for display (e.g. "article.author_id = person.id")
        on_expr: String,
        /// JOIN type (Inner, Left, Right, Full)
        join_type: JoinType,
    },
    /// Nested loop join
    NestedLoopJoin {
        outer: Box<QueryPlanNode>,
        inner: Box<QueryPlanNode>,
        /// Left-side key field
        left_key: String,
        /// Right-side key field
        right_key: String,
        /// Original raw expression
        on_expr: String,
        /// JOIN type
        join_type: JoinType,
    },
    /// Filter operation
    Filter {
        input: Box<QueryPlanNode>,
        predicate: String,
    },
    /// Project specific fields
    Project {
        input: Box<QueryPlanNode>,
        fields: Vec<String>,
    },
    /// Sort operation
    Sort {
        input: Box<QueryPlanNode>,
        order_by: Vec<(String, bool)>, // (field, ascending)
    },
    /// Limit/Offset
    LimitOffset {
        input: Box<QueryPlanNode>,
        limit: Option<usize>,
        offset: Option<usize>,
    },
    /// Vector similarity search
    VectorSearch {
        store_alias: String,
        query_vector: Vec<f32>,
        top_k: usize,
        min_score: Option<f32>,
        estimated_rows: usize,
    },
    /// Vector similarity join: search vectors, then join results with another store
    ///
    /// This node performs a two-phase operation:
    /// 1. Execute a vector similarity search on the vector store
    /// 2. Join the vector results with another store's results using key-based matching
    ///
    /// Typical use case: find similar vectors, then enrich with graph entity details.
    VectorSimilarityJoin {
        /// The non-vector side (e.g., a GraphTraverse or StoreScan)
        left: Box<QueryPlanNode>,
        /// Vector store alias
        vector_store_alias: String,
        /// Query vector for similarity search
        query_vector: Vec<f32>,
        /// Number of top results to retrieve from vector search
        top_k: usize,
        /// Minimum similarity score threshold
        min_score: Option<f32>,
        /// Key field on the left side to match (e.g., "id")
        left_key: String,
        /// Key field on the vector side to match (e.g., "id" or a metadata field like "entity_id")
        right_key: String,
        /// JOIN type
        join_type: JoinType,
    },
    /// Union of multiple results
    Union {
        inputs: Vec<QueryPlanNode>,
    },
}

/// Result transform specification
#[derive(Debug, Clone)]
pub struct ResultTransform {
    /// Fields to include in output
    pub fields: Vec<String>,
    /// Field renames: old_name -> new_name
    pub renames: HashMap<String, String>,
    /// Computed fields: field_name -> expression
    pub computed: HashMap<String, String>,
}

/// Complete query plan
#[derive(Debug, Clone)]
pub struct QueryPlan {
    /// Root node of the plan
    pub root: QueryPlanNode,
    /// Result transformation
    pub transform: Option<ResultTransform>,
    /// Estimated total cost
    pub estimated_cost: f64,
    /// Estimated result count
    pub estimated_rows: usize,
    /// Stores involved
    pub stores_involved: Vec<String>,
}

/// Query planner for cross-store queries
pub struct QueryPlanner {
    /// Store statistics for cost estimation
    store_stats: HashMap<String, StoreStatistics>,
}

#[derive(Debug, Clone, Default)]
struct StoreStatistics {
    row_count: usize,
    avg_row_size: usize,
    index_available: bool,
}

impl QueryPlanner {
    /// Create a new query planner
    pub fn new() -> Self {
        Self {
            store_stats: HashMap::new(),
        }
    }
    
    /// Update statistics for a store
    pub fn update_stats(&mut self, store_alias: &str, row_count: usize, avg_row_size: usize, has_index: bool) {
        self.store_stats.insert(store_alias.to_string(), StoreStatistics {
            row_count,
            avg_row_size,
            index_available: has_index,
        });
    }
    
    /// Plan a single-store query
    pub fn plan_single(&self, store_alias: &str, store_type: StoreType, query: &UniversalQuery) -> Result<QueryPlan> {
        let stats = self.store_stats.get(store_alias)
            .cloned()
            .unwrap_or_default();
        
        // Vector search: use specialized VectorSearch node for better cost estimation
        if store_type == StoreType::Vector {
            if let Some(ref vs) = query.vector_search {
                let estimated_rows = vs.top_k.min(stats.row_count.max(1));
                
                let mut root = QueryPlanNode::VectorSearch {
                    store_alias: store_alias.to_string(),
                    query_vector: vs.query_vector.clone(),
                    top_k: vs.top_k,
                    min_score: vs.min_score,
                    estimated_rows,
                };
                
                // Apply limit/offset on top of vector results
                if query.limit.is_some() || query.offset.is_some() {
                    root = QueryPlanNode::LimitOffset {
                        input: Box::new(root),
                        limit: query.limit,
                        offset: query.offset,
                    };
                }
                
                let cost = self.estimate_cost(&root);
                
                return Ok(QueryPlan {
                    root,
                    transform: None,
                    estimated_cost: cost,
                    estimated_rows,
                    stores_involved: vec![store_alias.to_string()],
                });
            }
        }
        
        // Estimate rows based on filters
        let selectivity = self.estimate_selectivity(&query.filters);
        let estimated_rows = ((stats.row_count as f64) * selectivity) as usize;
        
        // Build plan
        let mut root = QueryPlanNode::StoreScan {
            store_alias: store_alias.to_string(),
            store_type,
            query: query.clone(),
            estimated_rows,
        };
        
        // Add sort if needed
        if let Some((field, order)) = &query.order_by {
            root = QueryPlanNode::Sort {
                input: Box::new(root),
                order_by: vec![(field.clone(), matches!(order, super::types::SortOrder::Asc))],
            };
        }
        
        // Add limit/offset
        if query.limit.is_some() || query.offset.is_some() {
            root = QueryPlanNode::LimitOffset {
                input: Box::new(root),
                limit: query.limit,
                offset: query.offset,
            };
        }
        
        let cost = self.estimate_cost(&root);
        
        Ok(QueryPlan {
            root,
            transform: None,
            estimated_cost: cost,
            estimated_rows,
            stores_involved: vec![store_alias.to_string()],
        })
    }
    
    /// Plan a cross-database query
    ///
    /// The planner:
    /// 1. Routes global filters to the correct target stores
    /// 2. Builds per-target StoreScan nodes (using target.query if present)
    /// 3. Parses JOIN `on` conditions and picks the best JoinStrategy
    /// 4. Builds a left-deep join tree
    /// 5. Appends post-join operations (Filter, Project, Sort, LimitOffset)
    pub fn plan_cross_db(&self, query: &CrossDbQuery) -> Result<QueryPlan> {
        if query.targets.is_empty() {
            return Err(FederationError::QueryError("CrossDbQuery must have at least one target".to_string()));
        }

        // ── 1. Classify global filters by target affinity ──────────────
        let (routed_filters, cross_filters) = self.route_filters(&query.filters, &query.targets);

        // ── 2. Build per-target StoreScan nodes ────────────────────────
        let mut plans: Vec<(String, QueryPlanNode)> = Vec::new();
        let mut stores_involved: Vec<String> = Vec::new();

        for target in &query.targets {
            stores_involved.push(target.db_alias.clone());

            // If the target carries its own query, use it; otherwise build one
            let store_query = if let Some(ref tq) = target.query {
                tq.clone()
            } else {
                let target_filters = routed_filters
                    .get(&target.db_alias)
                    .cloned()
                    .unwrap_or_default();

                UniversalQuery {
                    target: target.type_name.clone(),
                    filters: target_filters,
                    limit: None,   // don't limit before join
                    offset: None,
                    order_by: None,
                    traverse: None,
                    vector_search: None,
                }
            };

            let stats = self.store_stats.get(&target.db_alias)
                .cloned()
                .unwrap_or_default();

            let selectivity = self.estimate_selectivity(&store_query.filters);
            let estimated_rows = ((stats.row_count as f64) * selectivity).max(1.0) as usize;

            plans.push((target.db_alias.clone(), QueryPlanNode::StoreScan {
                store_alias: target.db_alias.clone(),
                store_type: target.store_type,
                query: store_query,
                estimated_rows,
            }));
        }

        // ── 3. Build join tree using JoinClauses ───────────────────────
        let root = if plans.len() == 1 {
            plans.pop().unwrap().1
        } else {
            self.build_join_tree_v2(plans, &query.joins, &query.targets)?
        };

        // ── 4. Post-join: cross-store filters that could not be pushed down
        let root = if cross_filters.is_empty() {
            root
        } else {
            let predicate = cross_filters.iter()
                .map(|f| Self::filter_to_string(f))
                .collect::<Vec<_>>()
                .join(" AND ");
            QueryPlanNode::Filter {
                input: Box::new(root),
                predicate,
            }
        };

        // ── 5. Post-join: SELECT (projection) ─────────────────────────
        let root = if let Some(ref fields) = query.select {
            QueryPlanNode::Project {
                input: Box::new(root),
                fields: fields.clone(),
            }
        } else {
            root
        };

        // ── 6. Post-join: ORDER BY ─────────────────────────────────────
        let root = if let Some((ref field, ref order)) = query.order_by {
            QueryPlanNode::Sort {
                input: Box::new(root),
                order_by: vec![(field.clone(), matches!(order, super::types::SortOrder::Asc))],
            }
        } else {
            root
        };

        // ── 7. Post-join: LIMIT / OFFSET ──────────────────────────────
        let root = if query.limit.is_some() || query.offset.is_some() {
            QueryPlanNode::LimitOffset {
                input: Box::new(root),
                limit: query.limit,
                offset: query.offset,
            }
        } else {
            root
        };

        let estimated_rows = self.estimate_join_rows(&root);
        let cost = self.estimate_cost(&root);

        Ok(QueryPlan {
            root,
            transform: None,
            estimated_cost: cost,
            estimated_rows,
            stores_involved,
        })
    }
    
    /// Analyze cross-type join and recommend strategy
    pub fn analyze_cross_type_join(
        &self,
        left_type: StoreType,
        right_type: StoreType,
        left_size: usize,
        right_size: usize,
        has_index: bool,
    ) -> JoinStrategy {
        // Determine optimal join strategy based on types and sizes
        match (left_type, right_type) {
            (StoreType::Graph, StoreType::Graph) => {
                // Same type - can use direct join if in same store
                JoinStrategy::DirectJoin {
                    use_index: has_index,
                }
            }
            (StoreType::Graph, StoreType::Document) | (StoreType::Document, StoreType::Graph) => {
                // Cross-type join - prefer index lookup if available
                if has_index {
                    JoinStrategy::IndexLookup {
                        index_store: "cross_db_index".to_string(),
                        lookup_type: "type_cross_relation".to_string(),
                    }
                } else if left_size < 1000 || right_size < 1000 {
                    // Small dataset - nested loop is fine
                    JoinStrategy::NestedLoop {
                        batch_size: 100,
                    }
                } else {
                    // Large dataset - use hash join
                    JoinStrategy::HashJoin {
                        build_side: if left_size < right_size { "left".to_string() } else { "right".to_string() },
                    }
                }
            }
            (StoreType::Vector, _) | (_, StoreType::Vector) => {
                // Vector joins always use index lookup (similarity-based)
                JoinStrategy::IndexLookup {
                    index_store: "vector_index".to_string(),
                    lookup_type: "similarity".to_string(),
                }
            }
            _ => {
                // Default to hash join
                JoinStrategy::HashJoin {
                    build_side: "left".to_string(),
                }
            }
        }
    }
    
    // ===== Private Helpers =====
    
    fn estimate_selectivity(&self, filters: &[super::types::Filter]) -> f64 {
        if filters.is_empty() {
            return 1.0;
        }
        
        // Simple selectivity estimation based on filter variant
        // In practice, this would use histograms and statistics
        let mut selectivity: f64 = 1.0;
        for filter in filters {
            selectivity *= match filter {
                super::types::Filter::Eq(..) => 0.1,
                super::types::Filter::Ne(..) => 0.9,
                super::types::Filter::Gt(..) | super::types::Filter::Lt(..) => 0.3,
                super::types::Filter::Gte(..) | super::types::Filter::Lte(..) => 0.5,
                super::types::Filter::Contains(..) => 0.2,
                super::types::Filter::In(..) => 0.2,
                super::types::Filter::And(inner) => self.estimate_selectivity(inner),
                super::types::Filter::Or(inner) => {
                    1.0 - inner.iter().fold(1.0, |acc, _| acc * 0.5)
                }
            };
        }
        
        selectivity.max(0.001) // At least 0.1%
    }
    
    fn estimate_cost(&self, node: &QueryPlanNode) -> f64 {
        match node {
            QueryPlanNode::StoreScan { estimated_rows, .. } => {
                *estimated_rows as f64 * 1.0
            }
            QueryPlanNode::GraphTraverse { max_depth, .. } => {
                100.0 * (*max_depth as f64).powi(2)
            }
            QueryPlanNode::IndexLookup { keys, .. } => {
                keys.len() as f64 * 0.1
            }
            QueryPlanNode::HashJoin { left, right, .. } => {
                self.estimate_cost(left) + self.estimate_cost(right) + 50.0
            }
            QueryPlanNode::NestedLoopJoin { outer, inner, .. } => {
                let outer_cost = self.estimate_cost(outer);
                let inner_cost = self.estimate_cost(inner);
                outer_cost + (outer_cost * inner_cost / 10.0)
            }
            QueryPlanNode::Filter { input, .. } => {
                self.estimate_cost(input) * 1.1
            }
            QueryPlanNode::Project { input, .. } => {
                self.estimate_cost(input)
            }
            QueryPlanNode::Sort { input, .. } => {
                let input_cost = self.estimate_cost(input);
                input_cost * (input_cost.log2().max(1.0))
            }
            QueryPlanNode::LimitOffset { input, .. } => {
                self.estimate_cost(input)
            }
            QueryPlanNode::VectorSearch { top_k, .. } => {
                // HNSW search is O(log N) per query, much cheaper than full scan
                (*top_k as f64) * 0.5
            }
            QueryPlanNode::VectorSimilarityJoin { left, top_k, .. } => {
                // Cost = vector search + left scan + hash join overhead
                self.estimate_cost(left) + (*top_k as f64) * 0.5 + 30.0
            }
            QueryPlanNode::Union { inputs } => {
                inputs.iter().map(|n| self.estimate_cost(n)).sum()
            }
        }
    }
    
    fn estimate_join_rows(&self, node: &QueryPlanNode) -> usize {
        match node {
            QueryPlanNode::StoreScan { estimated_rows, .. } => *estimated_rows,
            QueryPlanNode::HashJoin { left, right, .. } => {
                let l = self.estimate_join_rows(left);
                let r = self.estimate_join_rows(right);
                ((l.min(r) as f64) * 0.1).max(1.0) as usize
            }
            QueryPlanNode::NestedLoopJoin { outer, inner, .. } => {
                let o = self.estimate_join_rows(outer);
                let i = self.estimate_join_rows(inner);
                ((o as f64) * (i as f64) * 0.01).max(1.0) as usize
            }
            QueryPlanNode::Filter { input, .. } => {
                ((self.estimate_join_rows(input) as f64) * 0.5).max(1.0) as usize
            }
            QueryPlanNode::Project { input, .. } => {
                self.estimate_join_rows(input)
            }
            QueryPlanNode::Sort { input, .. } => {
                self.estimate_join_rows(input)
            }
            QueryPlanNode::LimitOffset { input, limit, .. } => {
                let input_rows = self.estimate_join_rows(input);
                limit.unwrap_or(input_rows).min(input_rows)
            }
            QueryPlanNode::VectorSearch { estimated_rows, .. } => *estimated_rows,
            QueryPlanNode::VectorSimilarityJoin { left, top_k, .. } => {
                let left_rows = self.estimate_join_rows(left);
                // Estimated: min(left_rows, top_k) matched pairs
                left_rows.min(*top_k)
            }
            QueryPlanNode::Union { inputs } => {
                inputs.iter().map(|n| self.estimate_join_rows(n)).sum()
            }
            _ => 100, // Default estimate
        }
    }
    
    // ===== Predicate Pushdown =====

    /// Route global filters to the correct target stores (predicate pushdown).
    ///
    /// **Rules:**
    /// 1. Filters with a target prefix (e.g. `person.name`) are pushed to
    ///    the matching target, with the prefix stripped.
    /// 2. `And(...)` filters are **decomposed**: each sub-condition is routed
    ///    independently, so parts can be pushed to different stores.
    /// 3. `Or(...)` filters can only be pushed if **all** branches belong to
    ///    the **same** target; otherwise they stay as a cross-filter.
    /// 4. Unqualified fields (no prefix) are matched by **store-type heuristics**:
    ///    `entity_type`/`name` → Graph, `doc_type`/`tags`/`content` → Document.
    /// 5. If still ambiguous and there's only one target, assign to it.
    fn route_filters(
        &self,
        filters: &[Filter],
        targets: &[QueryTarget],
    ) -> (HashMap<String, Vec<Filter>>, Vec<Filter>) {
        let mut routed: HashMap<String, Vec<Filter>> = HashMap::new();
        let mut cross: Vec<Filter> = Vec::new();

        for filter in filters {
            self.route_single_filter(filter, targets, &mut routed, &mut cross);
        }

        (routed, cross)
    }

    /// Route a single filter, handling AND decomposition recursively.
    fn route_single_filter(
        &self,
        filter: &Filter,
        targets: &[QueryTarget],
        routed: &mut HashMap<String, Vec<Filter>>,
        cross: &mut Vec<Filter>,
    ) {
        match filter {
            // Decompose AND: each sub-condition can go to a different store
            Filter::And(inner) => {
                for sub in inner {
                    self.route_single_filter(sub, targets, routed, cross);
                }
            }
            // OR: can only push if ALL branches target the same store
            Filter::Or(inner) => {
                let mut targets_found: Vec<Option<String>> = Vec::new();
                for sub in inner {
                    let field = Self::filter_field(sub);
                    targets_found.push(
                        self.find_target_for_field(&field, targets)
                            .map(|t| t.db_alias.clone())
                    );
                }
                // Check if all branches target the same store
                let first = targets_found.first().and_then(|v| v.as_ref());
                let all_same = first.is_some() && targets_found.iter().all(|v| v.as_ref() == first);
                if all_same {
                    let alias = first.unwrap().clone();
                    let stripped = Filter::Or(
                        inner.iter().map(Self::strip_prefix_filter).collect()
                    );
                    routed.entry(alias).or_default().push(stripped);
                } else {
                    cross.push(filter.clone());
                }
            }
            // Leaf filters: try prefix match, then heuristic match
            _ => {
                let field_name = Self::filter_field(filter);
                if let Some(target) = self.find_target_for_field(&field_name, targets) {
                    let stripped = Self::strip_prefix_filter(filter);
                    routed.entry(target.db_alias.clone())
                        .or_default()
                        .push(stripped);
                } else {
                    cross.push(filter.clone());
                }
            }
        }
    }

    /// Extract the primary field name from a filter variant.
    fn filter_field(filter: &Filter) -> String {
        match filter {
            Filter::Eq(f, _) | Filter::Ne(f, _)
            | Filter::Gt(f, _) | Filter::Gte(f, _)
            | Filter::Lt(f, _) | Filter::Lte(f, _)
            | Filter::Contains(f, _) => f.clone(),
            Filter::In(f, _) => f.clone(),
            Filter::And(_) | Filter::Or(_) => String::new(),
        }
    }

    /// Find which target a prefixed field belongs to.
    ///
    /// Matching order:
    /// 1. Explicit prefix: `person.name` → find target with type_name/db_alias "person"
    /// 2. Store-type heuristic: `entity_type`→Graph, `doc_type`→Document, `vector`→Vector
    /// 3. Single target: if only one target exists, assign to it
    fn find_target_for_field<'a>(
        &self,
        field: &str,
        targets: &'a [QueryTarget],
    ) -> Option<&'a QueryTarget> {
        // 1. Explicit prefix match
        if let Some(dot_pos) = field.find('.') {
            let prefix = &field[..dot_pos];
            let found = targets.iter().find(|t| {
                t.type_name.eq_ignore_ascii_case(prefix) || t.db_alias.eq_ignore_ascii_case(prefix)
            });
            if found.is_some() {
                return found;
            }
        }

        // 2. Store-type heuristic for unqualified fields
        let bare_field = if let Some(pos) = field.find('.') {
            &field[pos + 1..]
        } else {
            field
        };

        let inferred_type = match bare_field {
            "entity_type" | "name" | "source" | "target" | "relation_type" => Some(StoreType::Graph),
            "doc_type" | "tags" | "content" => Some(StoreType::Document),
            "vector" | "score" | "embedding" | "similarity" | "distance" | "dimension" => Some(StoreType::Vector),
            _ => None,
        };

        if let Some(st) = inferred_type {
            let found = targets.iter().find(|t| t.store_type == st);
            if found.is_some() {
                return found;
            }
        }

        // 3. Single target fallback
        if targets.len() == 1 {
            Some(&targets[0])
        } else {
            None
        }
    }

    /// Strip the `prefix.` from a filter's field name, returning a new Filter.
    fn strip_prefix_filter(filter: &Filter) -> Filter {
        let strip = |f: &str| -> String {
            if let Some(pos) = f.find('.') {
                f[pos + 1..].to_string()
            } else {
                f.to_string()
            }
        };

        match filter {
            Filter::Eq(f, v)       => Filter::Eq(strip(f), v.clone()),
            Filter::Ne(f, v)       => Filter::Ne(strip(f), v.clone()),
            Filter::Gt(f, v)       => Filter::Gt(strip(f), v.clone()),
            Filter::Gte(f, v)      => Filter::Gte(strip(f), v.clone()),
            Filter::Lt(f, v)       => Filter::Lt(strip(f), v.clone()),
            Filter::Lte(f, v)      => Filter::Lte(strip(f), v.clone()),
            Filter::In(f, v)       => Filter::In(strip(f), v.clone()),
            Filter::Contains(f, v) => Filter::Contains(strip(f), v.clone()),
            Filter::And(inner)     => Filter::And(inner.iter().map(Self::strip_prefix_filter).collect()),
            Filter::Or(inner)      => Filter::Or(inner.iter().map(Self::strip_prefix_filter).collect()),
        }
    }

    /// Convert a Filter to a human-readable predicate string (for post-join Filter nodes).
    fn filter_to_string(filter: &Filter) -> String {
        match filter {
            Filter::Eq(f, v)       => format!("{} = {}", f, v),
            Filter::Ne(f, v)       => format!("{} != {}", f, v),
            Filter::Gt(f, v)       => format!("{} > {}", f, v),
            Filter::Gte(f, v)      => format!("{} >= {}", f, v),
            Filter::Lt(f, v)       => format!("{} < {}", f, v),
            Filter::Lte(f, v)      => format!("{} <= {}", f, v),
            Filter::In(f, vs)      => format!("{} IN ({:?})", f, vs),
            Filter::Contains(f, v) => format!("{} CONTAINS '{}'", f, v),
            Filter::And(inner)     => inner.iter()
                .map(Self::filter_to_string)
                .collect::<Vec<_>>()
                .join(" AND "),
            Filter::Or(inner)      => format!("({})", inner.iter()
                .map(Self::filter_to_string)
                .collect::<Vec<_>>()
                .join(" OR ")),
        }
    }

    // ===== JOIN Condition Parsing =====

    /// Parsed result of a JOIN `on` expression.
    ///
    /// Given `"article.author_id = person.id"`, produces:
    /// - `left_prefix`:  `"article"`
    /// - `left_field`:   `"author_id"`
    /// - `right_prefix`: `"person"`
    /// - `right_field`:  `"id"`
    fn parse_join_condition(on: &str) -> (String, String, String, String) {
        // Normalise: trim, replace multiple spaces, handle `=` or `==`
        let on = on.trim();
        let parts: Vec<&str> = on.splitn(2, '=').collect();
        let (left_expr, right_expr) = if parts.len() == 2 {
            (parts[0].trim().trim_end_matches('='), parts[1].trim().trim_start_matches('='))
        } else {
            (on, "id")
        };

        let split_qualified = |expr: &str| -> (String, String) {
            if let Some(pos) = expr.rfind('.') {
                (expr[..pos].to_string(), expr[pos + 1..].to_string())
            } else {
                (String::new(), expr.to_string())
            }
        };

        let (lp, lf) = split_qualified(left_expr.trim());
        let (rp, rf) = split_qualified(right_expr.trim());
        (lp, lf, rp, rf)
    }

    // ===== Join Tree Construction V2 =====

    /// Build a strategy-aware, left-deep join tree.
    ///
    /// For each `JoinClause` it:
    /// 1. Parses the `on` expression into left/right keys
    /// 2. Resolves the right-side plan node by matching the join target
    /// 3. Calls `analyze_cross_type_join` to select the best strategy
    /// 4. Creates a HashJoin or NestedLoopJoin node accordingly
    fn build_join_tree_v2(
        &self,
        mut plans: Vec<(String, QueryPlanNode)>,
        joins: &[JoinClause],
        targets: &[QueryTarget],
    ) -> Result<QueryPlanNode> {
        if plans.is_empty() {
            return Err(FederationError::QueryError("No plans to join".to_string()));
        }
        if plans.len() == 1 {
            return Ok(plans.pop().unwrap().1);
        }

        // Start with the first plan as the left side
        let (_, mut result) = plans.remove(0);

        for (i, (alias, plan)) in plans.into_iter().enumerate() {
            let join_clause = joins.get(i);

            // ── Parse on condition ────────────────────────────────────
            let (on_expr, left_key, right_key, join_type) = if let Some(jc) = join_clause {
                let (_lp, lf, _rp, rf) = Self::parse_join_condition(&jc.on);
                (jc.on.clone(), lf, rf, jc.join_type)
            } else {
                ("id = id".to_string(), "id".to_string(), "id".to_string(), JoinType::Inner)
            };

            // ── Determine store types for strategy selection ──────────
            let left_type = Self::node_store_type(&result);
            let right_type = Self::node_store_type(&plan);
            let left_rows = self.estimate_join_rows(&result);
            let right_rows = self.estimate_join_rows(&plan);

            let right_target = targets.iter().find(|t| t.db_alias == alias);
            let has_index = right_target
                .map(|t| self.store_stats.get(&t.db_alias).map_or(false, |s| s.index_available))
                .unwrap_or(false);

            let strategy = self.analyze_cross_type_join(
                left_type, right_type, left_rows, right_rows, has_index,
            );

            // ── Build the appropriate join node ──────────────────────

            // Special case: if the right side involves vector search, use VectorSimilarityJoin
            let vector_search_spec = if let QueryPlanNode::StoreScan { store_type: StoreType::Vector, query, store_alias, .. } = &plan {
                query.vector_search.as_ref().map(|vs| (store_alias.clone(), vs.clone()))
            } else if let QueryPlanNode::VectorSearch { store_alias, query_vector, top_k, min_score, .. } = &plan {
                Some((store_alias.clone(), super::types::VectorSearchSpec {
                    query_vector: query_vector.clone(),
                    top_k: *top_k,
                    min_score: *min_score,
                    metadata_filters: vec![],
                }))
            } else {
                None
            };

            result = if let Some((vs_alias, vs_spec)) = vector_search_spec {
                // Use VectorSimilarityJoin for vector-involving joins
                QueryPlanNode::VectorSimilarityJoin {
                    left: Box::new(result),
                    vector_store_alias: vs_alias,
                    query_vector: vs_spec.query_vector,
                    top_k: vs_spec.top_k,
                    min_score: vs_spec.min_score,
                    left_key,
                    right_key,
                    join_type,
                }
            } else if matches!(strategy, JoinStrategy::NestedLoop { .. }) {
                QueryPlanNode::NestedLoopJoin {
                    outer: Box::new(result),
                    inner: Box::new(plan),
                    left_key,
                    right_key,
                    on_expr,
                    join_type,
                }
            } else {
                // HashJoin / DirectJoin / IndexLookup all use hash-join execution
                QueryPlanNode::HashJoin {
                    left: Box::new(result),
                    right: Box::new(plan),
                    left_key,
                    right_key,
                    on_expr,
                    join_type,
                }
            };
        }

        Ok(result)
    }

    /// Extract the StoreType from a plan node (best effort).
    fn node_store_type(node: &QueryPlanNode) -> StoreType {
        match node {
            QueryPlanNode::StoreScan { store_type, .. } => *store_type,
            QueryPlanNode::GraphTraverse { .. } => StoreType::Graph,
            QueryPlanNode::HashJoin { left, .. } => Self::node_store_type(left),
            QueryPlanNode::NestedLoopJoin { outer, .. } => Self::node_store_type(outer),
            QueryPlanNode::Filter { input, .. }
            | QueryPlanNode::Project { input, .. }
            | QueryPlanNode::Sort { input, .. }
            | QueryPlanNode::LimitOffset { input, .. } => Self::node_store_type(input),
            QueryPlanNode::VectorSimilarityJoin { .. } => StoreType::Vector,
            QueryPlanNode::Union { inputs } => inputs.first()
                .map(Self::node_store_type)
                .unwrap_or(StoreType::Graph),
            _ => StoreType::Graph,
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::types::SortOrder;
    
    #[test]
    fn test_single_store_plan() {
        let mut planner = QueryPlanner::new();
        planner.update_stats("social", 10000, 200, true);
        
        let query = UniversalQuery {
            target: "person".to_string(),
            filters: vec![],
            limit: Some(10),
            offset: None,
            order_by: None,
            traverse: None,
            vector_search: None,
        };
        
        let plan = planner.plan_single("social", StoreType::Graph, &query).unwrap();
        
        assert!(plan.estimated_cost > 0.0);
        assert_eq!(plan.stores_involved.len(), 1);
    }
    
    #[test]
    fn test_cross_type_join_strategy() {
        let planner = QueryPlanner::new();
        
        // Graph-Document with index
        let strategy = planner.analyze_cross_type_join(
            StoreType::Graph,
            StoreType::Document,
            10000,
            5000,
            true,
        );
        assert!(matches!(strategy, JoinStrategy::IndexLookup { .. }));
        
        // Graph-Document without index, large datasets
        let strategy = planner.analyze_cross_type_join(
            StoreType::Graph,
            StoreType::Document,
            10000,
            5000,
            false,
        );
        assert!(matches!(strategy, JoinStrategy::HashJoin { .. }));
        
        // Small dataset → NestedLoop
        let strategy = planner.analyze_cross_type_join(
            StoreType::Graph,
            StoreType::Document,
            500,
            200,
            false,
        );
        assert!(matches!(strategy, JoinStrategy::NestedLoop { .. }));
    }

    #[test]
    fn test_parse_join_condition() {
        // Standard qualified condition
        let (lp, lf, rp, rf) = QueryPlanner::parse_join_condition("article.author_id = person.id");
        assert_eq!(lp, "article");
        assert_eq!(lf, "author_id");
        assert_eq!(rp, "person");
        assert_eq!(rf, "id");
        
        // Unqualified fields
        let (lp, lf, rp, rf) = QueryPlanner::parse_join_condition("author_id = id");
        assert_eq!(lp, "");
        assert_eq!(lf, "author_id");
        assert_eq!(rp, "");
        assert_eq!(rf, "id");
        
        // With extra spaces
        let (lp, lf, rp, rf) = QueryPlanner::parse_join_condition("  docs.ref_id  =  social.entity_id  ");
        assert_eq!(lp, "docs");
        assert_eq!(lf, "ref_id");
        assert_eq!(rp, "social");
        assert_eq!(rf, "entity_id");
    }

    #[test]
    fn test_filter_routing() {
        let planner = QueryPlanner::new();
        
        let targets = vec![
            QueryTarget::graph("social", "person"),
            QueryTarget::document("knowledge", "article"),
        ];
        
        let filters = vec![
            Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
            Filter::Contains("article.title".to_string(), "Rust".to_string()),
            Filter::Gt("score".to_string(), serde_json::json!(80)),
        ];
        
        let (routed, cross) = planner.route_filters(&filters, &targets);
        
        // "person.name" → social (person target)
        assert_eq!(routed.get("social").map_or(0, |v| v.len()), 1);
        // "article.title" → knowledge (article target)
        assert_eq!(routed.get("knowledge").map_or(0, |v| v.len()), 1);
        // "score" has no prefix → cross filter
        assert_eq!(cross.len(), 1);
        
        // Verify prefix was stripped
        if let Some(social_filters) = routed.get("social") {
            assert!(matches!(&social_filters[0], Filter::Eq(f, _) if f == "name"));
        }
    }

    #[test]
    fn test_plan_cross_db_basic() {
        let mut planner = QueryPlanner::new();
        planner.update_stats("social", 10000, 200, true);
        planner.update_stats("knowledge", 5000, 500, false);
        
        let query = CrossDbQuery {
            targets: vec![
                QueryTarget::graph("social", "person"),
                QueryTarget::document("knowledge", "article"),
            ],
            joins: vec![
                JoinClause {
                    target: QueryTarget::document("knowledge", "article"),
                    on: "person.id = article.author_id".to_string(),
                    join_type: JoinType::Inner,
                },
            ],
            filters: vec![
                Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
            ],
            select: Some(vec!["person.name".to_string(), "article.title".to_string()]),
            order_by: Some(("article.title".to_string(), SortOrder::Asc)),
            limit: Some(20),
            offset: None,
            streaming: false,
            batch_size: None,
        };
        
        let plan = planner.plan_cross_db(&query).unwrap();
        
        // Should involve 2 stores
        assert_eq!(plan.stores_involved.len(), 2);
        assert!(plan.stores_involved.contains(&"social".to_string()));
        assert!(plan.stores_involved.contains(&"knowledge".to_string()));
        
        // Cost and rows should be positive
        assert!(plan.estimated_cost > 0.0);
        assert!(plan.estimated_rows > 0);
    }

    #[test]
    fn test_plan_cross_db_with_post_join_ops() {
        let planner = QueryPlanner::new();
        
        let query = CrossDbQuery {
            targets: vec![
                QueryTarget::graph("social", "person"),
                QueryTarget::document("docs", "article"),
            ],
            joins: vec![
                JoinClause {
                    target: QueryTarget::document("docs", "article"),
                    on: "person.id = article.author_id".to_string(),
                    join_type: JoinType::Left,
                },
            ],
            filters: vec![],
            select: Some(vec!["name".to_string(), "title".to_string()]),
            order_by: Some(("name".to_string(), SortOrder::Desc)),
            limit: Some(10),
            offset: Some(5),
            streaming: false,
            batch_size: None,
        };
        
        let plan = planner.plan_cross_db(&query).unwrap();
        
        // The root should be LimitOffset wrapping Sort wrapping Project wrapping HashJoin
        assert!(matches!(plan.root, QueryPlanNode::LimitOffset { .. }));
        if let QueryPlanNode::LimitOffset { input, limit, offset } = &plan.root {
            assert_eq!(*limit, Some(10));
            assert_eq!(*offset, Some(5));
            assert!(matches!(**input, QueryPlanNode::Sort { .. }));
        }
    }

    #[test]
    fn test_plan_cross_db_empty_targets_error() {
        let planner = QueryPlanner::new();
        let query = CrossDbQuery::default();
        assert!(planner.plan_cross_db(&query).is_err());
    }

    #[test]
    fn test_plan_cross_db_single_target() {
        let planner = QueryPlanner::new();
        let query = CrossDbQuery {
            targets: vec![QueryTarget::graph("social", "person")],
            joins: vec![],
            filters: vec![Filter::Eq("name".to_string(), serde_json::json!("Bob"))],
            select: None,
            order_by: None,
            limit: Some(5),
            offset: None,
            streaming: false,
            batch_size: None,
        };
        
        let plan = planner.plan_cross_db(&query).unwrap();
        assert_eq!(plan.stores_involved.len(), 1);
        // Single target: filter should be pushed down, root is LimitOffset(StoreScan)
        assert!(matches!(plan.root, QueryPlanNode::LimitOffset { .. }));
    }

    #[test]
    fn test_filter_to_string() {
        let s = QueryPlanner::filter_to_string(
            &Filter::Eq("name".to_string(), serde_json::json!("Alice")),
        );
        assert_eq!(s, "name = \"Alice\"");
        
        let s = QueryPlanner::filter_to_string(
            &Filter::Contains("title".to_string(), "Rust".to_string()),
        );
        assert_eq!(s, "title CONTAINS 'Rust'");
    }

    // ===== Predicate Pushdown Tests =====

    #[test]
    fn test_pushdown_and_decomposition() {
        let planner = QueryPlanner::new();
        let targets = vec![
            QueryTarget::graph("social", "person"),
            QueryTarget::document("knowledge", "article"),
        ];

        // AND with mixed targets: should decompose and push each to correct store
        let filters = vec![
            Filter::And(vec![
                Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
                Filter::Contains("article.title".to_string(), "Rust".to_string()),
                Filter::Gt("score".to_string(), serde_json::json!(80)),
            ]),
        ];

        let (routed, cross) = planner.route_filters(&filters, &targets);

        // person.name → social
        assert_eq!(routed.get("social").map_or(0, |v| v.len()), 1);
        // article.title → knowledge
        assert_eq!(routed.get("knowledge").map_or(0, |v| v.len()), 1);
        // score → ambiguous → cross
        assert_eq!(cross.len(), 1);
    }

    #[test]
    fn test_pushdown_or_same_target() {
        let planner = QueryPlanner::new();
        let targets = vec![
            QueryTarget::graph("social", "person"),
            QueryTarget::document("knowledge", "article"),
        ];

        // OR where all branches target the same store → can push down
        let filters = vec![
            Filter::Or(vec![
                Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
                Filter::Eq("person.name".to_string(), serde_json::json!("Bob")),
            ]),
        ];

        let (routed, cross) = planner.route_filters(&filters, &targets);

        // Both branches target "social" → pushed down as Or(...)
        assert_eq!(routed.get("social").map_or(0, |v| v.len()), 1);
        assert!(cross.is_empty());

        // Verify it's an Or variant with stripped prefixes
        if let Some(social_filters) = routed.get("social") {
            assert!(matches!(&social_filters[0], Filter::Or(inner) if inner.len() == 2));
        }
    }

    #[test]
    fn test_pushdown_or_mixed_targets_stays_cross() {
        let planner = QueryPlanner::new();
        let targets = vec![
            QueryTarget::graph("social", "person"),
            QueryTarget::document("knowledge", "article"),
        ];

        // OR where branches target different stores → cannot push down
        let filters = vec![
            Filter::Or(vec![
                Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
                Filter::Eq("article.title".to_string(), serde_json::json!("Rust")),
            ]),
        ];

        let (routed, cross) = planner.route_filters(&filters, &targets);

        assert!(routed.is_empty() || routed.values().all(|v| v.is_empty()));
        assert_eq!(cross.len(), 1);
    }

    #[test]
    fn test_pushdown_heuristic_unqualified_fields() {
        let planner = QueryPlanner::new();
        let targets = vec![
            QueryTarget::graph("social", "person"),
            QueryTarget::document("knowledge", "article"),
        ];

        // "entity_type" → Graph heuristic, "doc_type" → Document heuristic
        let filters = vec![
            Filter::Eq("entity_type".to_string(), serde_json::json!("person")),
            Filter::Eq("doc_type".to_string(), serde_json::json!("article")),
        ];

        let (routed, cross) = planner.route_filters(&filters, &targets);

        // entity_type → social (Graph)
        assert_eq!(routed.get("social").map_or(0, |v| v.len()), 1);
        // doc_type → knowledge (Document)
        assert_eq!(routed.get("knowledge").map_or(0, |v| v.len()), 1);
        // Nothing left as cross
        assert!(cross.is_empty());
    }

    #[test]
    fn test_pushdown_nested_and_in_cross_query() {
        // Verify that AND decomposition works end-to-end in plan_cross_db
        let mut planner = QueryPlanner::new();
        planner.update_stats("social", 10000, 200, true);
        planner.update_stats("docs", 5000, 500, false);

        let query = CrossDbQuery {
            targets: vec![
                QueryTarget::graph("social", "person"),
                QueryTarget::document("docs", "article"),
            ],
            joins: vec![
                JoinClause {
                    target: QueryTarget::document("docs", "article"),
                    on: "person.id = article.author_id".to_string(),
                    join_type: JoinType::Inner,
                },
            ],
            // AND with two pushable conditions + one cross condition
            filters: vec![
                Filter::And(vec![
                    Filter::Eq("person.name".to_string(), serde_json::json!("Alice")),
                    Filter::Contains("article.title".to_string(), "Rust".to_string()),
                    Filter::Gt("combined_score".to_string(), serde_json::json!(90)),
                ]),
            ],
            select: None,
            order_by: None,
            limit: Some(10),
            offset: None,
            streaming: false,
            batch_size: None,
        };

        let plan = planner.plan_cross_db(&query).unwrap();
        assert_eq!(plan.stores_involved.len(), 2);
        // Plan should have a Filter node for the cross condition "combined_score > 90"
        // wrapped around the join, then LimitOffset on top
        assert!(matches!(plan.root, QueryPlanNode::LimitOffset { .. }));
    }
}
