qql-core 0.4.1

use super::{
    ComparisonOp, FilterExpr, PointId, PointIdPredicate, PointSelector, Prefetch, PrefetchSource,
    QueryExpr, QueryStmt, ShardKey, Stmt, Value,
};
use crate::error::QqlError;
use alloc::boxed::Box;
use alloc::string::ToString;

impl Stmt {
    /// Custom shard routing key for this statement, if any.
    ///
    /// Corresponds to QQL `SHARD` on DML, lowered to request-level
    /// `shard_key` (REST) / `ShardKeySelector` (gRPC) — never inside `Filter`.
    /// Keyword and numeric forms are preserved (`ShardKey::Number(101)` reads
    /// back as a number, never coerced to `"101"`).
    pub fn shard_key(&self) -> Option<&ShardKey> {
        match self {
            Self::Query(query) => query.shard_key.as_ref(),
            Self::Scroll(scroll) => scroll.shard_key.as_ref(),
            Self::Count(count) => count.shard_key.as_ref(),
            Self::Facet(facet) => facet.shard_key.as_ref(),
            Self::Upsert(upsert) => upsert.shard_key.as_ref(),
            Self::Delete(delete) => delete.shard_key.as_ref(),
            Self::ClearPayload(clear) => clear.shard_key.as_ref(),
            Self::DeletePayload(delete) => delete.shard_key.as_ref(),
            Self::DeleteVector(delete) => delete.shard_key.as_ref(),
            Self::UpdateVector(update) => update.shard_key.as_ref(),
            Self::UpdatePayload(update) => update.shard_key.as_ref(),
            Self::Batch(batch) => {
                // Members of one batch RPC normally share routing; report a
                // key only when every member agrees so callers never act on
                // a partial view.
                let mut keys = batch.statements.iter().map(Stmt::shard_key);
                match keys.next() {
                    None => None,
                    Some(first) => {
                        if keys.all(|key| key == first) {
                            first
                        } else {
                            None
                        }
                    }
                }
            }
            _ => None,
        }
    }

    /// Set custom shard routing (same field as QQL `SHARD`).
    ///
    /// Prefer writing the `SHARD` clause in the query when the tenant is known
    /// at authoring time. Use this setter only when the host resolves the key
    /// after parse (e.g. from auth context) without re-stringifying QQL.
    ///
    /// On `QUERY`, recurses into CTEs and nested prefetch queries so routing
    /// matches a top-level `SHARD` clause. `None` (or an empty keyword) clears
    /// the key.
    /// Returns `false` for statement types that cannot carry routing (DDL, SHOW).
    pub fn set_shard_key(&mut self, shard_key: Option<ShardKey>) -> bool {
        let shard_key = shard_key.filter(|k| !matches!(k, ShardKey::Keyword(s) if s.is_empty()));
        match self {
            Self::Query(query) => {
                apply_query_shard(query, shard_key.as_ref());
                true
            }
            Self::Scroll(scroll) => {
                scroll.shard_key = shard_key;
                true
            }
            Self::Count(count) => {
                count.shard_key = shard_key;
                true
            }
            Self::Facet(facet) => {
                facet.shard_key = shard_key;
                true
            }
            Self::Upsert(upsert) => {
                upsert.shard_key = shard_key;
                true
            }
            Self::Delete(delete) => {
                delete.shard_key = shard_key;
                true
            }
            Self::ClearPayload(clear) => {
                clear.shard_key = shard_key;
                true
            }
            Self::DeletePayload(delete) => {
                delete.shard_key = shard_key;
                true
            }
            Self::DeleteVector(delete) => {
                delete.shard_key = shard_key;
                true
            }
            Self::UpdateVector(update) => {
                update.shard_key = shard_key;
                true
            }
            Self::UpdatePayload(update) => {
                update.shard_key = shard_key;
                true
            }
            Self::Batch(batch) => {
                // Pre-check before mutating: a mixed batch (e.g. containing
                // DDL/SHOW) must report failure without half-applying keys to
                // the members that precede the unsupported one.
                if !batch.statements.iter().all(can_carry_shard_key) {
                    return false;
                }
                for member in &mut batch.statements {
                    member.set_shard_key(shard_key.clone());
                }
                true
            }
            _ => false,
        }
    }
}

/// Whether a statement can carry `SHARD` routing: mirrors the capable arms
/// of [`Stmt::set_shard_key`]. Used to pre-check batches before mutating so a
/// mixed batch never half-applies.
fn can_carry_shard_key(statement: &Stmt) -> bool {
    match statement {
        Stmt::Query(_)
        | Stmt::Scroll(_)
        | Stmt::Count(_)
        | Stmt::Facet(_)
        | Stmt::Upsert(_)
        | Stmt::Delete(_)
        | Stmt::ClearPayload(_)
        | Stmt::DeletePayload(_)
        | Stmt::DeleteVector(_)
        | Stmt::UpdateVector(_)
        | Stmt::UpdatePayload(_) => true,
        Stmt::Batch(batch) => batch.statements.iter().all(can_carry_shard_key),
        _ => false,
    }
}

/// Apply shard routing to a query and nested CTE / prefetch queries.
fn apply_query_shard(query: &mut QueryStmt, key: Option<&ShardKey>) {
    query.shard_key = key.cloned();
    for cte in &mut query.ctes {
        apply_query_shard(&mut cte.query, key);
    }
    if let Some(prefetches) = expression_prefetch(&mut query.expression) {
        for prefetch in prefetches {
            if let PrefetchSource::Query(nested) = &mut prefetch.source {
                apply_query_shard(nested, key);
            }
        }
    }
}

/// Injects a typed field comparison into a statement (CTEs and prefetches included), fail-closed.
///
/// `QUERY` merges the comparison into the top-level filter, every CTE, and
/// every prefetch stage; `SCROLL` / `COUNT` / `FACET` merge into their filter;
/// selector statements (`DELETE`, `CLEAR PAYLOAD`, `DELETE PAYLOAD`,
/// `DELETE VECTOR`, `UPDATE PAYLOAD`) fold the filter into their point
/// selector. `BATCH` recurses into every member.
///
/// `UPSERT` has no filter to merge into, so the field is stamped onto each
/// point's payload instead, last-writer-wins: a pre-existing conflicting value
/// is overwritten, never merged. Overwriting (rather than skipping on
/// conflict) is what makes tenant stamping sound — a caller-supplied payload
/// value must not survive the injected policy value. Only `Eq` on a non-`id`
/// payload field is stampable; anything else fails closed, as does an unbound
/// whole-point placeholder (`:name` / `?`), which has no payload yet —
/// silently skipping it would drop a security filter, so callers must bind
/// point parameters before injection.
///
/// `UPDATE VECTOR` is an explicit deny, not an oversight: point-vector
/// replacement addresses points by ID list and carries no selector, so there
/// is nothing to merge a tenant filter into. Filter by IDs before building
/// the statement instead. DDL, `SHOW`, quota, and any other statement type
/// without a filter or payload surface fail closed with
/// `QQL-VALIDATION-FILTER-INJECT` rather than silently no-oping, so a policy
/// bypass can never hide behind an unsupported statement kind.
pub fn inject_filter(
    statement: &mut Stmt,
    field: &str,
    operator: ComparisonOp,
    value: Value,
) -> Result<(), QqlError> {
    let filter = build_filter(field, operator, value.clone())?;
    match statement {
        Stmt::Query(query) => inject_query(query, &filter),
        Stmt::Batch(batch) => {
            for member in &mut batch.statements {
                inject_filter(member, field, operator, value.clone())?;
            }
        }
        Stmt::Scroll(scroll) => merge_filter(&mut scroll.filter, filter),
        Stmt::Delete(delete) => merge_selector(&mut delete.selector, filter),
        Stmt::Count(count) => merge_filter(&mut count.filter, filter),
        Stmt::Facet(facet) => merge_filter(&mut facet.filter, filter),
        Stmt::ClearPayload(clear) => merge_selector(&mut clear.selector, filter),
        Stmt::DeletePayload(del) => merge_selector(&mut del.selector, filter),
        Stmt::DeleteVector(del_vec) => merge_selector(&mut del_vec.selector, filter),
        Stmt::UpdatePayload(update) => merge_selector(&mut update.selector, filter),
        // Explicit deny (see rustdoc): no selector exists to merge into.
        Stmt::UpdateVector(_) => {
            return Err(QqlError::validation(
                "QQL-VALIDATION-FILTER-INJECT",
                "inject_filter does not apply to this statement type (UPDATE VECTOR): point-vector replacement has no selector; address points by ID",
                None,
            ));
        }
        Stmt::Upsert(_) if operator != ComparisonOp::Eq || field.eq_ignore_ascii_case("id") => {
            return Err(QqlError::validation(
                "QQL-VALIDATION-FILTER-INJECT",
                "inject_filter into UPSERT requires Eq on a non-id payload field",
                None,
            ));
        }
        Stmt::Upsert(upsert) => {
            for point in &mut upsert.points {
                // A whole-point placeholder has no payload yet (see rustdoc).
                let inline = match point {
                    crate::ast::PointEntry::Inline(inline) => inline,
                    crate::ast::PointEntry::Param(name, _) => {
                        return Err(QqlError::validation(
                            "QQL-VALIDATION-FILTER-INJECT",
                            alloc::format!(
                                "cannot inject filter into unbound point parameter ':{name}'; bind point parameters before filter injection"
                            ),
                            None,
                        ));
                    }
                    crate::ast::PointEntry::PositionalParam(idx, _) => {
                        return Err(QqlError::validation(
                            "QQL-VALIDATION-FILTER-INJECT",
                            alloc::format!(
                                "cannot inject filter into unbound point parameter '?{}'; bind point parameters before filter injection",
                                *idx + 1
                            ),
                            None,
                        ));
                    }
                };
                if let Some((_, current)) = inline
                    .payload
                    .iter_mut()
                    .find(|(key, _)| key.eq_ignore_ascii_case(field))
                {
                    // Last-writer-wins (see rustdoc): overwrite, never merge.
                    *current = value.clone();
                } else {
                    inline.payload.push((field.to_string(), value.clone()));
                }
            }
        }
        other => {
            return Err(QqlError::validation(
                "QQL-VALIDATION-FILTER-INJECT",
                format!(
                    "inject_filter does not apply to this statement type ({})",
                    other.stmt_kind()
                ),
                None,
            ));
        }
    }
    Ok(())
}

fn build_filter(field: &str, operator: ComparisonOp, value: Value) -> Result<FilterExpr, QqlError> {
    if field.eq_ignore_ascii_case("id") {
        if operator != ComparisonOp::Eq {
            return Err(QqlError::validation(
                "QQL-VALIDATION-ID-PREDICATE",
                "point ID injection supports equality only",
                None,
            ));
        }
        let id = match value {
            Value::Int(value) if value >= 0 => PointId::Number(value as u64),
            Value::UInt(value) => PointId::Number(value),
            Value::Str(value) => PointId::String(value),
            _ => {
                return Err(QqlError::validation(
                    "QQL-VALIDATION-POINT-ID",
                    "point IDs must be unsigned integers or strings",
                    None,
                ));
            }
        };
        Ok(FilterExpr::PointId(PointIdPredicate::Eq(id)))
    } else {
        Ok(FilterExpr::Compare {
            field: field.to_string(),
            op: operator,
            value,
        })
    }
}

fn inject_query(query: &mut QueryStmt, filter: &FilterExpr) {
    merge_filter(&mut query.filter, filter.clone());
    for cte in &mut query.ctes {
        inject_query(&mut cte.query, filter);
    }
    if let Some(prefetches) = expression_prefetch(&mut query.expression) {
        for prefetch in prefetches {
            merge_filter(&mut prefetch.filter, filter.clone());
            if let PrefetchSource::Query(query) = &mut prefetch.source {
                inject_query(query, filter);
            }
        }
    }
}

fn expression_prefetch(expression: &mut QueryExpr) -> Option<&mut Vec<Prefetch>> {
    match expression {
        QueryExpr::Nearest { prefetch, .. }
        | QueryExpr::Recommend { prefetch, .. }
        | QueryExpr::Context { prefetch, .. }
        | QueryExpr::Discover { prefetch, .. }
        | QueryExpr::Fusion { prefetch, .. }
        | QueryExpr::Formula { prefetch, .. }
        | QueryExpr::RelevanceFeedback { prefetch, .. }
        | QueryExpr::Rerank { prefetch, .. }
        | QueryExpr::CrossRerank { prefetch, .. } => Some(prefetch),
        QueryExpr::Points { .. }
        | QueryExpr::OrderBy { .. }
        | QueryExpr::SampleRandom
        | QueryExpr::Hybrid { .. } => None,
    }
}

fn merge_selector(selector: &mut PointSelector, filter: FilterExpr) {
    let current = match core::mem::replace(selector, PointSelector::Ids(Vec::new())) {
        PointSelector::Id(id) => FilterExpr::PointId(PointIdPredicate::Eq(id)),
        PointSelector::Ids(ids) => FilterExpr::PointId(PointIdPredicate::In(ids)),
        PointSelector::Filter(existing) => *existing,
    };
    *selector = PointSelector::Filter(Box::new(and(current, filter)));
}

fn merge_filter(current: &mut Option<Box<FilterExpr>>, filter: FilterExpr) {
    *current = Some(Box::new(match current.take() {
        Some(current) => and(*current, filter),
        None => filter,
    }));
}

fn and(left: FilterExpr, right: FilterExpr) -> FilterExpr {
    match left {
        FilterExpr::And { mut operands } => {
            operands.push(right);
            FilterExpr::And { operands }
        }
        left => FilterExpr::And {
            operands: alloc::vec![left, right],
        },
    }
}