linkrs-query 0.1.0

//! Vector Search AST Definitions
//!
//! This module defines the Abstract Syntax Tree (AST) nodes for vector search queries,
//! including CREATE VECTOR INDEX, SEARCH VECTOR, and related statements.

use linkrs_core::types::expr::contextual::ContextualExpression;
use linkrs_core::types::span::Span;
use serde::{Deserialize, Serialize};

use super::stmt::OrderByClause;
use super::types::{LimitClause, SkipClause};

// ============================================================================
// Vector Index DDL Statements
// ============================================================================

/// CREATE VECTOR INDEX statement
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateVectorIndex {
    pub span: Span,
    pub index_name: String,
    pub schema_name: String,
    pub field_name: String,
    pub config: VectorIndexConfig,
    pub if_not_exists: bool,
}

/// Quantization kind for vector index (mirrors `simvec::QuantizationType` / Qdrant builders)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum QuantizationKind {
    Scalar,
    Binary,
    Product,
}

/// Compression ratio for product quantization (mirrors `simvec::CompressionRatio` / Qdrant)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum CompressionRatioKind {
    X4,
    X8,
    X16,
    X32,
    X64,
}

/// Vector index configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorIndexConfig {
    pub vector_size: usize,
    pub distance: VectorDistance,
    pub hnsw_m: Option<usize>,
    pub hnsw_ef_construct: Option<usize>,
    /// Quantization kind: None = disabled (exact f32), Some = scalar/binary/product
    pub quantization: Option<QuantizationKind>,
    /// Scalar only: quantile in (0,1], default 0.99
    pub quantile: Option<f32>,
    /// Product only: compression ratio, default X4
    pub compression: Option<CompressionRatioKind>,
    /// Always keep quantized vectors in RAM (scalar/binary/product)
    pub always_ram: Option<bool>,
}

/// DROP VECTOR INDEX statement
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DropVectorIndex {
    pub span: Span,
    pub index_name: String,
    pub if_exists: bool,
}

// ============================================================================
// Vector Search DML Statements
// ============================================================================

/// SEARCH VECTOR statement
#[derive(Debug, Clone)]
pub struct SearchVectorStatement {
    pub span: Span,
    pub index_name: String,
    pub query: VectorQueryExpr,
    pub threshold: Option<f32>,
    pub where_clause: Option<ContextualExpression>,
    pub order_clause: Option<OrderByClause>,
    pub limit: Option<LimitClause>,
    pub skip: Option<SkipClause>,
    pub yield_clause: Option<VectorYieldClause>,
}

/// Vector query expression
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorQueryExpr {
    pub span: Span,
    pub query_type: VectorQueryType,
    pub query_data: String,
}

/// Vector query type
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum VectorQueryType {
    /// Direct vector: [0.1, 0.2, ...]
    Vector,
    /// Text query (requires embedding service)
    Text,
    /// Parameter reference: $param_name
    Parameter,
}

/// Vector distance metric
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum VectorDistance {
    Cosine,
    Euclidean,
    Dot,
    Manhattan,
}

/// YIELD clause for vector search
#[derive(Debug, Clone)]
pub struct VectorYieldClause {
    pub items: Vec<VectorYieldItem>,
}

/// Yield item for vector search
#[derive(Debug, Clone)]
pub struct VectorYieldItem {
    pub expr: ContextualExpression,
    pub alias: Option<String>,
}

// ============================================================================
// MATCH and LOOKUP extensions for vector search
// ============================================================================

/// MATCH with vector search
#[derive(Debug, Clone)]
pub struct MatchVector {
    pub span: Span,
    pub pattern: String,
    pub vector_condition: VectorMatchCondition,
    pub yield_clause: Option<VectorYieldClause>,
}

/// Vector match condition in WHERE clause
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorMatchCondition {
    pub field: String,
    pub query: VectorQueryExpr,
    pub threshold: Option<f32>,
}

/// LOOKUP with vector search
#[derive(Debug, Clone)]
pub struct LookupVector {
    pub span: Span,
    pub schema_name: String,
    pub index_name: String,
    pub query: VectorQueryExpr,
    pub yield_clause: Option<VectorYieldClause>,
    pub limit: Option<usize>,
}

// ============================================================================
// Helper Functions
// ============================================================================

impl CreateVectorIndex {
    pub fn new(
        span: Span,
        index_name: String,
        schema_name: String,
        field_name: String,
        config: VectorIndexConfig,
    ) -> Self {
        Self {
            span,
            index_name,
            schema_name,
            field_name,
            config,
            if_not_exists: false,
        }
    }

    pub fn with_if_not_exists(mut self, if_not_exists: bool) -> Self {
        self.if_not_exists = if_not_exists;
        self
    }
}

impl DropVectorIndex {
    pub fn new(span: Span, index_name: String) -> Self {
        Self {
            span,
            index_name,
            if_exists: false,
        }
    }

    pub fn with_if_exists(mut self, if_exists: bool) -> Self {
        self.if_exists = if_exists;
        self
    }
}

impl VectorQueryExpr {
    pub fn vector(span: Span, vector_str: String) -> Self {
        Self {
            span,
            query_type: VectorQueryType::Vector,
            query_data: vector_str,
        }
    }

    pub fn text(span: Span, text: String) -> Self {
        Self {
            span,
            query_type: VectorQueryType::Text,
            query_data: text,
        }
    }

    pub fn parameter(span: Span, param_name: String) -> Self {
        Self {
            span,
            query_type: VectorQueryType::Parameter,
            query_data: param_name,
        }
    }
}

impl VectorIndexConfig {
    pub fn new(vector_size: usize, distance: VectorDistance) -> Self {
        Self {
            vector_size,
            distance,
            hnsw_m: None,
            hnsw_ef_construct: None,
            quantization: None,
            quantile: None,
            compression: None,
            always_ram: None,
        }
    }

    pub fn with_hnsw(mut self, m: usize, ef_construct: usize) -> Self {
        self.hnsw_m = Some(m);
        self.hnsw_ef_construct = Some(ef_construct);
        self
    }

    pub fn with_quantization(
        mut self,
        quantization: QuantizationKind,
        quantile: Option<f32>,
        compression: Option<CompressionRatioKind>,
        always_ram: Option<bool>,
    ) -> Self {
        self.quantization = Some(quantization);
        self.quantile = quantile;
        self.compression = compression;
        self.always_ram = always_ram;
        self
    }
}

impl CompressionRatioKind {
    pub fn as_str(&self) -> &'static str {
        match self {
            CompressionRatioKind::X4 => "x4",
            CompressionRatioKind::X8 => "x8",
            CompressionRatioKind::X16 => "x16",
            CompressionRatioKind::X32 => "x32",
            CompressionRatioKind::X64 => "x64",
        }
    }
}