kimberlite-query 0.4.2

//! Query planner: transforms parsed SQL into execution plans.
//!
//! The planner analyzes predicates to select the optimal access path:
//! - `PointLookup`: When all primary key columns have equality predicates
//! - `RangeScan`: When primary key has range predicates
//! - `TableScan`: Fallback for non-indexed predicates

use std::ops::Bound;

use crate::error::{QueryError, Result};
use crate::key_encoder::{encode_key, successor_key};
use crate::parser::{
    CaseWhenArm, ComputedColumn, OrderByClause, ParsedSelect, Predicate, PredicateValue,
};
use crate::plan::{
    CaseColumnDef, CaseWhenClause, Filter, FilterCondition, FilterOp, QueryPlan, ScanOrder,
    SortSpec,
};
use crate::schema::{ColumnName, Schema, TableDef};
use crate::value::Value;

/// Creates table metadata from a table definition.
#[inline]
fn create_metadata(table_def: &TableDef, table_name: String) -> crate::plan::TableMetadata {
    crate::plan::TableMetadata {
        table_id: table_def.table_id,
        table_name,
        columns: table_def.columns.clone(),
        primary_key: table_def.primary_key.clone(),
    }
}

/// Builds a point lookup plan.
#[inline]
fn build_point_lookup_plan(
    table_def: &TableDef,
    table_name: String,
    key_values: &[Value],
    column_indices: Vec<usize>,
    column_names: Vec<ColumnName>,
) -> QueryPlan {
    let key = encode_key(key_values);
    QueryPlan::PointLookup {
        metadata: create_metadata(table_def, table_name),
        key,
        columns: column_indices,
        column_names,
    }
}

/// Builds a range scan plan.
#[inline]
#[allow(clippy::too_many_arguments)]
fn build_range_scan_plan(
    table_def: &TableDef,
    table_name: String,
    start_key: Bound<kimberlite_store::Key>,
    end_key: Bound<kimberlite_store::Key>,
    remaining_predicates: &[ResolvedPredicate],
    limit: Option<usize>,
    order_by: &[OrderByClause],
    column_indices: Vec<usize>,
    column_names: Vec<ColumnName>,
) -> Result<QueryPlan> {
    let filter = build_filter(table_def, remaining_predicates, &table_name)?;
    let order = determine_scan_order(order_by, table_def);

    // Build sort spec for client-side sorting when ORDER BY doesn't match scan order
    let needs_client_sort = !order_by.is_empty()
        && order_by
            .iter()
            .any(|clause| !table_def.is_primary_key(&clause.column));
    let order_by_spec = if needs_client_sort {
        build_sort_spec(order_by, table_def, &table_name)?
    } else {
        None
    };

    Ok(QueryPlan::RangeScan {
        metadata: create_metadata(table_def, table_name),
        start: start_key,
        end: end_key,
        filter,
        limit,
        order,
        order_by: order_by_spec,
        columns: column_indices,
        column_names,
    })
}

/// Builds an index scan plan.
#[inline]
#[allow(clippy::too_many_arguments)]
fn build_index_scan_plan(
    table_def: &TableDef,
    table_name: String,
    index_id: u64,
    index_name: String,
    start_key: Bound<kimberlite_store::Key>,
    end_key: Bound<kimberlite_store::Key>,
    remaining_predicates: &[ResolvedPredicate],
    limit: Option<usize>,
    order_by: &[OrderByClause],
    column_indices: Vec<usize>,
    column_names: Vec<ColumnName>,
) -> Result<QueryPlan> {
    let filter = build_filter(table_def, remaining_predicates, &table_name)?;
    let order = determine_scan_order(order_by, table_def);

    // Build sort spec for client-side sorting when ORDER BY doesn't match scan order
    let needs_client_sort = !order_by.is_empty()
        && order_by
            .iter()
            .any(|clause| !table_def.is_primary_key(&clause.column));
    let order_by_spec = if needs_client_sort {
        build_sort_spec(order_by, table_def, &table_name)?
    } else {
        None
    };

    Ok(QueryPlan::IndexScan {
        metadata: create_metadata(table_def, table_name),
        index_id,
        index_name,
        start: start_key,
        end: end_key,
        filter,
        limit,
        order,
        order_by: order_by_spec,
        columns: column_indices,
        column_names,
    })
}

/// Builds a table scan plan.
#[inline]
fn build_table_scan_plan(
    table_def: &TableDef,
    table_name: String,
    all_predicates: &[ResolvedPredicate],
    limit: Option<usize>,
    order_by: &[OrderByClause],
    column_indices: Vec<usize>,
    column_names: Vec<ColumnName>,
) -> Result<QueryPlan> {
    let filter = build_filter(table_def, all_predicates, &table_name)?;
    let order = build_sort_spec(order_by, table_def, &table_name)?;

    Ok(QueryPlan::TableScan {
        metadata: create_metadata(table_def, table_name),
        filter,
        limit,
        order,
        columns: column_indices,
        column_names,
    })
}

/// Wraps a base plan with an aggregate plan if needed.
#[inline]
fn wrap_with_aggregate(
    base_plan: QueryPlan,
    table_def: &TableDef,
    table_name: String,
    parsed: &ParsedSelect,
) -> Result<QueryPlan> {
    // For DISTINCT without explicit GROUP BY, group by all selected columns
    let group_by_columns = if parsed.distinct && parsed.group_by.is_empty() {
        parsed
            .columns
            .clone()
            .unwrap_or_else(|| table_def.columns.iter().map(|c| c.name.clone()).collect())
    } else {
        parsed.group_by.clone()
    };

    // For DISTINCT, aggregates should be empty (just deduplication)
    let aggregates = if parsed.distinct && parsed.aggregates.is_empty() {
        vec![]
    } else {
        parsed.aggregates.clone()
    };

    // Resolve GROUP BY column indices
    let mut group_by_indices = Vec::new();
    for col_name in &group_by_columns {
        let (idx, _) =
            table_def
                .find_column(col_name)
                .ok_or_else(|| QueryError::ColumnNotFound {
                    table: table_name.clone(),
                    column: col_name.to_string(),
                })?;
        group_by_indices.push(idx);
    }

    // Build result column names: GROUP BY columns + aggregate results
    let mut result_columns = group_by_columns.clone();
    for agg in &aggregates {
        let agg_name = match agg {
            crate::parser::AggregateFunction::CountStar => "COUNT(*)".to_string(),
            crate::parser::AggregateFunction::Count(col) => format!("COUNT({col})"),
            crate::parser::AggregateFunction::Sum(col) => format!("SUM({col})"),
            crate::parser::AggregateFunction::Avg(col) => format!("AVG({col})"),
            crate::parser::AggregateFunction::Min(col) => format!("MIN({col})"),
            crate::parser::AggregateFunction::Max(col) => format!("MAX({col})"),
        };
        result_columns.push(ColumnName::new(agg_name));
    }

    Ok(QueryPlan::Aggregate {
        metadata: create_metadata(table_def, table_name),
        source: Box::new(base_plan),
        group_by_cols: group_by_indices,
        group_by_names: group_by_columns,
        aggregates,
        column_names: result_columns,
        having: parsed.having.clone(),
    })
}

/// Plans a parsed SELECT statement.
pub fn plan_query(schema: &Schema, parsed: &ParsedSelect, params: &[Value]) -> Result<QueryPlan> {
    if parsed.joins.is_empty() {
        // Single-table query - existing logic
        plan_single_table_query(schema, parsed, params)
    } else {
        // Multi-table query - new JOIN logic
        plan_join_query(schema, parsed, params)
    }
}

/// Plans a single-table query (no JOINs).
fn plan_single_table_query(
    schema: &Schema,
    parsed: &ParsedSelect,
    params: &[Value],
) -> Result<QueryPlan> {
    // Look up table
    let table_name = parsed.table.clone();
    let table_def = schema
        .get_table(&table_name.clone().into())
        .ok_or_else(|| QueryError::TableNotFound(table_name.clone()))?;

    // Resolve predicate values (substitute parameters)
    let resolved_predicates = resolve_predicates(&parsed.predicates, params)?;

    // Resolve columns for the query (accounts for aggregate requirements)
    let (column_indices, column_names) = resolve_query_columns(table_def, parsed, &table_name)?;

    // Check if we need aggregates
    let needs_aggregate = !parsed.aggregates.is_empty()
        || !parsed.group_by.is_empty()
        || parsed.distinct
        || !parsed.having.is_empty();

    // Analyze predicates to determine access path
    let access_path = analyze_access_path(table_def, &resolved_predicates);

    // Build the base scan plan
    let base_plan = build_scan_plan(
        access_path,
        table_def,
        table_name.clone(),
        parsed,
        column_indices,
        column_names,
    )?;

    // Wrap in an aggregate plan if needed
    let plan_after_agg = if needs_aggregate {
        wrap_with_aggregate(base_plan, table_def, table_name.clone(), parsed)?
    } else {
        base_plan
    };

    // Wrap in a Materialize plan if CASE WHEN computed columns are present
    if !parsed.case_columns.is_empty() {
        let existing_columns = plan_after_agg.column_names().to_vec();
        let case_columns = resolve_case_columns_for_join(&parsed.case_columns, &existing_columns)?;
        let mut output_columns = existing_columns;
        output_columns.extend(case_columns.iter().map(|c| c.alias.clone()));

        Ok(QueryPlan::Materialize {
            source: Box::new(plan_after_agg),
            filter: None,
            case_columns,
            order: None,
            limit: None,
            column_names: output_columns,
        })
    } else {
        Ok(plan_after_agg)
    }
}

/// Plans a multi-table query with JOINs.
fn plan_join_query(schema: &Schema, parsed: &ParsedSelect, params: &[Value]) -> Result<QueryPlan> {
    // Build left-deep join tree
    let mut current_plan = plan_table_access(schema, &parsed.table, params)?;

    for join in &parsed.joins {
        // Plan right table access
        let right_plan = plan_table_access(schema, &join.table, params)?;

        // Build join conditions from ON clause
        let on_conditions =
            build_join_conditions(&join.on_condition, schema, &parsed.table, &join.table)?;

        // Merge column names from both tables (left then right)
        let left_columns = current_plan.column_names().to_vec();
        let right_columns = right_plan.column_names().to_vec();
        let mut all_columns = left_columns.clone();
        all_columns.extend(right_columns);

        // Build join node
        current_plan = QueryPlan::Join {
            join_type: join.join_type.clone(),
            left: Box::new(current_plan),
            right: Box::new(right_plan),
            on_conditions,
            columns: vec![], // All columns from both tables
            column_names: all_columns,
        };
    }

    // Collect the full combined column list from the join tree
    let combined_columns = current_plan.column_names().to_vec();

    // Resolve WHERE predicates against the combined column list
    let resolved_predicates = resolve_predicates(&parsed.predicates, params)?;
    let filter = build_filter_for_join(&resolved_predicates, &combined_columns)?;

    // Resolve ORDER BY against the combined column list
    let order = build_sort_spec_for_join(&parsed.order_by, &combined_columns)?;

    // Resolve CASE WHEN computed columns
    let case_columns = resolve_case_columns_for_join(&parsed.case_columns, &combined_columns)?;

    // Determine output column names: selected columns from combined list + computed aliases
    let output_columns: Vec<ColumnName> = match &parsed.columns {
        None => {
            // SELECT * — all combined columns + computed aliases
            let mut out = combined_columns.clone();
            out.extend(case_columns.iter().map(|c| c.alias.clone()));
            out
        }
        Some(selected) => {
            // Validate that each selected column exists in the combined list
            for col in selected {
                if !combined_columns.iter().any(|c| c == col) {
                    return Err(QueryError::ColumnNotFound {
                        table: parsed.table.clone(),
                        column: col.to_string(),
                    });
                }
            }
            let mut out = selected.clone();
            out.extend(case_columns.iter().map(|c| c.alias.clone()));
            out
        }
    };

    let needs_materialize =
        filter.is_some() || order.is_some() || parsed.limit.is_some() || !case_columns.is_empty();

    if needs_materialize {
        Ok(QueryPlan::Materialize {
            source: Box::new(current_plan),
            filter,
            case_columns,
            order,
            limit: parsed.limit,
            column_names: output_columns,
        })
    } else {
        Ok(current_plan)
    }
}

/// Builds a filter from resolved predicates against a named column list.
///
/// Used for JOIN queries where we don't have a single `TableDef` — instead we
/// resolve column names against the combined left+right column list.
fn build_filter_for_join(
    predicates: &[ResolvedPredicate],
    columns: &[ColumnName],
) -> Result<Option<Filter>> {
    if predicates.is_empty() {
        return Ok(None);
    }

    let filters: Result<Vec<_>> = predicates
        .iter()
        .map(|p| build_filter_for_join_predicate(p, columns))
        .collect();

    Ok(Some(Filter::and(filters?)))
}

fn build_filter_for_join_predicate(
    pred: &ResolvedPredicate,
    columns: &[ColumnName],
) -> Result<Filter> {
    if let ResolvedOp::Or(left_preds, right_preds) = &pred.op {
        let left_filter = build_filter_for_join(left_preds, columns)?.ok_or_else(|| {
            QueryError::UnsupportedFeature("OR left side has no predicates".to_string())
        })?;
        let right_filter = build_filter_for_join(right_preds, columns)?.ok_or_else(|| {
            QueryError::UnsupportedFeature("OR right side has no predicates".to_string())
        })?;
        Ok(Filter::or(vec![left_filter, right_filter]))
    } else {
        let condition = build_filter_condition_for_join(pred, columns)?;
        Ok(Filter::single(condition))
    }
}

fn build_filter_condition_for_join(
    pred: &ResolvedPredicate,
    columns: &[ColumnName],
) -> Result<FilterCondition> {
    let col_idx = columns
        .iter()
        .position(|c| c == &pred.column)
        .ok_or_else(|| QueryError::ColumnNotFound {
            table: "(join)".to_string(),
            column: pred.column.to_string(),
        })?;

    let (op, value) = match &pred.op {
        ResolvedOp::Eq(v) => (FilterOp::Eq, v.clone()),
        ResolvedOp::Lt(v) => (FilterOp::Lt, v.clone()),
        ResolvedOp::Le(v) => (FilterOp::Le, v.clone()),
        ResolvedOp::Gt(v) => (FilterOp::Gt, v.clone()),
        ResolvedOp::Ge(v) => (FilterOp::Ge, v.clone()),
        ResolvedOp::In(vals) => (FilterOp::In(vals.clone()), Value::Null),
        ResolvedOp::Like(pattern) => (FilterOp::Like(pattern.clone()), Value::Null),
        ResolvedOp::IsNull => (FilterOp::IsNull, Value::Null),
        ResolvedOp::IsNotNull => (FilterOp::IsNotNull, Value::Null),
        ResolvedOp::Or(_, _) => {
            return Err(QueryError::UnsupportedFeature(
                "OR predicates must be handled at filter level".to_string(),
            ));
        }
    };

    Ok(FilterCondition {
        column_idx: col_idx,
        op,
        value,
    })
}

/// Builds a sort specification from ORDER BY clauses against a named column list.
fn build_sort_spec_for_join(
    order_by: &[OrderByClause],
    columns: &[ColumnName],
) -> Result<Option<SortSpec>> {
    if order_by.is_empty() {
        return Ok(None);
    }

    let mut sort_cols = Vec::with_capacity(order_by.len());
    for clause in order_by {
        let idx = columns
            .iter()
            .position(|c| c == &clause.column)
            .ok_or_else(|| QueryError::ColumnNotFound {
                table: "(join)".to_string(),
                column: clause.column.to_string(),
            })?;
        let order = if clause.ascending {
            ScanOrder::Ascending
        } else {
            ScanOrder::Descending
        };
        sort_cols.push((idx, order));
    }

    Ok(Some(SortSpec { columns: sort_cols }))
}

/// Resolves CASE WHEN computed columns against a named column list.
fn resolve_case_columns_for_join(
    case_columns: &[ComputedColumn],
    columns: &[ColumnName],
) -> Result<Vec<CaseColumnDef>> {
    case_columns
        .iter()
        .map(|cc| resolve_single_case_column(cc, columns))
        .collect()
}

fn resolve_single_case_column(
    cc: &ComputedColumn,
    columns: &[ColumnName],
) -> Result<CaseColumnDef> {
    let when_clauses: Result<Vec<_>> = cc
        .when_clauses
        .iter()
        .map(|arm| resolve_case_when_arm(arm, columns))
        .collect();

    Ok(CaseColumnDef {
        alias: cc.alias.clone(),
        when_clauses: when_clauses?,
        else_value: cc.else_value.clone(),
    })
}

fn resolve_case_when_arm(arm: &CaseWhenArm, columns: &[ColumnName]) -> Result<CaseWhenClause> {
    // Resolve predicates (no params in CASE conditions)
    let resolved = resolve_predicates(&arm.condition, &[])?;
    let filter = build_filter_for_join(&resolved, columns)?.ok_or_else(|| {
        QueryError::UnsupportedFeature("CASE WHEN condition has no predicates".to_string())
    })?;

    Ok(CaseWhenClause {
        condition: filter,
        result: arm.result.clone(),
    })
}

/// Plans a table access for a single table (used in JOINs).
fn plan_table_access(schema: &Schema, table_name: &str, _params: &[Value]) -> Result<QueryPlan> {
    let table_def = schema
        .get_table(&table_name.into())
        .ok_or_else(|| QueryError::TableNotFound(table_name.to_string()))?;

    // For JOIN table access, just do a full table scan
    // (In the future, we could optimize this based on join predicates)
    let all_column_indices: Vec<usize> = (0..table_def.columns.len()).collect();
    let all_column_names: Vec<ColumnName> =
        table_def.columns.iter().map(|c| c.name.clone()).collect();

    Ok(QueryPlan::TableScan {
        metadata: create_metadata(table_def, table_name.to_string()),
        filter: None,
        limit: None,
        order: None,
        columns: all_column_indices,
        column_names: all_column_names,
    })
}

/// Builds join conditions from ON clause predicates.
///
/// JOIN predicates can have column references on both sides (e.g., users.id = orders.user_id).
/// These need to be resolved to indices in the concatenated row [left_cols..., right_cols...].
fn build_join_conditions(
    predicates: &[Predicate],
    schema: &Schema,
    left_table: &str,
    right_table: &str,
) -> Result<Vec<crate::plan::JoinCondition>> {
    let left_table_def = schema
        .get_table(&left_table.into())
        .ok_or_else(|| QueryError::TableNotFound(left_table.to_string()))?;
    let right_table_def = schema
        .get_table(&right_table.into())
        .ok_or_else(|| QueryError::TableNotFound(right_table.to_string()))?;

    let left_col_count = left_table_def.columns.len();

    predicates
        .iter()
        .map(|pred| {
            build_single_join_condition(
                pred,
                left_table,
                left_table_def,
                right_table,
                right_table_def,
                left_col_count,
            )
        })
        .collect()
}

/// Builds a single join condition from a JOIN predicate.
///
/// Handles column-to-column comparisons by resolving qualified/unqualified column names
/// to indices in the concatenated row [left_cols..., right_cols...].
fn build_single_join_condition(
    pred: &Predicate,
    left_table: &str,
    left_table_def: &TableDef,
    right_table: &str,
    right_table_def: &TableDef,
    left_col_count: usize,
) -> Result<crate::plan::JoinCondition> {
    use crate::plan::JoinOp;

    // Extract column name, operator, and right side from predicate
    let (left_col_name, op, right_value) = match pred {
        Predicate::Eq(col, val) => (col, JoinOp::Eq, val),
        Predicate::Lt(col, val) => (col, JoinOp::Lt, val),
        Predicate::Le(col, val) => (col, JoinOp::Le, val),
        Predicate::Gt(col, val) => (col, JoinOp::Gt, val),
        Predicate::Ge(col, val) => (col, JoinOp::Ge, val),
        _ => {
            return Err(QueryError::UnsupportedFeature(
                "only equality and comparison operators supported in JOIN ON clause".to_string(),
            ));
        }
    };

    // Resolve left column to index in concatenated row
    let left_col_idx = resolve_join_column(left_col_name, left_table, left_table_def, 0)?;

    // Right side must be a column reference for JOIN conditions
    match right_value {
        PredicateValue::ColumnRef(ref_str) => {
            // Column-to-column comparison
            let right_col_idx = resolve_join_column_ref(
                ref_str,
                left_table,
                left_table_def,
                right_table,
                right_table_def,
                left_col_count,
            )?;

            Ok(crate::plan::JoinCondition {
                left_col_idx,
                right_col_idx,
                op,
            })
        }
        _ => Err(QueryError::UnsupportedFeature(
            "JOIN ON clause requires column-to-column comparisons (e.g., users.id = orders.user_id)".to_string(),
        )),
    }
}

/// Resolves a column name to an index in the concatenated row.
fn resolve_join_column(
    col_name: &ColumnName,
    table_name: &str,
    table_def: &TableDef,
    offset: usize,
) -> Result<usize> {
    let (idx, _) = table_def
        .find_column(col_name)
        .ok_or_else(|| QueryError::ColumnNotFound {
            table: table_name.to_string(),
            column: col_name.to_string(),
        })?;
    Ok(offset + idx)
}

/// Resolves a qualified/unqualified column reference to an index in the concatenated row.
fn resolve_join_column_ref(
    ref_str: &str,
    left_table: &str,
    left_table_def: &TableDef,
    right_table: &str,
    right_table_def: &TableDef,
    left_col_count: usize,
) -> Result<usize> {
    // Parse qualified reference: "table.column" or just "column"
    if let Some((table, column)) = ref_str.split_once('.') {
        // Qualified: table.column
        if table == left_table {
            resolve_join_column(&column.into(), left_table, left_table_def, 0)
        } else if table == right_table {
            resolve_join_column(&column.into(), right_table, right_table_def, left_col_count)
        } else {
            Err(QueryError::TableNotFound(table.to_string()))
        }
    } else {
        // Unqualified: just "column" - try right table first (common pattern)
        if let Ok(idx) = resolve_join_column(
            &ref_str.into(),
            right_table,
            right_table_def,
            left_col_count,
        ) {
            Ok(idx)
        } else {
            resolve_join_column(&ref_str.into(), left_table, left_table_def, 0)
        }
    }
}

/// Builds a scan plan from the analyzed access path.
#[inline]
fn build_scan_plan(
    access_path: AccessPath,
    table_def: &TableDef,
    table_name: String,
    parsed: &ParsedSelect,
    column_indices: Vec<usize>,
    column_names: Vec<ColumnName>,
) -> Result<QueryPlan> {
    match access_path {
        AccessPath::PointLookup { key_values } => Ok(build_point_lookup_plan(
            table_def,
            table_name,
            &key_values,
            column_indices,
            column_names,
        )),
        AccessPath::RangeScan {
            start_key,
            end_key,
            remaining_predicates,
        } => build_range_scan_plan(
            table_def,
            table_name,
            start_key,
            end_key,
            &remaining_predicates,
            parsed.limit,
            &parsed.order_by,
            column_indices,
            column_names,
        ),
        AccessPath::IndexScan {
            index_id,
            index_name,
            start_key,
            end_key,
            remaining_predicates,
        } => build_index_scan_plan(
            table_def,
            table_name,
            index_id,
            index_name,
            start_key,
            end_key,
            &remaining_predicates,
            parsed.limit,
            &parsed.order_by,
            column_indices,
            column_names,
        ),
        AccessPath::TableScan {
            predicates: all_predicates,
        } => build_table_scan_plan(
            table_def,
            table_name,
            &all_predicates,
            parsed.limit,
            &parsed.order_by,
            column_indices,
            column_names,
        ),
    }
}

/// Resolves column selection for queries, accounting for aggregate requirements.
#[inline]
fn resolve_query_columns(
    table_def: &TableDef,
    parsed: &ParsedSelect,
    table_name: &str,
) -> Result<(Vec<usize>, Vec<ColumnName>)> {
    let needs_aggregate = !parsed.aggregates.is_empty()
        || !parsed.group_by.is_empty()
        || parsed.distinct
        || !parsed.having.is_empty();

    // For aggregate queries, the source plan must fetch ALL columns
    // so the executor can access columns by table-level indices
    if needs_aggregate {
        resolve_columns(table_def, None, table_name)
    } else {
        resolve_columns(table_def, parsed.columns.as_ref(), table_name)
    }
}

/// Resolves column selection to indices and names.
fn resolve_columns(
    table_def: &TableDef,
    columns: Option<&Vec<ColumnName>>,
    table_name: &str,
) -> Result<(Vec<usize>, Vec<ColumnName>)> {
    match columns {
        None => {
            // SELECT * - return all columns
            let indices: Vec<usize> = (0..table_def.columns.len()).collect();
            let names: Vec<ColumnName> = table_def.columns.iter().map(|c| c.name.clone()).collect();
            Ok((indices, names))
        }
        Some(cols) => {
            let mut indices = Vec::with_capacity(cols.len());
            let mut names = Vec::with_capacity(cols.len());

            for col in cols {
                let (idx, col_def) =
                    table_def
                        .find_column(col)
                        .ok_or_else(|| QueryError::ColumnNotFound {
                            table: table_name.to_string(),
                            column: col.to_string(),
                        })?;
                indices.push(idx);
                names.push(col_def.name.clone());
            }

            Ok((indices, names))
        }
    }
}

/// Resolved predicate with concrete values (parameters substituted).
#[derive(Debug, Clone)]
struct ResolvedPredicate {
    column: ColumnName,
    op: ResolvedOp,
}

#[derive(Debug, Clone)]
enum ResolvedOp {
    Eq(Value),
    Lt(Value),
    Le(Value),
    Gt(Value),
    Ge(Value),
    In(Vec<Value>),
    Like(String),
    IsNull,
    IsNotNull,
    Or(Vec<ResolvedPredicate>, Vec<ResolvedPredicate>),
}

/// Resolves predicates by substituting parameter values.
fn resolve_predicates(
    predicates: &[Predicate],
    params: &[Value],
) -> Result<Vec<ResolvedPredicate>> {
    predicates
        .iter()
        .map(|p| resolve_predicate(p, params))
        .collect()
}

fn resolve_predicate(predicate: &Predicate, params: &[Value]) -> Result<ResolvedPredicate> {
    match predicate {
        Predicate::Eq(col, val) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Eq(resolve_value(val, params)?),
        }),
        Predicate::Lt(col, val) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Lt(resolve_value(val, params)?),
        }),
        Predicate::Le(col, val) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Le(resolve_value(val, params)?),
        }),
        Predicate::Gt(col, val) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Gt(resolve_value(val, params)?),
        }),
        Predicate::Ge(col, val) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Ge(resolve_value(val, params)?),
        }),
        Predicate::In(col, vals) => {
            let resolved: Result<Vec<_>> = vals.iter().map(|v| resolve_value(v, params)).collect();
            Ok(ResolvedPredicate {
                column: col.clone(),
                op: ResolvedOp::In(resolved?),
            })
        }
        Predicate::Like(col, pattern) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::Like(pattern.clone()),
        }),
        Predicate::IsNull(col) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::IsNull,
        }),
        Predicate::IsNotNull(col) => Ok(ResolvedPredicate {
            column: col.clone(),
            op: ResolvedOp::IsNotNull,
        }),
        Predicate::Or(left_preds, right_preds) => {
            // For OR, we use a dummy column (empty string) since OR can span multiple columns
            let left_resolved = resolve_predicates(left_preds, params)?;
            let right_resolved = resolve_predicates(right_preds, params)?;
            Ok(ResolvedPredicate {
                column: ColumnName::new(String::new()),
                op: ResolvedOp::Or(left_resolved, right_resolved),
            })
        }
    }
}

fn resolve_value(val: &PredicateValue, params: &[Value]) -> Result<Value> {
    match val {
        PredicateValue::Int(v) => Ok(Value::BigInt(*v)),
        PredicateValue::String(s) => Ok(Value::Text(s.clone())),
        PredicateValue::Bool(b) => Ok(Value::Boolean(*b)),
        PredicateValue::Null => Ok(Value::Null),
        PredicateValue::Literal(v) => Ok(v.clone()),
        PredicateValue::Param(idx) => {
            // Parameters are 1-indexed in SQL
            let zero_idx = idx.checked_sub(1).ok_or(QueryError::ParameterNotFound(0))?;
            params
                .get(zero_idx)
                .cloned()
                .ok_or(QueryError::ParameterNotFound(*idx))
        }
        PredicateValue::ColumnRef(_) => Err(QueryError::UnsupportedFeature(
            "column references in WHERE clause not supported (use JOIN ON for column-to-column comparisons)".to_string(),
        )),
    }
}

/// Access path determined by predicate analysis.
enum AccessPath {
    /// Point lookup on primary key.
    PointLookup { key_values: Vec<Value> },
    /// Range scan on primary key.
    RangeScan {
        start_key: Bound<kimberlite_store::Key>,
        end_key: Bound<kimberlite_store::Key>,
        remaining_predicates: Vec<ResolvedPredicate>,
    },
    /// Index scan on a secondary index.
    IndexScan {
        index_id: u64,
        index_name: String,
        start_key: Bound<kimberlite_store::Key>,
        end_key: Bound<kimberlite_store::Key>,
        remaining_predicates: Vec<ResolvedPredicate>,
    },
    /// Full table scan.
    TableScan { predicates: Vec<ResolvedPredicate> },
}

/// Analyzes predicates to determine the optimal access path.
fn analyze_access_path(table_def: &TableDef, predicates: &[ResolvedPredicate]) -> AccessPath {
    let pk_columns = &table_def.primary_key;

    if pk_columns.is_empty() {
        // No primary key - must do table scan
        return AccessPath::TableScan {
            predicates: predicates.to_vec(),
        };
    }

    // Check for point lookup: all PK columns have equality predicates
    let mut pk_values: Vec<Option<Value>> = vec![None; pk_columns.len()];
    let mut non_pk_predicates = Vec::new();

    for pred in predicates {
        if let Some(pk_pos) = table_def.primary_key_position(&pred.column) {
            if let ResolvedOp::Eq(val) = &pred.op {
                pk_values[pk_pos] = Some(val.clone());
                continue;
            }
        }
        non_pk_predicates.push(pred.clone());
    }

    // Check if we have all PK columns with equality
    if pk_values.iter().all(Option::is_some) {
        let key_values: Vec<Value> = pk_values.into_iter().flatten().collect();
        return AccessPath::PointLookup { key_values };
    }

    // Check for range scan: first PK column(s) have predicates
    // For simplicity, only handle single-column PK range scans for now
    if pk_columns.len() == 1 {
        let pk_col = &pk_columns[0];
        let pk_predicates: Vec<_> = predicates.iter().filter(|p| &p.column == pk_col).collect();

        if !pk_predicates.is_empty() {
            let bounds_result = compute_range_bounds(&pk_predicates);

            // If we have useful bounds (not both unbounded), use range scan
            let has_bounds = !matches!(
                (&bounds_result.start, &bounds_result.end),
                (Bound::Unbounded, Bound::Unbounded)
            );

            if has_bounds {
                // Collect remaining predicates: non-PK predicates + unconverted PK predicates
                let mut remaining: Vec<_> = predicates
                    .iter()
                    .filter(|p| &p.column != pk_col)
                    .cloned()
                    .collect();
                remaining.extend(bounds_result.unconverted);

                return AccessPath::RangeScan {
                    start_key: bounds_result.start,
                    end_key: bounds_result.end,
                    remaining_predicates: remaining,
                };
            }
            // If no useful bounds (e.g., only IN predicates), fall through to index scan check
        }
    }

    // Check for index scan: find indexes that can be used
    let index_candidates = find_usable_indexes(table_def, predicates);
    if let Some((best_index, start, end, remaining)) = select_best_index(&index_candidates) {
        return AccessPath::IndexScan {
            index_id: best_index.index_id,
            index_name: best_index.name.clone(),
            start_key: start,
            end_key: end,
            remaining_predicates: remaining,
        };
    }

    // Fall back to table scan
    AccessPath::TableScan {
        predicates: predicates.to_vec(),
    }
}

/// Result of computing range bounds from predicates.
struct RangeBoundsResult {
    start: Bound<kimberlite_store::Key>,
    end: Bound<kimberlite_store::Key>,
    /// Predicates that couldn't be converted to bounds (e.g., IN).
    unconverted: Vec<ResolvedPredicate>,
}

/// Computes range bounds from predicates on a single column.
fn compute_range_bounds(predicates: &[&ResolvedPredicate]) -> RangeBoundsResult {
    let mut lower: Option<(Value, bool)> = None; // (value, inclusive)
    let mut upper: Option<(Value, bool)> = None;
    let mut unconverted = Vec::new();

    for pred in predicates {
        match &pred.op {
            ResolvedOp::Eq(val) => {
                // Exact match - both bounds are this value
                lower = Some((val.clone(), true));
                upper = Some((val.clone(), true));
            }
            ResolvedOp::Gt(val) => {
                lower = Some((val.clone(), false));
            }
            ResolvedOp::Ge(val) => {
                lower = Some((val.clone(), true));
            }
            ResolvedOp::Lt(val) => {
                upper = Some((val.clone(), false));
            }
            ResolvedOp::Le(val) => {
                upper = Some((val.clone(), true));
            }
            ResolvedOp::In(_)
            | ResolvedOp::Like(_)
            | ResolvedOp::IsNull
            | ResolvedOp::IsNotNull
            | ResolvedOp::Or(_, _) => {
                // These can't be converted to range bounds - add to filter
                unconverted.push((*pred).clone());
            }
        }
    }

    let start = match lower {
        Some((val, true)) => Bound::Included(encode_key(&[val])),
        Some((val, false)) => Bound::Excluded(encode_key(&[val])),
        None => Bound::Unbounded,
    };

    let end = match upper {
        Some((val, true)) => {
            // For inclusive upper bound, we need the successor key
            Bound::Excluded(successor_key(&encode_key(&[val])))
        }
        Some((val, false)) => Bound::Excluded(encode_key(&[val])),
        None => Bound::Unbounded,
    };

    RangeBoundsResult {
        start,
        end,
        unconverted,
    }
}

/// Candidate index with its bounds and remaining predicates.
struct IndexCandidate<'a> {
    index_def: &'a crate::schema::IndexDef,
    start: Bound<kimberlite_store::Key>,
    end: Bound<kimberlite_store::Key>,
    remaining: Vec<ResolvedPredicate>,
    score: usize,
}

/// Finds indexes that can be used for the given predicates.
///
/// Returns a list of index candidates with their computed bounds.
/// Only includes indexes where the first column has predicates.
fn find_usable_indexes<'a>(
    table_def: &'a TableDef,
    predicates: &[ResolvedPredicate],
) -> Vec<IndexCandidate<'a>> {
    let mut candidates = Vec::new();
    let max_iterations = 100; // Bounded iteration limit

    for (iter_count, index_def) in table_def.indexes().iter().enumerate() {
        // Bounded iteration check
        if iter_count >= max_iterations {
            break;
        }

        // Skip empty indexes
        if index_def.columns.is_empty() {
            continue;
        }

        // Check if first column has predicates
        let first_col = &index_def.columns[0];
        let first_col_predicates: Vec<_> = predicates
            .iter()
            .filter(|p| &p.column == first_col)
            .collect();

        if first_col_predicates.is_empty() {
            continue;
        }

        // Compute range bounds for this index
        let bounds_result = compute_range_bounds(&first_col_predicates);

        // Skip if both bounds are unbounded (no useful range)
        if matches!(
            (&bounds_result.start, &bounds_result.end),
            (Bound::Unbounded, Bound::Unbounded)
        ) {
            continue;
        }

        // Collect remaining predicates (non-index predicates + unconverted)
        let mut remaining: Vec<_> = predicates
            .iter()
            .filter(|p| !index_def.columns.contains(&p.column))
            .cloned()
            .collect();
        remaining.extend(bounds_result.unconverted);

        // Score this index
        let score = score_index(index_def, predicates);

        candidates.push(IndexCandidate {
            index_def,
            start: bounds_result.start,
            end: bounds_result.end,
            remaining,
            score,
        });
    }

    candidates
}

/// Scores an index based on predicate coverage.
///
/// Returns higher scores for indexes that cover more predicates with better match types.
fn score_index(index_def: &crate::schema::IndexDef, predicates: &[ResolvedPredicate]) -> usize {
    let mut score = 0;
    let max_columns = 10; // Bounded iteration limit

    for (iter_count, index_col) in index_def.columns.iter().enumerate() {
        // Bounded iteration check
        if iter_count >= max_columns {
            break;
        }

        for pred in predicates {
            if &pred.column == index_col {
                match &pred.op {
                    ResolvedOp::Eq(_) => score += 10, // Equality predicates are best
                    ResolvedOp::Lt(_)
                    | ResolvedOp::Le(_)
                    | ResolvedOp::Gt(_)
                    | ResolvedOp::Ge(_) => {
                        score += 5; // Range predicates are good
                    }
                    _ => score += 1, // Other predicates have minor benefit
                }
            }
        }
    }

    score
}

/// Return type for index selection with index definition, key bounds, and remaining predicates.
type BestIndexResult<'a> = (
    &'a crate::schema::IndexDef,
    Bound<kimberlite_store::Key>,
    Bound<kimberlite_store::Key>,
    Vec<ResolvedPredicate>,
);

/// Selects the best index from candidates.
///
/// Returns the index with the highest score, breaking ties by fewest remaining predicates.
fn select_best_index<'a>(candidates: &'a [IndexCandidate<'a>]) -> Option<BestIndexResult<'a>> {
    if candidates.is_empty() {
        return None;
    }

    // Find the maximum score
    let max_score = candidates.iter().map(|c| c.score).max().unwrap_or(0);

    // Filter to candidates with max score
    let best_candidates: Vec<_> = candidates.iter().filter(|c| c.score == max_score).collect();

    // Among ties, select the one with fewest remaining predicates
    let best = best_candidates
        .iter()
        .min_by_key(|c| (c.remaining.len(), c.index_def.columns.len()))?;

    Some((
        best.index_def,
        best.start.clone(),
        best.end.clone(),
        best.remaining.clone(),
    ))
}

/// Builds a filter from remaining predicates.
fn build_filter(
    table_def: &TableDef,
    predicates: &[ResolvedPredicate],
    table_name: &str,
) -> Result<Option<Filter>> {
    if predicates.is_empty() {
        return Ok(None);
    }

    let filters: Result<Vec<_>> = predicates
        .iter()
        .map(|p| build_filter_from_predicate(table_def, p, table_name))
        .collect();

    Ok(Some(Filter::and(filters?)))
}

/// Builds a filter from a single resolved predicate.
/// Handles OR predicates recursively.
fn build_filter_from_predicate(
    table_def: &TableDef,
    pred: &ResolvedPredicate,
    table_name: &str,
) -> Result<Filter> {
    if let ResolvedOp::Or(left_preds, right_preds) = &pred.op {
        // Recursively build filters for left and right sides
        let left_filter = build_filter(table_def, left_preds, table_name)?.ok_or_else(|| {
            QueryError::UnsupportedFeature("OR left side has no predicates".to_string())
        })?;
        let right_filter = build_filter(table_def, right_preds, table_name)?.ok_or_else(|| {
            QueryError::UnsupportedFeature("OR right side has no predicates".to_string())
        })?;

        Ok(Filter::or(vec![left_filter, right_filter]))
    } else {
        // For non-OR predicates, build a FilterCondition
        let condition = build_filter_condition(table_def, pred, table_name)?;
        Ok(Filter::single(condition))
    }
}

fn build_filter_condition(
    table_def: &TableDef,
    pred: &ResolvedPredicate,
    table_name: &str,
) -> Result<FilterCondition> {
    let (col_idx, _) =
        table_def
            .find_column(&pred.column)
            .ok_or_else(|| QueryError::ColumnNotFound {
                table: table_name.to_string(),
                column: pred.column.to_string(),
            })?;

    let (op, value) = match &pred.op {
        ResolvedOp::Eq(v) => (FilterOp::Eq, v.clone()),
        ResolvedOp::Lt(v) => (FilterOp::Lt, v.clone()),
        ResolvedOp::Le(v) => (FilterOp::Le, v.clone()),
        ResolvedOp::Gt(v) => (FilterOp::Gt, v.clone()),
        ResolvedOp::Ge(v) => (FilterOp::Ge, v.clone()),
        ResolvedOp::In(vals) => (FilterOp::In(vals.clone()), Value::Null), // Value unused for In
        ResolvedOp::Like(pattern) => (FilterOp::Like(pattern.clone()), Value::Null),
        ResolvedOp::IsNull => (FilterOp::IsNull, Value::Null),
        ResolvedOp::IsNotNull => (FilterOp::IsNotNull, Value::Null),
        ResolvedOp::Or(_, _) => {
            // OR predicates need special handling - they can't be represented as a single FilterCondition
            return Err(QueryError::UnsupportedFeature(
                "OR predicates must be handled at filter level, not as individual conditions"
                    .to_string(),
            ));
        }
    };

    Ok(FilterCondition {
        column_idx: col_idx,
        op,
        value,
    })
}

/// Determines scan order from ORDER BY for range scans.
fn determine_scan_order(order_by: &[OrderByClause], table_def: &TableDef) -> ScanOrder {
    if order_by.is_empty() {
        return ScanOrder::Ascending;
    }

    // Check if first ORDER BY column is in the primary key
    let first = &order_by[0];
    if table_def.is_primary_key(&first.column) {
        if first.ascending {
            ScanOrder::Ascending
        } else {
            ScanOrder::Descending
        }
    } else {
        ScanOrder::Ascending
    }
}

/// Builds a sort specification for table scans.
fn build_sort_spec(
    order_by: &[OrderByClause],
    table_def: &TableDef,
    table_name: &str,
) -> Result<Option<SortSpec>> {
    if order_by.is_empty() {
        return Ok(None);
    }

    let mut columns = Vec::with_capacity(order_by.len());

    for clause in order_by {
        let (col_idx, _) =
            table_def
                .find_column(&clause.column)
                .ok_or_else(|| QueryError::ColumnNotFound {
                    table: table_name.to_string(),
                    column: clause.column.to_string(),
                })?;

        let order = if clause.ascending {
            ScanOrder::Ascending
        } else {
            ScanOrder::Descending
        };

        columns.push((col_idx, order));
    }

    Ok(Some(SortSpec { columns }))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::{ParsedStatement, parse_statement};
    use crate::schema::{ColumnDef, DataType, SchemaBuilder};
    use kimberlite_store::TableId;

    fn parse_test_select(sql: &str) -> ParsedSelect {
        match parse_statement(sql).unwrap() {
            ParsedStatement::Select(s) => s,
            _ => panic!("expected SELECT statement"),
        }
    }

    fn test_schema() -> Schema {
        SchemaBuilder::new()
            .table(
                "users",
                TableId::new(1),
                vec![
                    ColumnDef::new("id", DataType::BigInt).not_null(),
                    ColumnDef::new("name", DataType::Text).not_null(),
                    ColumnDef::new("age", DataType::BigInt),
                ],
                vec!["id".into()],
            )
            .build()
    }

    #[test]
    fn test_plan_point_lookup() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM users WHERE id = 42");
        let plan = plan_query(&schema, &parsed, &[]).unwrap();

        assert!(matches!(plan, QueryPlan::PointLookup { .. }));
    }

    #[test]
    fn test_plan_range_scan() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM users WHERE id > 10");
        let plan = plan_query(&schema, &parsed, &[]).unwrap();

        assert!(matches!(plan, QueryPlan::RangeScan { .. }));
    }

    #[test]
    fn test_plan_table_scan() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM users WHERE name = 'alice'");
        let plan = plan_query(&schema, &parsed, &[]).unwrap();

        assert!(matches!(plan, QueryPlan::TableScan { .. }));
    }

    #[test]
    fn test_plan_with_params() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM users WHERE id = $1");
        let plan = plan_query(&schema, &parsed, &[Value::BigInt(42)]).unwrap();

        assert!(matches!(plan, QueryPlan::PointLookup { .. }));
    }

    #[test]
    fn test_plan_missing_param() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM users WHERE id = $1");
        let result = plan_query(&schema, &parsed, &[]);

        assert!(matches!(result, Err(QueryError::ParameterNotFound(1))));
    }

    #[test]
    fn test_plan_unknown_table() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT * FROM unknown");
        let result = plan_query(&schema, &parsed, &[]);

        assert!(matches!(result, Err(QueryError::TableNotFound(_))));
    }

    #[test]
    fn test_plan_unknown_column() {
        let schema = test_schema();
        let parsed = parse_test_select("SELECT unknown FROM users");
        let result = plan_query(&schema, &parsed, &[]);

        assert!(matches!(result, Err(QueryError::ColumnNotFound { .. })));
    }
}