tegdb 0.5.0

//! Query planner and optimizer for TegDB
//!
//! This module provides a simple rule-based query planner that takes parsed SQL statements
//! and generates optimized execution plans. The planner focuses on fast, predictable
//! optimizations without complex cost estimation.

use crate::parser::{
    ComparisonOperator, Condition, CreateTableStatement, DeleteStatement, DropTableStatement,
    InsertStatement, OrderByClause, OrderByItem, SelectStatement, SqlValue, Statement,
    UpdateStatement,
};
use crate::query_processor::{ColumnInfo, TableSchema};
use crate::Result;
use std::collections::HashMap;
use std::rc::Rc;

/// Vector similarity functions for vector search
#[derive(Debug, Clone)]
pub enum VectorSimilarityFunction {
    CosineSimilarity,
    EuclideanDistance,
    DotProduct,
}

/// A query execution plan generated by the planner
#[derive(Debug, Clone)]
pub enum ExecutionPlan {
    /// Direct primary key lookup (most efficient) - exact equality on single PK column
    PrimaryKeyLookup {
        table: String,
        pk_value: SqlValue,
        selected_columns: Vec<crate::parser::Expression>,
        additional_filter: Option<Condition>,
    },
    /// Primary key range scan - range conditions on PK column
    TableRangeScan {
        table: String,
        selected_columns: Vec<crate::parser::Expression>,
        pk_range: PkRange,
        additional_filter: Option<Condition>,
        limit: Option<u64>,
    },
    /// Full table scan - no PK conditions or non-PK filters
    TableScan {
        table: String,
        selected_columns: Vec<crate::parser::Expression>,
        filter: Option<Condition>,
        limit: Option<u64>,
    },
    /// Insert operation plan
    Insert {
        table: String,
        rows: Vec<HashMap<String, SqlValue>>,
        conflict_resolution: ConflictResolution,
    },
    /// Update operation plan
    Update {
        table: String,
        assignments: Vec<Assignment>,
        scan_plan: Box<ExecutionPlan>,
    },
    /// Delete operation plan
    Delete {
        table: String,
        scan_plan: Box<ExecutionPlan>,
    },
    /// DDL operations
    CreateTable {
        table: String,
        schema: TableSchema,
    },
    DropTable {
        table: String,
        if_exists: bool,
    },
    CreateIndex {
        index_name: String,
        table_name: String,
        column_name: String,
        unique: bool,
    },
    DropIndex {
        index_name: String,
        if_exists: bool,
    },
    CreateExtension {
        extension_name: String,
        library_path: Option<String>,
    },
    DropExtension {
        extension_name: String,
    },
    /// Transaction control
    Begin,
    Commit,
    Rollback,
    /// Index scan operation plan
    IndexScan {
        table: String,
        index: String,
        column_value: crate::parser::SqlValue,
        selected_columns: Vec<crate::parser::Expression>,
        additional_filter: Option<Condition>,
    },
    /// Vector similarity search plan
    VectorSearch {
        table: String,
        index: String,
        query_vector: Vec<f64>,
        similarity_function: VectorSimilarityFunction,
        k: usize,
        selected_columns: Vec<crate::parser::Expression>,
        additional_filter: Option<Condition>,
    },
    /// Sort operation plan
    Sort {
        input_plan: Box<ExecutionPlan>,
        order_by_items: Vec<OrderByItem>,
        schema: std::rc::Rc<TableSchema>,
        query_schema: crate::query_processor::QuerySchema,
        limit: Option<u64>,
    },
}

/// Primary key range for range scans
#[derive(Debug, Clone)]
pub struct PkRange {
    pub start_bound: Option<PkBound>,
    pub end_bound: Option<PkBound>,
}

/// Primary key bound (inclusive or exclusive)
#[derive(Debug, Clone)]
pub struct PkBound {
    pub value: SqlValue,
    pub inclusive: bool,
}

/// Assignment for update operations
#[derive(Debug, Clone)]
pub struct Assignment {
    pub column: String,
    pub value: crate::parser::Expression,
}

/// Conflict resolution strategy for inserts
#[derive(Debug, Clone)]
pub enum ConflictResolution {
    Error,   // Fail on constraint violation
    Ignore,  // Skip conflicting rows
    Replace, // Replace existing rows
}

/// Simple rule-based query planner
pub struct QueryPlanner {
    /// Table schemas for optimization decisions
    table_schemas: HashMap<String, Rc<TableSchema>>,
}

impl QueryPlanner {
    /// Create a new query planner with table schemas
    pub fn new(table_schemas: HashMap<String, Rc<TableSchema>>) -> Self {
        Self { table_schemas }
    }

    /// Find a suitable index for a WHERE condition
    fn find_suitable_index(
        &self,
        table_name: &str,
        condition: &crate::parser::Condition,
    ) -> Option<(String, String, crate::parser::SqlValue)> {
        // Get table schema
        let schema = self.table_schemas.get(table_name)?;

        // Check each index to see if it can be used
        for index in &schema.indexes {
            if index.index_type != crate::parser::IndexType::BTree {
                continue;
            }
            if let Some(value) = Self::extract_indexable_value(condition, &index.column_name) {
                return Some((index.name.clone(), index.column_name.clone(), value));
            }
        }
        None
    }

    /// Extract a value that can be used for index lookup from a condition
    fn extract_indexable_value(
        condition: &crate::parser::Condition,
        column_name: &str,
    ) -> Option<crate::parser::SqlValue> {
        match condition {
            crate::parser::Condition::Comparison {
                left,
                operator,
                right,
            } => {
                if let crate::parser::Expression::Column(col_name) = left {
                    if col_name == column_name
                        && *operator == crate::parser::ComparisonOperator::Equal
                    {
                        Some(right.clone())
                    } else {
                        None
                    }
                } else {
                    None
                }
            }
            crate::parser::Condition::And(left, right) => {
                Self::extract_indexable_value(left, column_name)
                    .or_else(|| Self::extract_indexable_value(right, column_name))
            }
            crate::parser::Condition::Or(left, right) => {
                // For OR conditions, we need both sides to use the same index
                let left_val = Self::extract_indexable_value(left, column_name)?;
                let right_val = Self::extract_indexable_value(right, column_name)?;
                if left_val == right_val {
                    Some(left_val)
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    /// Generate an optimized execution plan for a SQL statement
    pub fn plan(&self, statement: Statement) -> Result<ExecutionPlan> {
        match statement {
            Statement::Select(select) => self.plan_select(select),
            Statement::Insert(insert) => self.plan_insert(insert),
            Statement::Update(update) => self.plan_update(update),
            Statement::Delete(delete) => self.plan_delete(delete),
            Statement::CreateTable(create) => self.plan_create_table(create),
            Statement::DropTable(drop) => self.plan_drop_table(drop),
            Statement::CreateIndex(create) => self.plan_create_index(create),
            Statement::DropIndex(drop) => self.plan_drop_index(drop),
            Statement::CreateExtension(create) => self.plan_create_extension(create),
            Statement::DropExtension(drop) => self.plan_drop_extension(drop),
            Statement::Begin => Ok(ExecutionPlan::Begin),
            Statement::Commit => Ok(ExecutionPlan::Commit),
            Statement::Rollback => Ok(ExecutionPlan::Rollback),
        }
    }

    /// Plan SELECT statement with index optimization
    fn plan_select(&self, select: SelectStatement) -> Result<ExecutionPlan> {
        // Try primary key lookup first (most efficient)
        if let Some(plan) = self.try_primary_key_lookup(&select)? {
            return self.wrap_with_sort(plan, &select.order_by);
        }

        // Try primary key range scan
        if let Some(plan) = self.try_primary_key_range_scan(&select)? {
            return self.wrap_with_sort(plan, &select.order_by);
        }

        // Try vector similarity search if we have vector indexes and similarity functions
        if let Some(plan) = self.try_vector_similarity_search(&select)? {
            return self.wrap_with_sort(plan, &select.order_by);
        }

        // Try index scan if we have a suitable index
        if let Some(where_clause) = &select.where_clause {
            if let Some((index_name, _column_name, value)) =
                self.find_suitable_index(&select.table, &where_clause.condition)
            {
                let expr_columns =
                    self.normalize_selected_columns(&select.table, &select.columns)?;

                // Extract additional filter conditions (non-indexed conditions)
                let additional_filter = self.extract_non_index_conditions(
                    &select.table,
                    &where_clause.condition,
                    &index_name,
                )?;

                let plan = ExecutionPlan::IndexScan {
                    table: select.table,
                    index: index_name,
                    column_value: value,
                    selected_columns: expr_columns,
                    additional_filter,
                };
                return self.wrap_with_sort(plan, &select.order_by);
            }
        }

        // Fall back to table scan
        let plan = self.plan_table_scan(&select)?;
        self.wrap_with_sort(plan, &select.order_by)
    }

    /// Try to create a primary key lookup plan (exact equality)
    fn try_primary_key_lookup(&self, select: &SelectStatement) -> Result<Option<ExecutionPlan>> {
        if let Some(ref where_clause) = select.where_clause {
            if let Some(pk_value) =
                self.extract_single_pk_equality(&select.table, &where_clause.condition)?
            {
                let expr_columns =
                    self.normalize_selected_columns(&select.table, &select.columns)?;
                let additional_filter =
                    self.extract_non_pk_conditions(&select.table, &where_clause.condition)?;
                return Ok(Some(ExecutionPlan::PrimaryKeyLookup {
                    table: select.table.clone(),
                    pk_value,
                    selected_columns: expr_columns,
                    additional_filter,
                }));
            }
        }
        Ok(None)
    }

    fn extract_single_pk_equality(
        &self,
        table_name: &str,
        condition: &Condition,
    ) -> Result<Option<SqlValue>> {
        let pk_columns = self.get_primary_key_columns(table_name)?;
        if pk_columns.len() != 1 {
            return Err(crate::Error::SqlError(
                "Composite primary keys are not supported for this operation".to_string(),
            ));
        }
        let pk_col = &pk_columns[0];
        let mut found = None;
        Self::find_pk_equality(condition, pk_col, &mut found);
        Ok(found)
    }

    fn find_pk_equality(cond: &Condition, pk_col: &str, found: &mut Option<SqlValue>) {
        match cond {
            Condition::Comparison {
                left: crate::parser::Expression::Column(col_name),
                operator,
                right,
            } if col_name == pk_col && matches!(operator, ComparisonOperator::Equal) => {
                *found = Some(right.clone());
            }
            Condition::Comparison { .. } => {}
            Condition::And(left, right) => {
                Self::find_pk_equality(left, pk_col, found);
                Self::find_pk_equality(right, pk_col, found);
            }
            _ => {}
        }
    }

    /// Try to create a primary key range scan plan
    fn try_primary_key_range_scan(
        &self,
        select: &SelectStatement,
    ) -> Result<Option<ExecutionPlan>> {
        if let Some(ref where_clause) = select.where_clause {
            if let Some(pk_range) =
                self.extract_pk_range_conditions(&select.table, &where_clause.condition)?
            {
                let additional_filter =
                    self.extract_non_pk_conditions(&select.table, &where_clause.condition)?;
                let expr_columns =
                    self.normalize_selected_columns(&select.table, &select.columns)?;

                return Ok(Some(ExecutionPlan::TableRangeScan {
                    table: select.table.clone(),
                    selected_columns: expr_columns,
                    pk_range,
                    additional_filter,
                    limit: select.limit,
                }));
            }
        }

        Ok(None)
    }

    /// Plan table scan (full scan)
    fn plan_table_scan(&self, select: &SelectStatement) -> Result<ExecutionPlan> {
        let expr_columns = self.normalize_selected_columns(&select.table, &select.columns)?;
        let filter = select.where_clause.as_ref().map(|w| w.condition.clone());

        Ok(ExecutionPlan::TableScan {
            table: select.table.clone(),
            selected_columns: expr_columns,
            filter,
            limit: select.limit,
        })
    }

    /// Plan INSERT statement
    fn plan_insert(&self, insert: InsertStatement) -> Result<ExecutionPlan> {
        // Convert expressions to values, then to HashMap format
        let mut rows = Vec::with_capacity(insert.values.len());
        for expr_row in insert.values {
            let mut row_data = HashMap::with_capacity(insert.columns.len());
            for (i, col_name) in insert.columns.iter().enumerate() {
                if let Some(expr) = expr_row.get(i) {
                    // Evaluate expression to get SqlValue
                    let context = HashMap::new(); // Empty context for INSERT
                    let value = expr.evaluate(&context).map_err(crate::Error::Other)?;
                    row_data.insert(col_name.clone(), value);
                }
            }
            rows.push(row_data);
        }

        Ok(ExecutionPlan::Insert {
            table: insert.table,
            rows,
            conflict_resolution: ConflictResolution::Error,
        })
    }

    /// Plan UPDATE statement
    fn plan_update(&self, update: UpdateStatement) -> Result<ExecutionPlan> {
        // Convert assignments to planner format
        let assignments: Vec<Assignment> = update
            .assignments
            .into_iter()
            .map(|a| Assignment {
                column: a.column,
                value: a.value,
            })
            .collect();

        // Always expand ["*"] to all columns for scan plan
        let expr_columns: Vec<crate::parser::Expression> = self.normalize_selected_columns(
            &update.table,
            &[crate::parser::Expression::Column("*".to_string())],
        )?;

        // Create scan plan for the rows to update
        let scan_plan = if let Some(ref where_clause) = update.where_clause {
            // Try to optimize the scan based on WHERE clause
            if let Some(pk_value) =
                self.extract_single_pk_equality(&update.table, &where_clause.condition)?
            {
                let pk_columns = self.get_primary_key_columns(&update.table)?;
                if pk_columns.len() == 1 {
                    ExecutionPlan::PrimaryKeyLookup {
                        table: update.table.clone(),
                        pk_value,
                        selected_columns: expr_columns.clone(),
                        additional_filter: None,
                    }
                } else {
                    // Try range scan
                    if let Some(pk_range) =
                        self.extract_pk_range_conditions(&update.table, &where_clause.condition)?
                    {
                        ExecutionPlan::TableRangeScan {
                            table: update.table.clone(),
                            selected_columns: expr_columns.clone(),
                            pk_range,
                            additional_filter: None,
                            limit: None,
                        }
                    } else {
                        ExecutionPlan::TableScan {
                            table: update.table.clone(),
                            selected_columns: expr_columns.clone(),
                            filter: Some(where_clause.condition.clone()),
                            limit: None,
                        }
                    }
                }
            } else {
                // Try range scan
                if let Some(pk_range) =
                    self.extract_pk_range_conditions(&update.table, &where_clause.condition)?
                {
                    ExecutionPlan::TableRangeScan {
                        table: update.table.clone(),
                        selected_columns: expr_columns.clone(),
                        pk_range,
                        additional_filter: None,
                        limit: None,
                    }
                } else {
                    ExecutionPlan::TableScan {
                        table: update.table.clone(),
                        selected_columns: expr_columns.clone(),
                        filter: Some(where_clause.condition.clone()),
                        limit: None,
                    }
                }
            }
        } else {
            // No WHERE clause - scan all rows
            ExecutionPlan::TableScan {
                table: update.table.clone(),
                selected_columns: expr_columns.clone(),
                filter: None,
                limit: None,
            }
        };

        Ok(ExecutionPlan::Update {
            table: update.table,
            assignments,
            scan_plan: Box::new(scan_plan),
        })
    }

    /// Plan DELETE statement
    fn plan_delete(&self, delete: DeleteStatement) -> Result<ExecutionPlan> {
        // Always expand ["*"] to all columns for scan plan
        let expr_columns: Vec<crate::parser::Expression> = self.normalize_selected_columns(
            &delete.table,
            &[crate::parser::Expression::Column("*".to_string())],
        )?;

        // Create scan plan for the rows to delete
        let scan_plan = if let Some(ref where_clause) = delete.where_clause {
            // Try to optimize the scan based on WHERE clause
            if let Some(pk_value) =
                self.extract_single_pk_equality(&delete.table, &where_clause.condition)?
            {
                let pk_columns = self.get_primary_key_columns(&delete.table)?;
                if pk_columns.len() == 1 {
                    ExecutionPlan::PrimaryKeyLookup {
                        table: delete.table.clone(),
                        pk_value,
                        selected_columns: expr_columns.clone(),
                        additional_filter: None,
                    }
                } else {
                    // Try range scan
                    if let Some(pk_range) =
                        self.extract_pk_range_conditions(&delete.table, &where_clause.condition)?
                    {
                        ExecutionPlan::TableRangeScan {
                            table: delete.table.clone(),
                            selected_columns: expr_columns.clone(),
                            pk_range,
                            additional_filter: None,
                            limit: None,
                        }
                    } else {
                        ExecutionPlan::TableScan {
                            table: delete.table.clone(),
                            selected_columns: expr_columns.clone(),
                            filter: Some(where_clause.condition.clone()),
                            limit: None,
                        }
                    }
                }
            } else {
                // Try range scan
                if let Some(pk_range) =
                    self.extract_pk_range_conditions(&delete.table, &where_clause.condition)?
                {
                    ExecutionPlan::TableRangeScan {
                        table: delete.table.clone(),
                        selected_columns: expr_columns.clone(),
                        pk_range,
                        additional_filter: None,
                        limit: None,
                    }
                } else {
                    ExecutionPlan::TableScan {
                        table: delete.table.clone(),
                        selected_columns: expr_columns.clone(),
                        filter: Some(where_clause.condition.clone()),
                        limit: None,
                    }
                }
            }
        } else {
            // No WHERE clause - scan all rows
            ExecutionPlan::TableScan {
                table: delete.table.clone(),
                selected_columns: expr_columns.clone(),
                filter: None,
                limit: None,
            }
        };

        Ok(ExecutionPlan::Delete {
            table: delete.table,
            scan_plan: Box::new(scan_plan),
        })
    }

    /// Plan CREATE TABLE statement
    fn plan_create_table(&self, create: CreateTableStatement) -> Result<ExecutionPlan> {
        let schema = TableSchema {
            name: create.table.clone(),
            columns: create
                .columns
                .iter()
                .map(|col| ColumnInfo {
                    name: col.name.clone(),
                    data_type: col.data_type.clone(),
                    constraints: col.constraints.clone(),
                    storage_offset: 0,
                    storage_size: 0,
                    storage_type_code: 0,
                })
                .collect(),
            indexes: vec![], // Initialize indexes as empty
        };

        Ok(ExecutionPlan::CreateTable {
            table: create.table,
            schema,
        })
    }

    /// Plan DROP TABLE statement
    fn plan_drop_table(&self, drop: DropTableStatement) -> Result<ExecutionPlan> {
        Ok(ExecutionPlan::DropTable {
            table: drop.table,
            if_exists: drop.if_exists,
        })
    }

    /// Plan CREATE INDEX statement
    fn plan_create_index(
        &self,
        create: crate::parser::CreateIndexStatement,
    ) -> Result<ExecutionPlan> {
        let table_name = create.table_name.clone();
        let column_name = create.column_name.clone();
        let unique = create.unique;

        if let Some(_schema) = self.table_schemas.get(&table_name) {
            // Note: We don't actually modify the schema here since that's done in the executor
            return Ok(ExecutionPlan::CreateIndex {
                index_name: create.index_name,
                table_name,
                column_name,
                unique,
            });
        }
        Err(crate::Error::TableNotFound(table_name))
    }

    /// Plan DROP INDEX statement
    fn plan_drop_index(&self, drop: crate::parser::DropIndexStatement) -> Result<ExecutionPlan> {
        let index_name = drop.index_name.clone();
        let if_exists = drop.if_exists;

        // Note: We don't actually modify the schema here since that's done in the executor
        Ok(ExecutionPlan::DropIndex {
            index_name,
            if_exists,
        })
    }

    /// Plan CREATE EXTENSION statement
    fn plan_create_extension(
        &self,
        create: crate::parser::CreateExtensionStatement,
    ) -> Result<ExecutionPlan> {
        Ok(ExecutionPlan::CreateExtension {
            extension_name: create.extension_name,
            library_path: create.library_path,
        })
    }

    /// Plan DROP EXTENSION statement
    fn plan_drop_extension(
        &self,
        drop: crate::parser::DropExtensionStatement,
    ) -> Result<ExecutionPlan> {
        Ok(ExecutionPlan::DropExtension {
            extension_name: drop.extension_name,
        })
    }

    /// Try to create a vector similarity search plan
    fn try_vector_similarity_search(
        &self,
        select: &SelectStatement,
    ) -> Result<Option<ExecutionPlan>> {
        // Check if we have ORDER BY with vector similarity functions
        if let Some(order_by) = &select.order_by {
            for item in &order_by.items {
                if let Some((similarity_func, query_vector, k)) =
                    self.extract_vector_similarity_from_expr(&item.expression)?
                {
                    // Find a suitable vector index for this table
                    if let Some(index_name) = self.find_vector_index(&select.table)? {
                        let expr_columns =
                            self.normalize_selected_columns(&select.table, &select.columns)?;

                        // Extract additional filter conditions
                        let additional_filter = select
                            .where_clause
                            .as_ref()
                            .map(|where_clause| where_clause.condition.clone());

                        return Ok(Some(ExecutionPlan::VectorSearch {
                            table: select.table.clone(),
                            index: index_name,
                            query_vector,
                            similarity_function: similarity_func,
                            k,
                            selected_columns: expr_columns,
                            additional_filter,
                        }));
                    }
                }
            }
        }

        // Check WHERE clause for vector similarity conditions
        if let Some(where_clause) = &select.where_clause {
            if let Some((similarity_func, query_vector, _)) =
                self.extract_vector_similarity_from_condition(&where_clause.condition)?
            {
                if let Some(index_name) = self.find_vector_index(&select.table)? {
                    let expr_columns =
                        self.normalize_selected_columns(&select.table, &select.columns)?;

                    // Extract additional filter conditions (non-vector conditions)
                    let additional_filter =
                        self.extract_non_vector_conditions(&where_clause.condition)?;

                    return Ok(Some(ExecutionPlan::VectorSearch {
                        table: select.table.clone(),
                        index: index_name,
                        query_vector,
                        similarity_function: similarity_func,
                        k: 100, // Default k for WHERE clause similarity
                        selected_columns: expr_columns,
                        additional_filter,
                    }));
                }
            }
        }

        Ok(None)
    }

    /// Extract vector similarity information from an expression
    fn extract_vector_similarity_from_expr(
        &self,
        expr: &crate::parser::Expression,
    ) -> Result<Option<(VectorSimilarityFunction, Vec<f64>, usize)>> {
        if let crate::parser::Expression::FunctionCall { name, args } = expr {
            let similarity_func = match name.to_uppercase().as_str() {
                "COSINE_SIMILARITY" => VectorSimilarityFunction::CosineSimilarity,
                "EUCLIDEAN_DISTANCE" => VectorSimilarityFunction::EuclideanDistance,
                "DOT_PRODUCT" => VectorSimilarityFunction::DotProduct,
                _ => return Ok(None),
            };

            if args.len() == 2 {
                if let (
                    crate::parser::Expression::Column(_),
                    crate::parser::Expression::Value(SqlValue::Vector(vector)),
                ) = (&args[0], &args[1])
                {
                    return Ok(Some((similarity_func, vector.clone(), 10))); // Default k=10
                }
            }
        }
        Ok(None)
    }

    /// Extract vector similarity information from a condition
    fn extract_vector_similarity_from_condition(
        &self,
        condition: &Condition,
    ) -> Result<Option<(VectorSimilarityFunction, Vec<f64>, usize)>> {
        if let Condition::Comparison {
            left,
            operator: _,
            right,
        } = condition
        {
            if let (
                crate::parser::Expression::FunctionCall { name, args },
                SqlValue::Real(_threshold),
            ) = (left, right)
            {
                let similarity_func = match name.to_uppercase().as_str() {
                    "COSINE_SIMILARITY" => VectorSimilarityFunction::CosineSimilarity,
                    "EUCLIDEAN_DISTANCE" => VectorSimilarityFunction::EuclideanDistance,
                    "DOT_PRODUCT" => VectorSimilarityFunction::DotProduct,
                    _ => return Ok(None),
                };

                if args.len() == 2 {
                    if let (
                        crate::parser::Expression::Column(_),
                        crate::parser::Expression::Value(SqlValue::Vector(vector)),
                    ) = (&args[0], &args[1])
                    {
                        return Ok(Some((similarity_func, vector.clone(), 100)));
                        // Default k=100 for WHERE
                    }
                }
            }
        }
        Ok(None)
    }

    /// Find a vector index for the given table
    fn find_vector_index(&self, table_name: &str) -> Result<Option<String>> {
        if let Some(schema) = self.table_schemas.get(table_name) {
            for index in &schema.indexes {
                if matches!(
                    index.index_type,
                    crate::parser::IndexType::HNSW
                        | crate::parser::IndexType::IVF
                        | crate::parser::IndexType::LSH
                ) {
                    return Ok(Some(index.name.clone()));
                }
            }
        }
        Ok(None)
    }

    /// Extract non-vector conditions from a WHERE clause
    fn extract_non_vector_conditions(&self, _condition: &Condition) -> Result<Option<Condition>> {
        // For now, return None - in a full implementation, we'd filter out vector similarity conditions
        // and return the remaining conditions
        Ok(None)
    }

    /// Extract primary key range conditions from WHERE clause
    fn extract_pk_range_conditions(
        &self,
        table_name: &str,
        condition: &Condition,
    ) -> Result<Option<PkRange>> {
        let pk_columns = self.get_primary_key_columns(table_name)?;

        // Check if we have any PK conditions (but not all for exact lookup)
        let mut pk_conditions = HashMap::new();
        Self::collect_pk_range_values(condition, &pk_columns, &mut pk_conditions);

        if pk_conditions.is_empty() {
            return Ok(None);
        }

        // Build range from conditions
        let mut start_bound = None;
        let mut end_bound = None;

        for (_column, conditions) in pk_conditions {
            for (operator, value) in conditions {
                match operator {
                    ComparisonOperator::GreaterThan => {
                        if start_bound.is_none() {
                            start_bound = Some(PkBound {
                                value: value.clone(),
                                inclusive: false,
                            });
                        }
                    }
                    ComparisonOperator::GreaterThanOrEqual => {
                        if start_bound.is_none() {
                            start_bound = Some(PkBound {
                                value: value.clone(),
                                inclusive: true,
                            });
                        }
                    }
                    ComparisonOperator::LessThan => {
                        if end_bound.is_none() {
                            end_bound = Some(PkBound {
                                value: value.clone(),
                                inclusive: false,
                            });
                        }
                    }
                    ComparisonOperator::LessThanOrEqual => {
                        if end_bound.is_none() {
                            end_bound = Some(PkBound {
                                value: value.clone(),
                                inclusive: true,
                            });
                        }
                    }
                    _ => {} // Ignore other operators for range scan
                }
            }
        }

        if start_bound.is_some() || end_bound.is_some() {
            Ok(Some(PkRange {
                start_bound,
                end_bound,
            }))
        } else {
            Ok(None)
        }
    }

    /// Recursively collect primary key range values from conditions
    fn collect_pk_range_values(
        condition: &Condition,
        pk_columns: &[String],
        pk_conditions: &mut HashMap<String, Vec<(ComparisonOperator, SqlValue)>>,
    ) {
        match condition {
            Condition::Comparison {
                left,
                operator,
                right,
            } => {
                if let crate::parser::Expression::Column(col_name) = left {
                    if pk_columns.contains(col_name) {
                        pk_conditions
                            .entry(col_name.clone())
                            .or_default()
                            .push((*operator, right.clone()));
                    }
                }
            }
            Condition::Between {
                column: _column,
                low: _low,
                high: _high,
            } => {
                if pk_columns.contains(_column) {
                    // Convert BETWEEN to >= and <= conditions for range optimization
                    pk_conditions
                        .entry(_column.clone())
                        .or_default()
                        .push((ComparisonOperator::GreaterThanOrEqual, _low.clone()));
                    pk_conditions
                        .entry(_column.clone())
                        .or_default()
                        .push((ComparisonOperator::LessThanOrEqual, _high.clone()));
                }
            }
            Condition::And(left_cond, right_cond) => {
                Self::collect_pk_range_values(left_cond, pk_columns, pk_conditions);
                Self::collect_pk_range_values(right_cond, pk_columns, pk_conditions);
            }
            Condition::Or(_, _) => {
                // For OR conditions, we cannot optimize with range scan
            }
        }
    }

    /// Extract non-PK or non-indexed conditions by excluding a set of column names.
    fn extract_conditions_excluding_columns(
        excluded_columns: &std::collections::HashSet<String>,
        condition: &Condition,
    ) -> Option<Condition> {
        match condition {
            Condition::Comparison { left, .. } => {
                if let crate::parser::Expression::Column(col_name) = left {
                    if excluded_columns.contains(col_name) {
                        None
                    } else {
                        Some(condition.clone())
                    }
                } else {
                    Some(condition.clone())
                }
            }
            Condition::Between { column, .. } => {
                if excluded_columns.contains(column) {
                    None
                } else {
                    Some(condition.clone())
                }
            }
            Condition::And(left, right) => {
                let l = Self::extract_conditions_excluding_columns(excluded_columns, left);
                let r = Self::extract_conditions_excluding_columns(excluded_columns, right);
                match (l, r) {
                    (Some(lc), Some(rc)) => Some(Condition::And(Box::new(lc), Box::new(rc))),
                    (Some(lc), None) => Some(lc),
                    (None, Some(rc)) => Some(rc),
                    (None, None) => None,
                }
            }
            Condition::Or(left, right) => {
                let l = Self::extract_conditions_excluding_columns(excluded_columns, left);
                let r = Self::extract_conditions_excluding_columns(excluded_columns, right);
                match (l, r) {
                    (Some(lc), Some(rc)) => Some(Condition::Or(Box::new(lc), Box::new(rc))),
                    (Some(lc), None) => Some(lc),
                    (None, Some(rc)) => Some(rc),
                    (None, None) => None,
                }
            }
        }
    }

    /// Extract non-primary key conditions for additional filtering
    fn extract_non_pk_conditions(
        &self,
        table_name: &str,
        condition: &Condition,
    ) -> Result<Option<Condition>> {
        let pk_columns = self.get_primary_key_columns(table_name)?;
        let excluded: std::collections::HashSet<String> = pk_columns.into_iter().collect();
        Ok(Self::extract_conditions_excluding_columns(
            &excluded, condition,
        ))
    }

    /// Extract non-indexed conditions for additional filtering
    fn extract_non_index_conditions(
        &self,
        table_name: &str,
        condition: &Condition,
        index_name: &str,
    ) -> Result<Option<Condition>> {
        let schema = self
            .table_schemas
            .get(table_name)
            .ok_or_else(|| crate::Error::Other(format!("Table '{table_name}' not found")))?;
        let indexed_column = schema
            .indexes
            .iter()
            .find(|idx| idx.name == index_name)
            .map(|idx| idx.column_name.clone())
            .ok_or_else(|| crate::Error::Other(format!("Index '{index_name}' not found")))?;
        let mut excluded = std::collections::HashSet::new();
        excluded.insert(indexed_column);
        Ok(Self::extract_conditions_excluding_columns(
            &excluded, condition,
        ))
    }

    /// Normalize selected columns: if ["*"] is given, expand to all columns from the schema
    fn normalize_selected_columns(
        &self,
        table_name: &str,
        columns: &[crate::parser::Expression],
    ) -> Result<Vec<crate::parser::Expression>> {
        if columns.len() == 1 {
            if let crate::parser::Expression::Column(ref star) = columns[0] {
                if star == "*" {
                    // Expand * to all columns from the schema
                    if let Some(schema) = self.table_schemas.get(table_name) {
                        return Ok(schema
                            .columns
                            .iter()
                            .map(|c| crate::parser::Expression::Column(c.name.clone()))
                            .collect());
                    } else {
                        return Err(crate::Error::Other(format!(
                            "Table '{table_name}' not found for column expansion"
                        )));
                    }
                }
            }
        }

        // Analyze function calls to extract referenced columns
        let mut all_columns = Vec::new();
        let mut referenced_columns = std::collections::HashSet::new();

        for expr in columns {
            all_columns.push(expr.clone());
            Self::extract_columns_from_expression(expr, &mut referenced_columns);
        }

        // Add any referenced columns that aren't already in the selection
        for column_name in referenced_columns {
            let column_expr = crate::parser::Expression::Column(column_name);
            if !all_columns.iter().any(|expr| expr == &column_expr) {
                all_columns.push(column_expr);
            }
        }

        Ok(all_columns)
    }

    fn extract_columns_from_expression(
        expr: &crate::parser::Expression,
        columns: &mut std::collections::HashSet<String>,
    ) {
        match expr {
            crate::parser::Expression::Column(name) => {
                columns.insert(name.clone());
            }
            crate::parser::Expression::FunctionCall { args, .. } => {
                for arg in args {
                    Self::extract_columns_from_expression(arg, columns);
                }
            }
            crate::parser::Expression::BinaryOp { left, right, .. } => {
                Self::extract_columns_from_expression(left, columns);
                Self::extract_columns_from_expression(right, columns);
            }
            crate::parser::Expression::Value(_) => {
                // No columns to extract from literal values
            }
            crate::parser::Expression::AggregateFunction { arg, .. } => {
                Self::extract_columns_from_expression(arg, columns);
            }
        }
    }

    /// Get primary key columns for a table
    fn get_primary_key_columns(&self, table_name: &str) -> Result<Vec<String>> {
        if let Some(schema) = self.table_schemas.get(table_name) {
            let pk_columns: Vec<String> = schema
                .columns
                .iter()
                .filter(|col| {
                    col.constraints
                        .contains(&crate::parser::ColumnConstraint::PrimaryKey)
                })
                .map(|col| col.name.clone())
                .collect();
            Ok(pk_columns)
        } else {
            Err(crate::Error::TableNotFound(table_name.to_string()))
        }
    }

    /// Update table schemas (called when DDL operations occur)
    pub fn update_table_schema(&mut self, table_name: String, schema: Rc<TableSchema>) {
        self.table_schemas.insert(table_name, schema);
    }

    /// Remove table schema (called when table is dropped)
    pub fn remove_table_schema(&mut self, table_name: &str) {
        self.table_schemas.remove(table_name);
    }

    /// Get current table schemas
    pub fn table_schemas(&self) -> &HashMap<String, Rc<TableSchema>> {
        &self.table_schemas
    }

    /// Wrap a plan with a Sort plan if ORDER BY is present
    fn wrap_with_sort(
        &self,
        plan: ExecutionPlan,
        order_by: &Option<OrderByClause>,
    ) -> Result<ExecutionPlan> {
        if let Some(order_by_clause) = order_by {
            let table_name = plan.primary_table().unwrap_or("");
            let schema = self.table_schemas.get(table_name).unwrap().clone();
            // Use the selected columns from the plan if possible
            let selected_columns = match &plan {
                ExecutionPlan::PrimaryKeyLookup {
                    selected_columns, ..
                } => selected_columns.clone(),
                ExecutionPlan::TableRangeScan {
                    selected_columns, ..
                } => selected_columns.clone(),
                ExecutionPlan::TableScan {
                    selected_columns, ..
                } => selected_columns.clone(),
                ExecutionPlan::IndexScan {
                    selected_columns, ..
                } => selected_columns.clone(),
                ExecutionPlan::VectorSearch {
                    selected_columns, ..
                } => selected_columns.clone(),
                _ => schema
                    .columns
                    .iter()
                    .map(|c| crate::parser::Expression::Column(c.name.clone()))
                    .collect(),
            };
            let query_schema = crate::query_processor::QuerySchema::new_with_expressions(
                &selected_columns,
                &schema,
            );

            // Extract LIMIT from the input plan
            let limit = match &plan {
                ExecutionPlan::TableScan { limit, .. } => *limit,
                ExecutionPlan::TableRangeScan { limit, .. } => *limit,
                _ => None,
            };

            Ok(ExecutionPlan::Sort {
                input_plan: Box::new(plan),
                order_by_items: order_by_clause.items.clone(),
                schema,
                query_schema,
                limit,
            })
        } else {
            Ok(plan)
        }
    }
}

/// Helper functions for plan analysis and debugging
impl ExecutionPlan {
    /// Get the main table involved in this plan
    pub fn primary_table(&self) -> Option<&str> {
        match self {
            ExecutionPlan::PrimaryKeyLookup { table, .. } => Some(table),
            ExecutionPlan::TableRangeScan { table, .. } => Some(table),
            ExecutionPlan::TableScan { table, .. } => Some(table),
            ExecutionPlan::Insert { table, .. } => Some(table),
            ExecutionPlan::Update { table, .. } => Some(table),
            ExecutionPlan::Delete { table, .. } => Some(table),
            ExecutionPlan::CreateTable { table, .. } => Some(table),
            ExecutionPlan::DropTable { table, .. } => Some(table),
            ExecutionPlan::CreateIndex { table_name, .. } => Some(table_name),
            ExecutionPlan::CreateExtension { .. } | ExecutionPlan::DropExtension { .. } => None,
            ExecutionPlan::IndexScan { table, .. } => Some(table),
            ExecutionPlan::VectorSearch { table, .. } => Some(table),
            ExecutionPlan::Sort { input_plan, .. } => input_plan.primary_table(),
            _ => None,
        }
    }

    /// Check if this plan involves a table scan
    pub fn requires_table_scan(&self) -> bool {
        match self {
            ExecutionPlan::TableScan { .. } => true,
            ExecutionPlan::TableRangeScan { .. } => true,
            ExecutionPlan::Update { scan_plan, .. } => scan_plan.requires_table_scan(),
            ExecutionPlan::Delete { scan_plan, .. } => scan_plan.requires_table_scan(),
            ExecutionPlan::IndexScan { .. } => true,
            ExecutionPlan::Sort { input_plan, .. } => input_plan.requires_table_scan(),
            _ => false,
        }
    }

    /// Get a human-readable description of the plan
    pub fn describe(&self) -> String {
        match self {
            ExecutionPlan::PrimaryKeyLookup {
                table, pk_value, ..
            } => {
                format!("Primary Key Lookup on {table} (key: {pk_value:?})")
            }
            ExecutionPlan::TableRangeScan {
                table, pk_range, ..
            } => {
                let range_desc = match (&pk_range.start_bound, &pk_range.end_bound) {
                    (Some(start), Some(end)) => {
                        let start_value = &start.value;
                        let end_value = &end.value;
                        format!("range: {start_value:?} to {end_value:?}")
                    }
                    (Some(start), None) => {
                        let value = &start.value;
                        format!("range: >= {value:?}")
                    }
                    (None, Some(end)) => {
                        let value = &end.value;
                        format!("range: <= {value:?}")
                    }
                    (None, None) => "range: all".to_string(),
                };
                format!("Table Range Scan on {table} ({range_desc})")
            }
            ExecutionPlan::TableScan {
                table,
                selected_columns,
                filter,
                limit,
            } => {
                let columns_str = selected_columns
                    .iter()
                    .map(|expr| match expr {
                        crate::parser::Expression::Column(name) => name.clone(),
                        _ => format!("{expr:?}"),
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("TableScan({table}, columns=[{columns_str}], filter={filter:?}, limit={limit:?})")
            }
            ExecutionPlan::Insert { table, rows, .. } => {
                let row_count = rows.len();
                format!("Insert into {table} ({row_count} rows)")
            }
            ExecutionPlan::Update {
                table, scan_plan, ..
            } => {
                let plan_desc = scan_plan.describe();
                format!("Update {table} via {plan_desc}")
            }
            ExecutionPlan::Delete {
                table, scan_plan, ..
            } => {
                let plan_desc = scan_plan.describe();
                format!("Delete from {table} via {plan_desc}")
            }
            ExecutionPlan::CreateTable { table, .. } => {
                format!("Create Table {table}")
            }
            ExecutionPlan::DropTable { table, .. } => {
                format!("Drop Table {table}")
            }
            ExecutionPlan::CreateIndex {
                index_name,
                table_name,
                column_name,
                unique,
            } => {
                format!("Create Index {index_name} on {table_name} (column: {column_name}, unique: {unique})")
            }
            ExecutionPlan::DropIndex {
                index_name,
                if_exists,
            } => {
                format!("Drop Index {index_name} (if_exists: {if_exists})")
            }
            ExecutionPlan::Begin => "Begin Transaction".to_string(),
            ExecutionPlan::Commit => "Commit Transaction".to_string(),
            ExecutionPlan::Rollback => "Rollback Transaction".to_string(),
            ExecutionPlan::CreateExtension { extension_name, .. } => {
                format!("Create Extension: {extension_name}")
            }
            ExecutionPlan::DropExtension { extension_name } => {
                format!("Drop Extension: {extension_name}")
            }
            ExecutionPlan::IndexScan {
                table,
                index,
                column_value,
                selected_columns,
                additional_filter: _,
            } => {
                let columns_str = selected_columns
                    .iter()
                    .map(|expr| match expr {
                        crate::parser::Expression::Column(name) => name.clone(),
                        _ => format!("{expr:?}"),
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                format!(
                    "Index Scan on {table} (index: {index}, column: {column_value:?}, selected: {columns_str})"
                )
            }
            ExecutionPlan::VectorSearch {
                table,
                index,
                query_vector,
                similarity_function,
                k,
                selected_columns,
                additional_filter: _,
            } => {
                let columns_str = selected_columns
                    .iter()
                    .map(|expr| match expr {
                        crate::parser::Expression::Column(name) => name.clone(),
                        _ => format!("{expr:?}"),
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                let func_name = match similarity_function {
                    VectorSimilarityFunction::CosineSimilarity => "COSINE_SIMILARITY",
                    VectorSimilarityFunction::EuclideanDistance => "EUCLIDEAN_DISTANCE",
                    VectorSimilarityFunction::DotProduct => "DOT_PRODUCT",
                };
                format!(
                    "Vector Search on {table} (index: {index}, function: {func_name}, k: {k}, vector: {query_vector:?}, selected: {columns_str})"
                )
            }
            ExecutionPlan::Sort {
                input_plan,
                order_by_items,
                ..
            } => {
                let order_desc = order_by_items
                    .iter()
                    .map(|item| {
                        let expr_str = match &item.expression {
                            crate::parser::Expression::Column(name) => name.clone(),
                            _ => format!("{expr:?}", expr = item.expression),
                        };
                        let direction = item.direction;
                        format!("{expr_str} {direction:?}")
                    })
                    .collect::<Vec<_>>()
                    .join(", ");
                let input_desc = input_plan.describe();
                format!("Sort on {input_desc} (order: {order_desc})")
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::{ColumnConstraint, DataType, SqlValue, Statement, WhereClause};
    use crate::query_processor::{ColumnInfo, TableSchema};
    use std::collections::HashMap;
    use std::rc::Rc;

    fn create_test_schema() -> HashMap<String, Rc<TableSchema>> {
        let mut schemas = HashMap::new();

        // Users table with single primary key
        let mut users_schema = TableSchema {
            name: "users".to_string(),
            columns: vec![
                ColumnInfo {
                    name: "id".to_string(),
                    data_type: DataType::Integer,
                    constraints: vec![ColumnConstraint::PrimaryKey],
                    storage_offset: 0,
                    storage_size: 0,
                    storage_type_code: 0,
                },
                ColumnInfo {
                    name: "name".to_string(),
                    data_type: DataType::Text(None),
                    constraints: vec![],
                    storage_offset: 0,
                    storage_size: 0,
                    storage_type_code: 0,
                },
                ColumnInfo {
                    name: "email".to_string(),
                    data_type: DataType::Text(None),
                    constraints: vec![],
                    storage_offset: 0,
                    storage_size: 0,
                    storage_type_code: 0,
                },
            ],
            indexes: vec![], // Initialize indexes as empty
        };
        let _ = crate::catalog::Catalog::compute_table_metadata(&mut users_schema);
        schemas.insert("users".to_string(), Rc::new(users_schema));

        schemas
    }

    #[test]
    fn test_primary_key_optimization() {
        let schemas = create_test_schema();
        let planner = QueryPlanner::new(schemas);

        let select = SelectStatement {
            columns: vec![
                crate::parser::Expression::Column("name".to_string()),
                crate::parser::Expression::Column("email".to_string()),
            ],
            table: "users".to_string(),
            where_clause: Some(WhereClause {
                condition: Condition::Comparison {
                    left: crate::parser::Expression::Column("id".to_string()),
                    operator: ComparisonOperator::Equal,
                    right: SqlValue::Integer(123),
                },
            }),
            order_by: None,
            limit: None,
        };

        let plan = planner.plan(Statement::Select(select)).unwrap();

        match plan {
            ExecutionPlan::PrimaryKeyLookup {
                table, pk_value, ..
            } => {
                assert_eq!(table, "users");
                assert_eq!(pk_value, SqlValue::Integer(123));
            }
            _ => panic!("Expected PrimaryKeyLookup plan"),
        }
    }

    #[test]
    fn test_primary_key_range_scan() {
        let schemas = create_test_schema();
        let planner = QueryPlanner::new(schemas);

        let select = SelectStatement {
            columns: vec![crate::parser::Expression::Column("*".to_string())],
            table: "users".to_string(),
            where_clause: Some(WhereClause {
                condition: Condition::And(
                    Box::new(Condition::Comparison {
                        left: crate::parser::Expression::Column("id".to_string()),
                        operator: ComparisonOperator::GreaterThan,
                        right: SqlValue::Integer(1),
                    }),
                    Box::new(Condition::Comparison {
                        left: crate::parser::Expression::Column("id".to_string()),
                        operator: ComparisonOperator::LessThan,
                        right: SqlValue::Integer(10),
                    }),
                ),
            }),
            order_by: None,
            limit: Some(10),
        };

        let plan = planner.plan(Statement::Select(select)).unwrap();

        match plan {
            ExecutionPlan::TableRangeScan {
                table, pk_range, ..
            } => {
                assert_eq!(table, "users");
                assert!(pk_range.start_bound.is_some());
                assert!(pk_range.end_bound.is_some());
            }
            _ => panic!("Expected TableRangeScan plan"),
        }
    }

    #[test]
    fn test_table_scan_fallback() {
        let schemas = create_test_schema();
        let planner = QueryPlanner::new(schemas);

        let select = SelectStatement {
            columns: vec![crate::parser::Expression::Column("*".to_string())],
            table: "users".to_string(),
            where_clause: Some(WhereClause {
                condition: Condition::Comparison {
                    left: crate::parser::Expression::Column("name".to_string()),
                    operator: ComparisonOperator::Equal,
                    right: SqlValue::Text("John".to_string()),
                },
            }),
            order_by: None,
            limit: Some(10),
        };

        let plan = planner.plan(Statement::Select(select)).unwrap();

        match plan {
            ExecutionPlan::TableScan { table, limit, .. } => {
                assert_eq!(table, "users");
                assert_eq!(limit, Some(10));
            }
            _ => panic!("Expected TableScan plan"),
        }
    }
}