aegis_query/

planner.rs

//! Aegis Query Planner
//!
//! Transforms analyzed SQL statements into optimized query plans.
//! Implements cost-based optimization with rule-based transformations.
//!
//! Key Features:
//! - Logical plan generation from AST
//! - Cost-based join ordering
//! - Predicate pushdown optimization
//! - Index selection
//!
//! @version 0.1.0
//! @author AutomataNexus Development Team

use crate::ast::{
    BinaryOperator, Expression, FromClause, JoinClause, JoinType, SelectColumn, SelectStatement,
    SetOperationType, Statement, TableReference,
};
use aegis_common::DataType;
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;

// =============================================================================
// Error Types
// =============================================================================

#[derive(Debug, Error)]
pub enum PlannerError {
    #[error("Table not found: {0}")]
    TableNotFound(String),

    #[error("Column not found: {0}")]
    ColumnNotFound(String),

    #[error("Unsupported operation: {0}")]
    UnsupportedOperation(String),

    #[error("Internal planner error: {0}")]
    Internal(String),
}

pub type PlannerResult<T> = Result<T, PlannerError>;

// =============================================================================
// Query Plan
// =============================================================================

/// A complete query plan ready for execution.
#[derive(Debug, Clone)]
pub struct QueryPlan {
    pub root: PlanNode,
    pub estimated_cost: f64,
    pub estimated_rows: u64,
}

/// A node in the query plan tree.
#[derive(Debug, Clone)]
pub enum PlanNode {
    Scan(ScanNode),
    Filter(FilterNode),
    Project(ProjectNode),
    Join(JoinNode),
    Aggregate(AggregateNode),
    Sort(SortNode),
    Limit(LimitNode),
    Empty,
    // DDL nodes
    CreateTable(CreateTableNode),
    DropTable(DropTableNode),
    AlterTable(AlterTableNode),
    CreateIndex(CreateIndexNode),
    DropIndex(DropIndexNode),
    // DML nodes
    Insert(InsertNode),
    Update(UpdateNode),
    Delete(DeleteNode),
    // Set operations
    SetOperation(SetOperationNode),
    // Transaction control
    BeginTransaction,
    CommitTransaction,
    RollbackTransaction,
}

/// Set operation node (UNION, INTERSECT, EXCEPT).
#[derive(Debug, Clone)]
pub struct SetOperationNode {
    pub op: SetOperationType,
    pub left: Box<PlanNode>,
    pub right: Box<PlanNode>,
}

// =============================================================================
// Plan Node Types
// =============================================================================

/// Table scan operation.
#[derive(Debug, Clone)]
pub struct ScanNode {
    pub table_name: String,
    pub alias: Option<String>,
    pub columns: Vec<String>,
    pub index_scan: Option<IndexScan>,
}

/// Index scan specification.
#[derive(Debug, Clone)]
pub struct IndexScan {
    pub index_name: String,
    pub key_range: KeyRange,
}

/// Key range for index scans.
#[derive(Debug, Clone)]
pub struct KeyRange {
    pub start: Option<PlanExpression>,
    pub start_inclusive: bool,
    pub end: Option<PlanExpression>,
    pub end_inclusive: bool,
}

/// Filter operation.
#[derive(Debug, Clone)]
pub struct FilterNode {
    pub input: Box<PlanNode>,
    pub predicate: PlanExpression,
}

/// Projection operation.
#[derive(Debug, Clone)]
pub struct ProjectNode {
    pub input: Box<PlanNode>,
    pub expressions: Vec<ProjectionExpr>,
}

/// A projection expression with optional alias.
#[derive(Debug, Clone)]
pub struct ProjectionExpr {
    pub expr: PlanExpression,
    pub alias: Option<String>,
}

/// Join operation.
#[derive(Debug, Clone)]
pub struct JoinNode {
    pub left: Box<PlanNode>,
    pub right: Box<PlanNode>,
    pub join_type: PlanJoinType,
    pub condition: Option<PlanExpression>,
    pub strategy: JoinStrategy,
}

/// Join type in plan.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanJoinType {
    Inner,
    Left,
    Right,
    Full,
    Cross,
}

/// Physical join strategy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JoinStrategy {
    NestedLoop,
    HashJoin,
    MergeJoin,
}

/// Aggregation operation.
#[derive(Debug, Clone)]
pub struct AggregateNode {
    pub input: Box<PlanNode>,
    pub group_by: Vec<PlanExpression>,
    pub aggregates: Vec<AggregateExpr>,
}

/// Aggregate expression.
#[derive(Debug, Clone)]
pub struct AggregateExpr {
    pub function: AggregateFunction,
    pub argument: Option<PlanExpression>,
    pub distinct: bool,
    pub alias: Option<String>,
}

/// Aggregate functions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggregateFunction {
    Count,
    Sum,
    Avg,
    Min,
    Max,
}

/// Sort operation.
#[derive(Debug, Clone)]
pub struct SortNode {
    pub input: Box<PlanNode>,
    pub order_by: Vec<SortKey>,
}

/// Sort key specification.
#[derive(Debug, Clone)]
pub struct SortKey {
    pub expr: PlanExpression,
    pub ascending: bool,
    pub nulls_first: bool,
}

/// Limit operation.
#[derive(Debug, Clone)]
pub struct LimitNode {
    pub input: Box<PlanNode>,
    pub limit: Option<u64>,
    pub offset: Option<u64>,
}

// =============================================================================
// DDL Plan Nodes
// =============================================================================

/// CREATE TABLE operation.
#[derive(Debug, Clone)]
pub struct CreateTableNode {
    pub table_name: String,
    pub columns: Vec<CreateColumnDef>,
    pub constraints: Vec<CreateTableConstraint>,
    pub if_not_exists: bool,
}

/// Column definition for CREATE TABLE.
#[derive(Debug, Clone)]
pub struct CreateColumnDef {
    pub name: String,
    pub data_type: DataType,
    pub nullable: bool,
    pub default: Option<PlanExpression>,
    pub primary_key: bool,
    pub unique: bool,
}

/// Table constraint for CREATE TABLE.
#[derive(Debug, Clone)]
pub enum CreateTableConstraint {
    PrimaryKey {
        columns: Vec<String>,
    },
    Unique {
        columns: Vec<String>,
    },
    ForeignKey {
        columns: Vec<String>,
        ref_table: String,
        ref_columns: Vec<String>,
    },
    Check {
        expression: PlanExpression,
    },
}

/// DROP TABLE operation.
#[derive(Debug, Clone)]
pub struct DropTableNode {
    pub table_name: String,
    pub if_exists: bool,
}

/// ALTER TABLE operation.
#[derive(Debug, Clone)]
pub struct AlterTableNode {
    pub table_name: String,
    pub operations: Vec<PlanAlterOperation>,
}

/// ALTER TABLE operation type.
#[derive(Debug, Clone)]
pub enum PlanAlterOperation {
    AddColumn(CreateColumnDef),
    DropColumn {
        name: String,
        if_exists: bool,
    },
    RenameColumn {
        old_name: String,
        new_name: String,
    },
    AlterColumn {
        name: String,
        data_type: Option<DataType>,
        set_not_null: Option<bool>,
        set_default: Option<Option<PlanExpression>>,
    },
    RenameTable {
        new_name: String,
    },
    AddConstraint(CreateTableConstraint),
    DropConstraint {
        name: String,
    },
}

/// CREATE INDEX operation.
#[derive(Debug, Clone)]
pub struct CreateIndexNode {
    pub index_name: String,
    pub table_name: String,
    pub columns: Vec<String>,
    pub unique: bool,
    pub if_not_exists: bool,
}

/// DROP INDEX operation.
#[derive(Debug, Clone)]
pub struct DropIndexNode {
    pub index_name: String,
    pub if_exists: bool,
}

// =============================================================================
// DML Plan Nodes
// =============================================================================

/// Source for INSERT data in a query plan.
#[derive(Debug, Clone)]
pub enum InsertPlanSource {
    /// Direct values (INSERT INTO ... VALUES ...)
    Values(Vec<Vec<PlanExpression>>),
    /// Subquery (INSERT INTO ... SELECT ...)
    Query(Box<PlanNode>),
}

/// INSERT operation.
#[derive(Debug, Clone)]
pub struct InsertNode {
    pub table_name: String,
    pub columns: Vec<String>,
    pub source: InsertPlanSource,
}

/// UPDATE operation.
#[derive(Debug, Clone)]
pub struct UpdateNode {
    pub table_name: String,
    pub assignments: Vec<(String, PlanExpression)>,
    pub where_clause: Option<PlanExpression>,
}

/// DELETE operation.
#[derive(Debug, Clone)]
pub struct DeleteNode {
    pub table_name: String,
    pub where_clause: Option<PlanExpression>,
}

// =============================================================================
// Plan Expressions
// =============================================================================

/// Expression in a query plan.
#[derive(Debug, Clone)]
pub enum PlanExpression {
    Column {
        table: Option<String>,
        name: String,
        data_type: DataType,
    },
    Literal(PlanLiteral),
    BinaryOp {
        left: Box<PlanExpression>,
        op: PlanBinaryOp,
        right: Box<PlanExpression>,
    },
    UnaryOp {
        op: PlanUnaryOp,
        expr: Box<PlanExpression>,
    },
    Function {
        name: String,
        args: Vec<PlanExpression>,
        return_type: DataType,
    },
    Cast {
        expr: Box<PlanExpression>,
        target_type: DataType,
    },
    IsNull {
        expr: Box<PlanExpression>,
        negated: bool,
    },
    /// CASE expression
    Case {
        operand: Option<Box<PlanExpression>>,
        conditions: Vec<(PlanExpression, PlanExpression)>,
        else_result: Option<Box<PlanExpression>>,
    },
    /// IN list expression (e.g., x IN (1, 2, 3))
    InList {
        expr: Box<PlanExpression>,
        list: Vec<PlanExpression>,
        negated: bool,
    },
    /// BETWEEN expression (e.g., x BETWEEN 1 AND 10)
    Between {
        expr: Box<PlanExpression>,
        low: Box<PlanExpression>,
        high: Box<PlanExpression>,
        negated: bool,
    },
    /// LIKE pattern matching expression
    Like {
        expr: Box<PlanExpression>,
        pattern: Box<PlanExpression>,
        negated: bool,
    },
    /// IN subquery expression (e.g., x IN (SELECT y FROM t))
    InSubquery {
        expr: Box<PlanExpression>,
        subquery: Box<PlanNode>,
        negated: bool,
    },
    /// EXISTS subquery expression
    Exists {
        subquery: Box<PlanNode>,
        negated: bool,
    },
    /// Scalar subquery (returns single value)
    ScalarSubquery(Box<PlanNode>),
    /// Placeholder for parameterized queries ($1, $2, etc.)
    Placeholder(usize),
}

/// Literal values in plan.
#[derive(Debug, Clone)]
pub enum PlanLiteral {
    Null,
    Boolean(bool),
    Integer(i64),
    Float(f64),
    String(String),
}

/// Binary operators in plan.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanBinaryOp {
    Add,
    Subtract,
    Multiply,
    Divide,
    Modulo,
    Equal,
    NotEqual,
    LessThan,
    LessThanOrEqual,
    GreaterThan,
    GreaterThanOrEqual,
    And,
    Or,
    Concat,
}

/// Unary operators in plan.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanUnaryOp {
    Not,
    Negative,
}

// =============================================================================
// Schema Information
// =============================================================================

/// Schema information for planning.
#[derive(Debug, Clone, Default)]
pub struct PlannerSchema {
    tables: HashMap<String, TableInfo>,
    indexes: HashMap<String, Vec<IndexInfo>>,
}

/// Table metadata for planning.
#[derive(Debug, Clone)]
pub struct TableInfo {
    pub name: String,
    pub columns: Vec<ColumnInfo>,
    pub row_count: u64,
}

/// Column metadata for planning.
#[derive(Debug, Clone)]
pub struct ColumnInfo {
    pub name: String,
    pub data_type: DataType,
    pub nullable: bool,
}

/// Index metadata for planning.
#[derive(Debug, Clone)]
pub struct IndexInfo {
    pub name: String,
    pub columns: Vec<String>,
    pub unique: bool,
}

impl PlannerSchema {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_table(&mut self, table: TableInfo) {
        self.tables.insert(table.name.clone(), table);
    }

    pub fn add_index(&mut self, table: &str, index: IndexInfo) {
        self.indexes
            .entry(table.to_string())
            .or_default()
            .push(index);
    }

    pub fn get_table(&self, name: &str) -> Option<&TableInfo> {
        self.tables.get(name)
    }

    pub fn get_indexes(&self, table: &str) -> Option<&Vec<IndexInfo>> {
        self.indexes.get(table)
    }
}

// =============================================================================
// Planner
// =============================================================================

/// Query planner that transforms AST into execution plans.
pub struct Planner {
    schema: Arc<PlannerSchema>,
}

impl Planner {
    pub fn new(schema: Arc<PlannerSchema>) -> Self {
        Self { schema }
    }

    /// Plan a statement.
    pub fn plan(&self, statement: &Statement) -> PlannerResult<QueryPlan> {
        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::AlterTable(alter) => self.plan_alter_table(alter),
            Statement::CreateIndex(create) => self.plan_create_index(create),
            Statement::DropIndex(drop) => self.plan_drop_index(drop),
            Statement::SetOperation(set_op) => self.plan_set_operation(set_op),
            Statement::Begin => Ok(QueryPlan {
                root: PlanNode::BeginTransaction,
                estimated_cost: 0.0,
                estimated_rows: 0,
            }),
            Statement::Commit => Ok(QueryPlan {
                root: PlanNode::CommitTransaction,
                estimated_cost: 0.0,
                estimated_rows: 0,
            }),
            Statement::Rollback => Ok(QueryPlan {
                root: PlanNode::RollbackTransaction,
                estimated_cost: 0.0,
                estimated_rows: 0,
            }),
        }
    }

    /// Plan a CREATE TABLE statement.
    fn plan_create_table(
        &self,
        create: &crate::ast::CreateTableStatement,
    ) -> PlannerResult<QueryPlan> {
        let mut columns: Vec<CreateColumnDef> = Vec::new();
        for col in &create.columns {
            let primary_key = col
                .constraints
                .iter()
                .any(|c| matches!(c, crate::ast::ColumnConstraint::PrimaryKey));
            let unique = col
                .constraints
                .iter()
                .any(|c| matches!(c, crate::ast::ColumnConstraint::Unique));

            let default = match &col.default {
                Some(e) => Some(self.plan_expression(e)?),
                None => None,
            };

            columns.push(CreateColumnDef {
                name: col.name.clone(),
                data_type: col.data_type.clone(),
                nullable: col.nullable,
                default,
                primary_key,
                unique,
            });
        }

        let constraints: Vec<CreateTableConstraint> = create
            .constraints
            .iter()
            .map(|c| match c {
                crate::ast::TableConstraint::PrimaryKey { columns } => {
                    Ok(CreateTableConstraint::PrimaryKey {
                        columns: columns.clone(),
                    })
                }
                crate::ast::TableConstraint::Unique { columns } => {
                    Ok(CreateTableConstraint::Unique {
                        columns: columns.clone(),
                    })
                }
                crate::ast::TableConstraint::ForeignKey {
                    columns,
                    ref_table,
                    ref_columns,
                } => Ok(CreateTableConstraint::ForeignKey {
                    columns: columns.clone(),
                    ref_table: ref_table.clone(),
                    ref_columns: ref_columns.clone(),
                }),
                crate::ast::TableConstraint::Check { expression } => {
                    Ok(CreateTableConstraint::Check {
                        expression: self.plan_expression(expression)?,
                    })
                }
            })
            .collect::<PlannerResult<Vec<_>>>()?;

        Ok(QueryPlan {
            root: PlanNode::CreateTable(CreateTableNode {
                table_name: create.name.clone(),
                columns,
                constraints,
                if_not_exists: create.if_not_exists,
            }),
            estimated_cost: 1.0,
            estimated_rows: 0,
        })
    }

    /// Plan a DROP TABLE statement.
    fn plan_drop_table(&self, drop: &crate::ast::DropTableStatement) -> PlannerResult<QueryPlan> {
        Ok(QueryPlan {
            root: PlanNode::DropTable(DropTableNode {
                table_name: drop.name.clone(),
                if_exists: drop.if_exists,
            }),
            estimated_cost: 1.0,
            estimated_rows: 0,
        })
    }

    /// Plan an ALTER TABLE statement.
    fn plan_alter_table(
        &self,
        alter: &crate::ast::AlterTableStatement,
    ) -> PlannerResult<QueryPlan> {
        let operations = alter
            .operations
            .iter()
            .map(|op| match op {
                crate::ast::AlterTableOperation::AddColumn(col) => {
                    Ok(PlanAlterOperation::AddColumn(CreateColumnDef {
                        name: col.name.clone(),
                        data_type: col.data_type.clone(),
                        nullable: col.nullable,
                        default: col
                            .default
                            .as_ref()
                            .map(|e| self.plan_expression(e))
                            .transpose()?,
                        primary_key: false,
                        unique: false,
                    }))
                }
                crate::ast::AlterTableOperation::DropColumn { name, if_exists } => {
                    Ok(PlanAlterOperation::DropColumn {
                        name: name.clone(),
                        if_exists: *if_exists,
                    })
                }
                crate::ast::AlterTableOperation::RenameColumn { old_name, new_name } => {
                    Ok(PlanAlterOperation::RenameColumn {
                        old_name: old_name.clone(),
                        new_name: new_name.clone(),
                    })
                }
                crate::ast::AlterTableOperation::AlterColumn {
                    name,
                    data_type,
                    set_not_null,
                    set_default,
                } => {
                    let default_expr = match set_default {
                        Some(Some(expr)) => Some(Some(self.plan_expression(expr)?)),
                        Some(None) => Some(None),
                        None => None,
                    };
                    Ok(PlanAlterOperation::AlterColumn {
                        name: name.clone(),
                        data_type: data_type.clone(),
                        set_not_null: *set_not_null,
                        set_default: default_expr,
                    })
                }
                crate::ast::AlterTableOperation::RenameTable { new_name } => {
                    Ok(PlanAlterOperation::RenameTable {
                        new_name: new_name.clone(),
                    })
                }
                crate::ast::AlterTableOperation::AddConstraint(constraint) => {
                    let plan_constraint = match constraint {
                        crate::ast::TableConstraint::PrimaryKey { columns } => {
                            CreateTableConstraint::PrimaryKey {
                                columns: columns.clone(),
                            }
                        }
                        crate::ast::TableConstraint::Unique { columns } => {
                            CreateTableConstraint::Unique {
                                columns: columns.clone(),
                            }
                        }
                        crate::ast::TableConstraint::ForeignKey {
                            columns,
                            ref_table,
                            ref_columns,
                        } => CreateTableConstraint::ForeignKey {
                            columns: columns.clone(),
                            ref_table: ref_table.clone(),
                            ref_columns: ref_columns.clone(),
                        },
                        crate::ast::TableConstraint::Check { expression } => {
                            CreateTableConstraint::Check {
                                expression: self.plan_expression(expression)?,
                            }
                        }
                    };
                    Ok(PlanAlterOperation::AddConstraint(plan_constraint))
                }
                crate::ast::AlterTableOperation::DropConstraint { name } => {
                    Ok(PlanAlterOperation::DropConstraint { name: name.clone() })
                }
            })
            .collect::<PlannerResult<Vec<_>>>()?;

        Ok(QueryPlan {
            root: PlanNode::AlterTable(AlterTableNode {
                table_name: alter.name.clone(),
                operations,
            }),
            estimated_cost: 1.0,
            estimated_rows: 0,
        })
    }

    /// Plan a CREATE INDEX statement.
    fn plan_create_index(
        &self,
        create: &crate::ast::CreateIndexStatement,
    ) -> PlannerResult<QueryPlan> {
        Ok(QueryPlan {
            root: PlanNode::CreateIndex(CreateIndexNode {
                index_name: create.name.clone(),
                table_name: create.table.clone(),
                columns: create.columns.clone(),
                unique: create.unique,
                if_not_exists: create.if_not_exists,
            }),
            estimated_cost: 1.0,
            estimated_rows: 0,
        })
    }

    /// Plan a DROP INDEX statement.
    fn plan_drop_index(&self, drop: &crate::ast::DropIndexStatement) -> PlannerResult<QueryPlan> {
        Ok(QueryPlan {
            root: PlanNode::DropIndex(DropIndexNode {
                index_name: drop.name.clone(),
                if_exists: drop.if_exists,
            }),
            estimated_cost: 1.0,
            estimated_rows: 0,
        })
    }

    /// Plan an INSERT statement.
    fn plan_insert(&self, insert: &crate::ast::InsertStatement) -> PlannerResult<QueryPlan> {
        let columns = insert.columns.clone().unwrap_or_default();

        let (source, estimated_rows) = match &insert.source {
            crate::ast::InsertSource::Values(rows) => {
                let values = rows
                    .iter()
                    .map(|row| {
                        row.iter()
                            .map(|expr| self.plan_expression(expr))
                            .collect::<PlannerResult<Vec<_>>>()
                    })
                    .collect::<PlannerResult<Vec<_>>>()?;
                let num_rows = values.len() as u64;
                (InsertPlanSource::Values(values), num_rows)
            }
            crate::ast::InsertSource::Query(select) => {
                // Plan the SELECT subquery
                let subquery_plan = self.plan_select(select)?;
                let estimated_rows = subquery_plan.estimated_rows;
                (
                    InsertPlanSource::Query(Box::new(subquery_plan.root)),
                    estimated_rows,
                )
            }
        };

        Ok(QueryPlan {
            root: PlanNode::Insert(InsertNode {
                table_name: insert.table.clone(),
                columns,
                source,
            }),
            estimated_cost: estimated_rows as f64,
            estimated_rows,
        })
    }

    /// Plan an UPDATE statement.
    fn plan_update(&self, update: &crate::ast::UpdateStatement) -> PlannerResult<QueryPlan> {
        let assignments: Vec<(String, PlanExpression)> = update
            .assignments
            .iter()
            .map(|a| Ok((a.column.clone(), self.plan_expression(&a.value)?)))
            .collect::<PlannerResult<Vec<_>>>()?;

        let where_clause = update
            .where_clause
            .as_ref()
            .map(|e| self.plan_expression(e))
            .transpose()?;

        // Estimate rows based on table info
        let estimated_rows = self
            .schema
            .get_table(&update.table)
            .map(|t| t.row_count)
            .unwrap_or(100);

        Ok(QueryPlan {
            root: PlanNode::Update(UpdateNode {
                table_name: update.table.clone(),
                assignments,
                where_clause,
            }),
            estimated_cost: estimated_rows as f64,
            estimated_rows,
        })
    }

    /// Plan a DELETE statement.
    fn plan_delete(&self, delete: &crate::ast::DeleteStatement) -> PlannerResult<QueryPlan> {
        let where_clause = delete
            .where_clause
            .as_ref()
            .map(|e| self.plan_expression(e))
            .transpose()?;

        // Estimate rows based on table info
        let estimated_rows = self
            .schema
            .get_table(&delete.table)
            .map(|t| t.row_count)
            .unwrap_or(100);

        Ok(QueryPlan {
            root: PlanNode::Delete(DeleteNode {
                table_name: delete.table.clone(),
                where_clause,
            }),
            estimated_cost: estimated_rows as f64,
            estimated_rows,
        })
    }

    /// Plan a set operation (UNION, INTERSECT, EXCEPT).
    fn plan_set_operation(
        &self,
        set_op: &crate::ast::SetOperationStatement,
    ) -> PlannerResult<QueryPlan> {
        let left_plan = self.plan(set_op.left.as_ref())?;
        let right_plan = self.plan(set_op.right.as_ref())?;
        let estimated_rows = left_plan.estimated_rows + right_plan.estimated_rows;
        let estimated_cost =
            left_plan.estimated_cost + right_plan.estimated_cost + estimated_rows as f64;

        Ok(QueryPlan {
            root: PlanNode::SetOperation(SetOperationNode {
                op: set_op.op,
                left: Box::new(left_plan.root),
                right: Box::new(right_plan.root),
            }),
            estimated_cost,
            estimated_rows,
        })
    }

    /// Plan a SELECT statement.
    fn plan_select(&self, select: &SelectStatement) -> PlannerResult<QueryPlan> {
        let mut plan = self.plan_from_clause(&select.from)?;

        if let Some(ref where_clause) = select.where_clause {
            let predicate = self.plan_expression(where_clause)?;
            plan = PlanNode::Filter(FilterNode {
                input: Box::new(plan),
                predicate,
            });
        }

        if !select.group_by.is_empty() || self.has_aggregates(&select.columns) {
            plan = self.plan_aggregation(plan, select)?;
        }

        if let Some(ref having) = select.having {
            let predicate = self.plan_expression(having)?;
            plan = PlanNode::Filter(FilterNode {
                input: Box::new(plan),
                predicate,
            });
        }

        plan = self.plan_projection(plan, &select.columns)?;

        if !select.order_by.is_empty() {
            let order_by: Vec<SortKey> = select
                .order_by
                .iter()
                .map(|item| {
                    Ok(SortKey {
                        expr: self.plan_expression(&item.expression)?,
                        ascending: item.ascending,
                        nulls_first: item.nulls_first.unwrap_or(!item.ascending),
                    })
                })
                .collect::<PlannerResult<Vec<_>>>()?;

            plan = PlanNode::Sort(SortNode {
                input: Box::new(plan),
                order_by,
            });
        }

        if select.limit.is_some() || select.offset.is_some() {
            plan = PlanNode::Limit(LimitNode {
                input: Box::new(plan),
                limit: select.limit,
                offset: select.offset,
            });
        }

        let (estimated_cost, estimated_rows) = self.estimate_cost(&plan);

        Ok(QueryPlan {
            root: plan,
            estimated_cost,
            estimated_rows,
        })
    }

    /// Plan FROM clause.
    fn plan_from_clause(&self, from: &Option<FromClause>) -> PlannerResult<PlanNode> {
        let from = match from {
            Some(f) => f,
            None => return Ok(PlanNode::Empty),
        };

        let mut plan = self.plan_table_reference(&from.source)?;

        for join in &from.joins {
            plan = self.plan_join(plan, join)?;
        }

        Ok(plan)
    }

    /// Plan a table reference.
    fn plan_table_reference(&self, table_ref: &TableReference) -> PlannerResult<PlanNode> {
        match table_ref {
            TableReference::Table { name, alias } => {
                let columns = if let Some(info) = self.schema.get_table(name) {
                    info.columns.iter().map(|c| c.name.clone()).collect()
                } else {
                    Vec::new()
                };

                Ok(PlanNode::Scan(ScanNode {
                    table_name: name.clone(),
                    alias: alias.clone(),
                    columns,
                    index_scan: None,
                }))
            }
            TableReference::Subquery { query, alias: _ } => self.plan_select(query).map(|p| p.root),
        }
    }

    /// Plan a join.
    fn plan_join(&self, left: PlanNode, join: &JoinClause) -> PlannerResult<PlanNode> {
        let right = self.plan_table_reference(&join.table)?;
        let condition = join
            .condition
            .as_ref()
            .map(|c| self.plan_expression(c))
            .transpose()?;

        let join_type = match join.join_type {
            JoinType::Inner => PlanJoinType::Inner,
            JoinType::Left => PlanJoinType::Left,
            JoinType::Right => PlanJoinType::Right,
            JoinType::Full => PlanJoinType::Full,
            JoinType::Cross => PlanJoinType::Cross,
        };

        let strategy = self.choose_join_strategy(&left, &right, &condition);

        Ok(PlanNode::Join(JoinNode {
            left: Box::new(left),
            right: Box::new(right),
            join_type,
            condition,
            strategy,
        }))
    }

    /// Choose join strategy based on cost estimation.
    fn choose_join_strategy(
        &self,
        _left: &PlanNode,
        _right: &PlanNode,
        condition: &Option<PlanExpression>,
    ) -> JoinStrategy {
        if condition.is_none() {
            return JoinStrategy::NestedLoop;
        }

        JoinStrategy::HashJoin
    }

    /// Check if select columns contain aggregates.
    fn has_aggregates(&self, columns: &[SelectColumn]) -> bool {
        columns.iter().any(|col| match col {
            SelectColumn::Expression { expr, .. } => self.expression_has_aggregate(expr),
            _ => false,
        })
    }

    /// Check if expression contains aggregate function.
    fn expression_has_aggregate(&self, expr: &Expression) -> bool {
        match expr {
            Expression::Function { name, .. } => {
                matches!(
                    name.to_uppercase().as_str(),
                    "COUNT" | "SUM" | "AVG" | "MIN" | "MAX"
                )
            }
            Expression::BinaryOp { left, right, .. } => {
                self.expression_has_aggregate(left) || self.expression_has_aggregate(right)
            }
            _ => false,
        }
    }

    /// Plan aggregation.
    fn plan_aggregation(
        &self,
        input: PlanNode,
        select: &SelectStatement,
    ) -> PlannerResult<PlanNode> {
        let group_by: Vec<PlanExpression> = select
            .group_by
            .iter()
            .map(|e| self.plan_expression(e))
            .collect::<PlannerResult<Vec<_>>>()?;

        let mut aggregates = Vec::new();
        for col in &select.columns {
            if let SelectColumn::Expression { expr, alias } = col {
                if let Some(agg) = self.extract_aggregate(expr, alias.clone())? {
                    aggregates.push(agg);
                }
            }
        }

        Ok(PlanNode::Aggregate(AggregateNode {
            input: Box::new(input),
            group_by,
            aggregates,
        }))
    }

    /// Extract aggregate from expression.
    fn extract_aggregate(
        &self,
        expr: &Expression,
        alias: Option<String>,
    ) -> PlannerResult<Option<AggregateExpr>> {
        match expr {
            Expression::Function {
                name,
                args,
                distinct,
            } => {
                let func = match name.to_uppercase().as_str() {
                    "COUNT" => AggregateFunction::Count,
                    "SUM" => AggregateFunction::Sum,
                    "AVG" => AggregateFunction::Avg,
                    "MIN" => AggregateFunction::Min,
                    "MAX" => AggregateFunction::Max,
                    _ => return Ok(None),
                };

                let argument = if args.is_empty() {
                    None
                } else {
                    Some(self.plan_expression(&args[0])?)
                };

                Ok(Some(AggregateExpr {
                    function: func,
                    argument,
                    distinct: *distinct,
                    alias,
                }))
            }
            _ => Ok(None),
        }
    }

    /// Plan projection.
    fn plan_projection(
        &self,
        input: PlanNode,
        columns: &[SelectColumn],
    ) -> PlannerResult<PlanNode> {
        let mut expressions = Vec::new();

        for col in columns {
            match col {
                SelectColumn::AllColumns => {
                    expressions.push(ProjectionExpr {
                        expr: PlanExpression::Column {
                            table: None,
                            name: "*".to_string(),
                            data_type: DataType::Any,
                        },
                        alias: None,
                    });
                }
                SelectColumn::TableAllColumns(table) => {
                    expressions.push(ProjectionExpr {
                        expr: PlanExpression::Column {
                            table: Some(table.clone()),
                            name: "*".to_string(),
                            data_type: DataType::Any,
                        },
                        alias: None,
                    });
                }
                SelectColumn::Expression { expr, alias } => {
                    expressions.push(ProjectionExpr {
                        expr: self.plan_expression(expr)?,
                        alias: alias.clone(),
                    });
                }
            }
        }

        Ok(PlanNode::Project(ProjectNode {
            input: Box::new(input),
            expressions,
        }))
    }

    /// Convert AST expression to plan expression.
    fn plan_expression(&self, expr: &Expression) -> PlannerResult<PlanExpression> {
        match expr {
            Expression::Literal(lit) => Ok(PlanExpression::Literal(match lit {
                crate::ast::Literal::Null => PlanLiteral::Null,
                crate::ast::Literal::Boolean(b) => PlanLiteral::Boolean(*b),
                crate::ast::Literal::Integer(i) => PlanLiteral::Integer(*i),
                crate::ast::Literal::Float(f) => PlanLiteral::Float(*f),
                crate::ast::Literal::String(s) => PlanLiteral::String(s.clone()),
            })),

            Expression::Column(col_ref) => Ok(PlanExpression::Column {
                table: col_ref.table.clone(),
                name: col_ref.column.clone(),
                data_type: DataType::Any,
            }),

            Expression::BinaryOp { left, op, right } => {
                let plan_op = match op {
                    BinaryOperator::Add => PlanBinaryOp::Add,
                    BinaryOperator::Subtract => PlanBinaryOp::Subtract,
                    BinaryOperator::Multiply => PlanBinaryOp::Multiply,
                    BinaryOperator::Divide => PlanBinaryOp::Divide,
                    BinaryOperator::Modulo => PlanBinaryOp::Modulo,
                    BinaryOperator::Equal => PlanBinaryOp::Equal,
                    BinaryOperator::NotEqual => PlanBinaryOp::NotEqual,
                    BinaryOperator::LessThan => PlanBinaryOp::LessThan,
                    BinaryOperator::LessThanOrEqual => PlanBinaryOp::LessThanOrEqual,
                    BinaryOperator::GreaterThan => PlanBinaryOp::GreaterThan,
                    BinaryOperator::GreaterThanOrEqual => PlanBinaryOp::GreaterThanOrEqual,
                    BinaryOperator::And => PlanBinaryOp::And,
                    BinaryOperator::Or => PlanBinaryOp::Or,
                    BinaryOperator::Concat => PlanBinaryOp::Concat,
                };

                Ok(PlanExpression::BinaryOp {
                    left: Box::new(self.plan_expression(left)?),
                    op: plan_op,
                    right: Box::new(self.plan_expression(right)?),
                })
            }

            Expression::UnaryOp { op, expr } => {
                let plan_op = match op {
                    crate::ast::UnaryOperator::Not => PlanUnaryOp::Not,
                    crate::ast::UnaryOperator::Negative => PlanUnaryOp::Negative,
                    crate::ast::UnaryOperator::Positive => {
                        return self.plan_expression(expr);
                    }
                };

                Ok(PlanExpression::UnaryOp {
                    op: plan_op,
                    expr: Box::new(self.plan_expression(expr)?),
                })
            }

            Expression::Function { name, args, .. } => {
                let planned_args: Vec<PlanExpression> = args
                    .iter()
                    .map(|a| self.plan_expression(a))
                    .collect::<PlannerResult<Vec<_>>>()?;

                Ok(PlanExpression::Function {
                    name: name.clone(),
                    args: planned_args,
                    return_type: DataType::Any,
                })
            }

            Expression::IsNull { expr, negated } => Ok(PlanExpression::IsNull {
                expr: Box::new(self.plan_expression(expr)?),
                negated: *negated,
            }),

            Expression::Cast { expr, data_type } => Ok(PlanExpression::Cast {
                expr: Box::new(self.plan_expression(expr)?),
                target_type: data_type.clone(),
            }),

            Expression::Case {
                operand,
                conditions,
                else_result,
            } => {
                let planned_operand = operand
                    .as_ref()
                    .map(|e| self.plan_expression(e))
                    .transpose()?
                    .map(Box::new);

                let planned_conditions = conditions
                    .iter()
                    .map(|(cond, result)| {
                        Ok((self.plan_expression(cond)?, self.plan_expression(result)?))
                    })
                    .collect::<PlannerResult<Vec<_>>>()?;

                let planned_else = else_result
                    .as_ref()
                    .map(|e| self.plan_expression(e))
                    .transpose()?
                    .map(Box::new);

                Ok(PlanExpression::Case {
                    operand: planned_operand,
                    conditions: planned_conditions,
                    else_result: planned_else,
                })
            }

            Expression::InList {
                expr,
                list,
                negated,
            } => {
                let planned_list = list
                    .iter()
                    .map(|e| self.plan_expression(e))
                    .collect::<PlannerResult<Vec<_>>>()?;

                Ok(PlanExpression::InList {
                    expr: Box::new(self.plan_expression(expr)?),
                    list: planned_list,
                    negated: *negated,
                })
            }

            Expression::Between {
                expr,
                low,
                high,
                negated,
            } => Ok(PlanExpression::Between {
                expr: Box::new(self.plan_expression(expr)?),
                low: Box::new(self.plan_expression(low)?),
                high: Box::new(self.plan_expression(high)?),
                negated: *negated,
            }),

            Expression::Like {
                expr,
                pattern,
                negated,
            } => Ok(PlanExpression::Like {
                expr: Box::new(self.plan_expression(expr)?),
                pattern: Box::new(self.plan_expression(pattern)?),
                negated: *negated,
            }),

            Expression::InSubquery {
                expr,
                subquery,
                negated,
            } => {
                let subquery_plan = self.plan_select(subquery)?;
                Ok(PlanExpression::InSubquery {
                    expr: Box::new(self.plan_expression(expr)?),
                    subquery: Box::new(subquery_plan.root),
                    negated: *negated,
                })
            }

            Expression::Exists { subquery, negated } => {
                let subquery_plan = self.plan_select(subquery)?;
                Ok(PlanExpression::Exists {
                    subquery: Box::new(subquery_plan.root),
                    negated: *negated,
                })
            }

            Expression::Subquery(subquery) => {
                let subquery_plan = self.plan_select(subquery)?;
                Ok(PlanExpression::ScalarSubquery(Box::new(subquery_plan.root)))
            }

            Expression::Placeholder(idx) => Ok(PlanExpression::Placeholder(*idx)),
        }
    }

    /// Estimate cost and row count for a plan.
    fn estimate_cost(&self, node: &PlanNode) -> (f64, u64) {
        match node {
            PlanNode::Scan(scan) => {
                let rows = self
                    .schema
                    .get_table(&scan.table_name)
                    .map(|t| t.row_count)
                    .unwrap_or(1000);
                let cost = rows as f64 * 1.0;
                (cost, rows)
            }

            PlanNode::Filter(filter) => {
                let (input_cost, input_rows) = self.estimate_cost(&filter.input);
                let selectivity = 0.1;
                let rows = (input_rows as f64 * selectivity) as u64;
                let cost = input_cost + input_rows as f64 * 0.1;
                (cost, rows.max(1))
            }

            PlanNode::Project(project) => {
                let (input_cost, input_rows) = self.estimate_cost(&project.input);
                let cost = input_cost + input_rows as f64 * 0.01;
                (cost, input_rows)
            }

            PlanNode::Join(join) => {
                let (left_cost, left_rows) = self.estimate_cost(&join.left);
                let (right_cost, right_rows) = self.estimate_cost(&join.right);

                let rows = match join.join_type {
                    PlanJoinType::Cross => left_rows * right_rows,
                    _ => (left_rows + right_rows) / 2,
                };

                let join_cost = match join.strategy {
                    JoinStrategy::NestedLoop => left_rows as f64 * right_rows as f64,
                    JoinStrategy::HashJoin => {
                        left_rows as f64 + right_rows as f64 + right_rows as f64 * 1.2
                    }
                    JoinStrategy::MergeJoin => {
                        left_rows as f64 * 2.0_f64.log2()
                            + right_rows as f64 * 2.0_f64.log2()
                            + (left_rows + right_rows) as f64
                    }
                };

                (left_cost + right_cost + join_cost, rows)
            }

            PlanNode::Aggregate(agg) => {
                let (input_cost, input_rows) = self.estimate_cost(&agg.input);
                let groups = if agg.group_by.is_empty() {
                    1
                } else {
                    (input_rows as f64).sqrt() as u64
                };
                let cost = input_cost + input_rows as f64 * 0.5;
                (cost, groups.max(1))
            }

            PlanNode::Sort(sort) => {
                let (input_cost, input_rows) = self.estimate_cost(&sort.input);
                let sort_cost = input_rows as f64 * (input_rows as f64).log2();
                (input_cost + sort_cost, input_rows)
            }

            PlanNode::Limit(limit) => {
                let (input_cost, input_rows) = self.estimate_cost(&limit.input);
                let rows = limit.limit.unwrap_or(input_rows).min(input_rows);
                (input_cost, rows)
            }

            PlanNode::Empty => (0.0, 1),

            // DDL operations have minimal cost
            PlanNode::CreateTable(_) => (1.0, 0),
            PlanNode::DropTable(_) => (1.0, 0),
            PlanNode::AlterTable(_) => (1.0, 0),
            PlanNode::CreateIndex(_) => (1.0, 0),
            PlanNode::DropIndex(_) => (1.0, 0),

            // DML operations
            PlanNode::Insert(insert) => {
                let rows = match &insert.source {
                    InsertPlanSource::Values(values) => values.len() as u64,
                    InsertPlanSource::Query(subquery) => {
                        let (_, estimated_rows) = self.estimate_cost(subquery);
                        estimated_rows
                    }
                };
                (rows as f64, rows)
            }
            PlanNode::Update(update) => {
                let rows = self
                    .schema
                    .get_table(&update.table_name)
                    .map(|t| t.row_count)
                    .unwrap_or(100);
                (rows as f64, rows)
            }
            PlanNode::Delete(delete) => {
                let rows = self
                    .schema
                    .get_table(&delete.table_name)
                    .map(|t| t.row_count)
                    .unwrap_or(100);
                (rows as f64, rows)
            }

            // Set operations
            PlanNode::SetOperation(set_op) => {
                let (left_cost, left_rows) = self.estimate_cost(&set_op.left);
                let (right_cost, right_rows) = self.estimate_cost(&set_op.right);
                let rows = left_rows + right_rows;
                let cost = left_cost + right_cost + rows as f64;
                (cost, rows)
            }

            // Transaction control
            PlanNode::BeginTransaction
            | PlanNode::CommitTransaction
            | PlanNode::RollbackTransaction => (0.0, 0),
        }
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{ColumnRef, FromClause, Literal, OrderByItem, TableReference};

    fn create_test_schema() -> Arc<PlannerSchema> {
        let mut schema = PlannerSchema::new();

        schema.add_table(TableInfo {
            name: "users".to_string(),
            columns: vec![
                ColumnInfo {
                    name: "id".to_string(),
                    data_type: DataType::Integer,
                    nullable: false,
                },
                ColumnInfo {
                    name: "name".to_string(),
                    data_type: DataType::Text,
                    nullable: false,
                },
                ColumnInfo {
                    name: "email".to_string(),
                    data_type: DataType::Text,
                    nullable: true,
                },
            ],
            row_count: 10000,
        });

        schema.add_table(TableInfo {
            name: "orders".to_string(),
            columns: vec![
                ColumnInfo {
                    name: "id".to_string(),
                    data_type: DataType::Integer,
                    nullable: false,
                },
                ColumnInfo {
                    name: "user_id".to_string(),
                    data_type: DataType::Integer,
                    nullable: false,
                },
                ColumnInfo {
                    name: "amount".to_string(),
                    data_type: DataType::Float,
                    nullable: false,
                },
            ],
            row_count: 50000,
        });

        Arc::new(schema)
    }

    #[test]
    fn test_plan_simple_select() {
        let schema = create_test_schema();
        let planner = Planner::new(schema);

        let select = SelectStatement {
            distinct: false,
            columns: vec![SelectColumn::AllColumns],
            from: Some(FromClause {
                source: TableReference::Table {
                    name: "users".to_string(),
                    alias: None,
                },
                joins: vec![],
            }),
            where_clause: None,
            group_by: vec![],
            having: None,
            order_by: vec![],
            limit: None,
            offset: None,
        };

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

        assert!(plan.estimated_rows > 0);
        assert!(plan.estimated_cost > 0.0);
    }

    #[test]
    fn test_plan_select_with_filter() {
        let schema = create_test_schema();
        let planner = Planner::new(schema);

        let select = SelectStatement {
            distinct: false,
            columns: vec![SelectColumn::AllColumns],
            from: Some(FromClause {
                source: TableReference::Table {
                    name: "users".to_string(),
                    alias: None,
                },
                joins: vec![],
            }),
            where_clause: Some(Expression::BinaryOp {
                left: Box::new(Expression::Column(ColumnRef {
                    table: None,
                    column: "id".to_string(),
                })),
                op: BinaryOperator::Equal,
                right: Box::new(Expression::Literal(Literal::Integer(1))),
            }),
            group_by: vec![],
            having: None,
            order_by: vec![],
            limit: None,
            offset: None,
        };

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

        match &plan.root {
            PlanNode::Project(p) => match p.input.as_ref() {
                PlanNode::Filter(_) => {}
                _ => panic!("Expected filter node"),
            },
            _ => panic!("Expected project node"),
        }
    }

    #[test]
    fn test_plan_select_with_order_by() {
        let schema = create_test_schema();
        let planner = Planner::new(schema);

        let select = SelectStatement {
            distinct: false,
            columns: vec![SelectColumn::AllColumns],
            from: Some(FromClause {
                source: TableReference::Table {
                    name: "users".to_string(),
                    alias: None,
                },
                joins: vec![],
            }),
            where_clause: None,
            group_by: vec![],
            having: None,
            order_by: vec![OrderByItem {
                expression: Expression::Column(ColumnRef {
                    table: None,
                    column: "name".to_string(),
                }),
                ascending: true,
                nulls_first: None,
            }],
            limit: Some(10),
            offset: None,
        };

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

        match &plan.root {
            PlanNode::Limit(l) => match l.input.as_ref() {
                PlanNode::Sort(_) => {}
                _ => panic!("Expected sort node"),
            },
            _ => panic!("Expected limit node"),
        }
    }

    #[test]
    fn test_plan_set_operation_union() {
        let schema = create_test_schema();
        let planner = Planner::new(schema);

        let set_op = crate::ast::SetOperationStatement {
            op: SetOperationType::Union,
            left: Box::new(Statement::Select(SelectStatement {
                columns: vec![SelectColumn::AllColumns],
                from: Some(FromClause {
                    source: TableReference::Table {
                        name: "users".to_string(),
                        alias: None,
                    },
                    joins: vec![],
                }),
                ..Default::default()
            })),
            right: Box::new(Statement::Select(SelectStatement {
                columns: vec![SelectColumn::AllColumns],
                from: Some(FromClause {
                    source: TableReference::Table {
                        name: "orders".to_string(),
                        alias: None,
                    },
                    joins: vec![],
                }),
                ..Default::default()
            })),
        };

        let plan = planner.plan(&Statement::SetOperation(set_op)).unwrap();

        match &plan.root {
            PlanNode::SetOperation(node) => {
                assert_eq!(node.op, SetOperationType::Union);
            }
            _ => panic!("Expected SetOperation node"),
        }
        assert!(plan.estimated_rows > 0);
        assert!(plan.estimated_cost > 0.0);
    }
}