aedb 0.3.1

use crate::catalog::Catalog;
use crate::catalog::namespace_key;
use crate::catalog::schema::IndexType;
use crate::catalog::types::Value;
use crate::query::error::QueryError;
use crate::storage::encoded_key::EncodedKey;
use std::collections::HashMap;

use super::composite_index::{CompositeSelectionCriteria, composite_prefix_index_lookup};
use super::index_diagnostics::{
    merge_selected_indexes_if_diagnostic, merge_trace_if_diagnostic,
    merge_trace_single_if_diagnostic, plan_trace_if_diagnostic, selected_indexes_if_diagnostic,
};
use super::index_lookup::{IndexLookup, extract_indexable_predicate};
use super::index_utils::{expr_implied_by_eq_constraints, intersect_pks, union_pks};
use super::predicate::collect_eq_constraints;

pub(super) struct IndexLookupContext<'a> {
    pub(super) catalog: &'a Catalog,
    pub(super) project_id: &'a str,
    pub(super) scope_id: &'a str,
    pub(super) table_name: &'a str,
    pub(super) table: &'a crate::storage::keyspace::TableData,
    pub(super) segment_store: Option<&'a crate::storage::kv_segment::KvSegmentStore>,
    pub(super) include_diagnostics: bool,
}

#[derive(Debug, Clone)]
pub(super) struct IndexLookupResult {
    pub pks: Vec<EncodedKey>,
    pub selected_indexes: Vec<String>,
    pub plan_trace: Vec<String>,
    pub predicate_exact: bool,
    /// Whether every candidate in `pks` is the key of a live row. True for
    /// index-backed lookups (postings are maintained transactionally with
    /// rows), false for probe lists built from caller-supplied values (a
    /// primary-key `IN` may name rows that never existed or were deleted).
    /// Candidate-count shortcuts — truncating to the page window, draining
    /// OFFSET keys before hydration — are only sound when this holds;
    /// otherwise a candidate that fails to hydrate would silently stand in
    /// for a real row that was cut off.
    pub pks_are_live_rows: bool,
}

#[allow(clippy::too_many_arguments)]
pub(super) fn indexed_pks_for_predicate_limited(
    catalog: &Catalog,
    project_id: &str,
    scope_id: &str,
    table_name: &str,
    table: &crate::storage::keyspace::TableData,
    segment_store: Option<&crate::storage::kv_segment::KvSegmentStore>,
    predicate: &crate::query::plan::Expr,
    candidate_limit: Option<usize>,
) -> Result<Option<IndexLookupResult>, QueryError> {
    let context = IndexLookupContext {
        catalog,
        project_id,
        scope_id,
        table_name,
        table,
        segment_store,
        include_diagnostics: false,
    };
    indexed_pks_for_predicate_inner(&context, predicate, candidate_limit)
}

pub(super) fn indexed_pks_for_predicate_with_trace(
    catalog: &Catalog,
    project_id: &str,
    scope_id: &str,
    table_name: &str,
    table: &crate::storage::keyspace::TableData,
    segment_store: Option<&crate::storage::kv_segment::KvSegmentStore>,
    predicate: &crate::query::plan::Expr,
) -> Result<Option<IndexLookupResult>, QueryError> {
    let context = IndexLookupContext {
        catalog,
        project_id,
        scope_id,
        table_name,
        table,
        segment_store,
        include_diagnostics: true,
    };
    indexed_pks_for_predicate_inner(&context, predicate, None)
}

fn indexed_pks_for_predicate_inner(
    context: &IndexLookupContext<'_>,
    predicate: &crate::query::plan::Expr,
    candidate_limit: Option<usize>,
) -> Result<Option<IndexLookupResult>, QueryError> {
    let mut result = indexed_pks_for_predicate_uncapped(context, predicate, candidate_limit)?;
    if let Some(result) = result.as_mut()
        && result.predicate_exact
        && result.pks_are_live_rows
        && let Some(limit) = candidate_limit
    {
        result.pks.truncate(limit);
    }
    Ok(result)
}

fn indexed_pks_for_predicate_uncapped(
    context: &IndexLookupContext<'_>,
    predicate: &crate::query::plan::Expr,
    candidate_limit: Option<usize>,
) -> Result<Option<IndexLookupResult>, QueryError> {
    use crate::query::plan::Expr;

    match predicate {
        Expr::And(lhs, rhs) => {
            let left = indexed_pks_for_predicate_uncapped(context, lhs, None)?;
            let right = indexed_pks_for_predicate_uncapped(context, rhs, None)?;
            return Ok(match (left, right) {
                (Some(left), Some(right)) => Some(IndexLookupResult {
                    // An intersection member appears on both sides, so one
                    // verified side is enough to vouch for it.
                    pks_are_live_rows: left.pks_are_live_rows || right.pks_are_live_rows,
                    pks: intersect_pks(left.pks, right.pks),
                    selected_indexes: merge_selected_indexes_if_diagnostic(
                        context.include_diagnostics,
                        left.selected_indexes,
                        right.selected_indexes,
                    ),
                    predicate_exact: left.predicate_exact && right.predicate_exact,
                    plan_trace: merge_trace_if_diagnostic(
                        context.include_diagnostics,
                        "AND predicate combines indexed candidates with intersection",
                        left.plan_trace,
                        right.plan_trace,
                    ),
                }),
                (Some(left), None) => Some(IndexLookupResult {
                    plan_trace: merge_trace_single_if_diagnostic(
                        context.include_diagnostics,
                        "AND predicate uses indexed left side; right side will be residual filter",
                        left.plan_trace,
                    ),
                    predicate_exact: false,
                    ..left
                }),
                (None, Some(right)) => Some(IndexLookupResult {
                    plan_trace: merge_trace_single_if_diagnostic(
                        context.include_diagnostics,
                        "AND predicate uses indexed right side; left side will be residual filter",
                        right.plan_trace,
                    ),
                    predicate_exact: false,
                    ..right
                }),
                (None, None) => None,
            });
        }
        Expr::Or(lhs, rhs) => {
            let left = indexed_pks_for_predicate_uncapped(context, lhs, None)?;
            let right = indexed_pks_for_predicate_uncapped(context, rhs, None)?;
            return Ok(match (left, right) {
                (Some(left), Some(right)) => Some(IndexLookupResult {
                    // A union member may come from either side, so both must
                    // be verified.
                    pks_are_live_rows: left.pks_are_live_rows && right.pks_are_live_rows,
                    pks: union_pks(left.pks, right.pks),
                    selected_indexes: merge_selected_indexes_if_diagnostic(
                        context.include_diagnostics,
                        left.selected_indexes,
                        right.selected_indexes,
                    ),
                    predicate_exact: left.predicate_exact && right.predicate_exact,
                    plan_trace: merge_trace_if_diagnostic(
                        context.include_diagnostics,
                        "OR predicate combines indexed candidates with union",
                        left.plan_trace,
                        right.plan_trace,
                    ),
                }),
                _ => None,
            });
        }
        _ => {}
    }

    let Some(lookup) = extract_indexable_predicate(predicate) else {
        let mut equalities = HashMap::new();
        let eq_only = collect_eq_constraints(predicate, &mut equalities);
        if !eq_only {
            return Ok(None);
        }
        return composite_prefix_index_lookup(
            context,
            &equalities,
            CompositeSelectionCriteria::default(),
            candidate_limit,
        );
    };
    let column = match &lookup {
        IndexLookup::Range { column, .. }
        | IndexLookup::Eq { column, .. }
        | IndexLookup::In { column, .. } => *column,
    };

    // Range/Like/Between lookups require an ordered index. Hash/UniqueHash stores
    // return no rows for range scans (see SecondaryIndex::scan_range_window_ordered),
    // so selecting one for a range predicate would silently drop matching rows.
    // Skip them here and fall back to a residual full-scan filter instead.
    let lookup_requires_ordered_index = matches!(lookup, IndexLookup::Range { .. });
    let selected_index_name = select_single_column_index(
        context.catalog,
        context.table,
        context.project_id,
        context.scope_id,
        context.table_name,
        predicate,
        column,
        lookup_requires_ordered_index,
    );

    let Some(index_name) = selected_index_name else {
        if let IndexLookup::Eq { column, value, .. } = lookup
            && let Some(result) = indexed_pks_for_leftmost_composite_eq(
                context,
                column,
                value,
                predicate,
                candidate_limit,
            )?
        {
            return Ok(Some(result));
        }
        if let IndexLookup::In { column, values, .. } = &lookup
            && let Some(result) = primary_key_in_probe(context, column, values)
        {
            return Ok(Some(result));
        }
        return Ok(None);
    };
    let Some(index) = context.table.indexes.get(index_name) else {
        return Ok(None);
    };
    let trace_column = context.include_diagnostics.then(|| column.to_string());

    let predicate_exact = lookup.predicate_exact();
    let lookup_limit = if predicate_exact {
        candidate_limit
    } else {
        None
    };
    let store = context.segment_store;
    let pks = match lookup {
        IndexLookup::Range { bounds, .. } => index.tier_scan_range_limit(
            bounds.0,
            bounds.1,
            lookup_limit.unwrap_or(usize::MAX),
            store,
        )?,
        IndexLookup::Eq { value, .. } => index.tier_scan_eq_limit(
            &EncodedKey::from_values(std::slice::from_ref(value)),
            lookup_limit.unwrap_or(usize::MAX),
            store,
        )?,
        IndexLookup::In { values, .. } => {
            let limit = lookup_limit.unwrap_or(usize::MAX);
            let mut pks = Vec::new();
            for value in values {
                let remaining = limit.saturating_sub(pks.len());
                if remaining == 0 {
                    break;
                }
                pks.extend(index.tier_scan_eq_limit(
                    &EncodedKey::from_values(std::slice::from_ref(value)),
                    remaining,
                    store,
                )?);
            }
            pks
        }
    };
    Ok(Some(IndexLookupResult {
        pks,
        selected_indexes: selected_indexes_if_diagnostic(context.include_diagnostics, index_name),
        predicate_exact,
        pks_are_live_rows: true,
        plan_trace: plan_trace_if_diagnostic(context.include_diagnostics, || {
            let column = trace_column.as_deref().unwrap_or("");
            format!("selected single-column index '{index_name}' for predicate on '{column}'")
        }),
    }))
}

/// Whether `value` is exactly of `col_type`, so its encoded form can stand in
/// for an equality comparison against a column of that type. Float is
/// deliberately excluded: `compare_values` and encoded-byte equality disagree
/// on NaN and signed zero.
fn value_is_exactly_column_type(
    value: &Value,
    col_type: &crate::catalog::types::ColumnType,
) -> bool {
    use crate::catalog::types::ColumnType;
    matches!(
        (value, col_type),
        (Value::Text(_), ColumnType::Text)
            | (Value::U8(_), ColumnType::U8)
            | (Value::U64(_), ColumnType::U64)
            | (Value::Integer(_), ColumnType::Integer)
            | (Value::Boolean(_), ColumnType::Boolean)
            | (Value::U256(_), ColumnType::U256)
            | (Value::I256(_), ColumnType::I256)
            | (Value::Blob(_), ColumnType::Blob)
            | (Value::Timestamp(_), ColumnType::Timestamp)
            | (Value::Json(_), ColumnType::Json)
    )
}

/// `IN` over a single-column primary key resolves to direct PK point probes
/// when no secondary index covers the column, instead of a full scan with a
/// residual filter. Probes are encoded, sorted, and deduped so the output
/// order (ascending PK) and duplicate handling match the scan path exactly.
/// Every listed value must be exactly of the PK column's type: the residual
/// `IN` filter coerces across numeric types via `compare_values`, which
/// type-tagged encoded probes cannot honor — mixed-type lists fall back to
/// the scan. The probe itself is tier-aware downstream (row fetches resolve
/// cold-tier copies per pk), so no residency gate is needed.
///
/// The candidates are caller-supplied values, not index postings: a listed
/// value may name a row that never existed or was since deleted. The result
/// is therefore marked `pks_are_live_rows: false` so candidate-count
/// shortcuts (page-window truncation, OFFSET draining) don't cut off real
/// matches behind probes that fail to hydrate.
fn primary_key_in_probe(
    context: &IndexLookupContext<'_>,
    column: &str,
    values: &[Value],
) -> Option<IndexLookupResult> {
    let schema = context.catalog.tables.get(&(
        namespace_key(context.project_id, context.scope_id),
        context.table_name.to_string(),
    ))?;
    if schema.primary_key.len() != 1 || schema.primary_key[0] != column {
        return None;
    }
    let pk_type = schema
        .columns
        .iter()
        .find(|c| c.name == column)
        .map(|c| &c.col_type)?;
    if !values
        .iter()
        .all(|value| value_is_exactly_column_type(value, pk_type))
    {
        return None;
    }
    let mut pks: Vec<EncodedKey> = values
        .iter()
        .map(|value| EncodedKey::from_values(std::slice::from_ref(value)))
        .collect();
    pks.sort();
    pks.dedup();
    Some(IndexLookupResult {
        pks,
        selected_indexes: Vec::new(),
        predicate_exact: true,
        pks_are_live_rows: false,
        plan_trace: plan_trace_if_diagnostic(context.include_diagnostics, || {
            format!("IN over primary key '{column}' resolves to direct point probes")
        }),
    })
}

/// Pick the first usable single-column index over `column`, honoring the
/// ordered-store requirement for range lookups and partial-index implication
/// against the predicate's equality constraints. Shared by the eager lookup
/// path and the chunked residual-candidate stream so selection can't diverge.
#[allow(clippy::too_many_arguments)]
fn select_single_column_index<'a>(
    catalog: &'a Catalog,
    table: &crate::storage::keyspace::TableData,
    project_id: &str,
    scope_id: &str,
    table_name: &str,
    predicate: &crate::query::plan::Expr,
    column: &str,
    requires_ordered: bool,
) -> Option<&'a str> {
    let ns = namespace_key(project_id, scope_id);
    let mut equalities = None;
    for ((p, t, idx_name), idx_def) in &catalog.indexes {
        if p != &ns
            || t != table_name
            || idx_def.columns.len() != 1
            || idx_def.columns[0] != column
            || !table.indexes.contains_key(idx_name)
            || (requires_ordered
                && !matches!(idx_def.index_type, IndexType::BTree | IndexType::Art))
        {
            continue;
        }
        if let Some(filter) = &idx_def.partial_filter {
            let equalities = equalities.get_or_insert_with(|| {
                let mut equalities = HashMap::new();
                collect_eq_constraints(predicate, &mut equalities);
                equalities
            });
            if !expr_implied_by_eq_constraints(filter, equalities) {
                continue;
            }
        }
        return Some(idx_name.as_str());
    }
    None
}

/// First window fetched by [`ChunkedIndexRangePks`]; doubles per fetch up to
/// [`MAX_CANDIDATE_CHUNK`] so sparse residual matches converge in O(n log n)
/// index-entry visits instead of O(n²/chunk).
const INITIAL_CANDIDATE_CHUNK: usize = 64;
const MAX_CANDIDATE_CHUNK: usize = 4096;

/// Streams the primary keys posted under an index-value range in index order,
/// fetching geometrically growing windows on demand instead of cloning every
/// matching posting up front. Only used over a snapshot-frozen resident index,
/// so successive offset windows observe the same postings.
pub(super) struct ChunkedIndexRangePks<'a> {
    index: &'a crate::storage::keyspace::SecondaryIndex,
    bounds: super::index_lookup::IndexBounds,
    next_offset: usize,
    chunk: usize,
    buffered: std::vec::IntoIter<EncodedKey>,
    exhausted: bool,
}

impl Iterator for ChunkedIndexRangePks<'_> {
    type Item = EncodedKey;

    fn next(&mut self) -> Option<EncodedKey> {
        loop {
            if let Some(pk) = self.buffered.next() {
                return Some(pk);
            }
            if self.exhausted {
                return None;
            }
            let batch = self.index.scan_range_window_ordered(
                self.bounds.0.clone(),
                self.bounds.1.clone(),
                self.next_offset,
                self.chunk,
                false,
            );
            if batch.len() < self.chunk {
                self.exhausted = true;
            }
            self.next_offset += batch.len();
            self.chunk = self.chunk.saturating_mul(2).min(MAX_CANDIDATE_CHUNK);
            if batch.is_empty() {
                return None;
            }
            self.buffered = batch.into_iter();
        }
    }
}

/// A residual-candidate stream plus the index it reads, for plan diagnostics.
pub(super) struct ResidualCandidateStream<'a> {
    pub(super) index_name: String,
    pub(super) pks: ChunkedIndexRangePks<'a>,
}

/// Chunked candidate streaming for a predicate that is a single residual
/// range-shaped leaf over an indexed column — in practice `LIKE` with a
/// non-exact prefix, the one leaf `extract_indexable_predicate` marks
/// non-exact. Every other leaf is predicate-exact and already served by the
/// bounded eager lookup. Returns `None` (falling back to the eager path) for
/// composed And/Or predicates, exact lookups, unindexed columns, and indexes
/// with a cold segment tier (the resident window walk would miss evicted
/// postings).
pub(super) fn residual_index_candidate_stream<'a>(
    catalog: &Catalog,
    project_id: &str,
    scope_id: &str,
    table_name: &str,
    table: &'a crate::storage::keyspace::TableData,
    predicate: &crate::query::plan::Expr,
) -> Option<ResidualCandidateStream<'a>> {
    let lookup = extract_indexable_predicate(predicate)?;
    let IndexLookup::Range {
        column,
        bounds,
        predicate_exact: false,
    } = lookup
    else {
        return None;
    };
    let index_name = select_single_column_index(
        catalog, table, project_id, scope_id, table_name, predicate, column, true,
    )?;
    let index = table.indexes.get(index_name)?;
    if index.has_cold_segments() {
        return None;
    }
    Some(ResidualCandidateStream {
        index_name: index_name.to_string(),
        pks: ChunkedIndexRangePks {
            index,
            bounds,
            next_offset: 0,
            chunk: INITIAL_CANDIDATE_CHUNK,
            buffered: Vec::new().into_iter(),
            exhausted: false,
        },
    })
}

fn indexed_pks_for_leftmost_composite_eq(
    context: &IndexLookupContext<'_>,
    column: &str,
    value: &Value,
    predicate: &crate::query::plan::Expr,
    candidate_limit: Option<usize>,
) -> Result<Option<IndexLookupResult>, QueryError> {
    let mut equalities = HashMap::new();
    collect_eq_constraints(predicate, &mut equalities);
    if !equalities.contains_key(column) {
        equalities.insert(column.to_string(), value.clone());
    }

    composite_prefix_index_lookup(
        context,
        &equalities,
        CompositeSelectionCriteria {
            first_column: Some(column),
            min_index_cols: 2,
        },
        candidate_limit,
    )
}