xz-rag 0.1.1

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use super::chunk::ChunkMetadata;

use crate::pipeline::channel::ChannelConfig;

// === Retrieval Request ===

/// Multi-channel retrieval request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrieveRequest {
    /// Search query string.
    pub query: String,
    /// Channel configurations for this request.
    pub channels: Vec<ChannelConfig>,
    /// Global filters applied across all channels.
    ///
    /// **Note on channel support**: `StructuredFilter` values are currently only applied by the
    /// `metadata` channel. The `semantic`, `bm25`, and `graph` channels silently ignore
    /// `global_filters` in their current implementation.
    pub global_filters: Vec<StructuredFilter>,
    /// Maximum number of results after fusion.
    pub top_k: usize,
    /// Optional namespace for scoping.
    pub namespace: Option<String>,
    /// Whether to include embedding vectors in results.
    pub include_embeddings: bool,
    /// Optional query preprocessing strategy.
    pub query_preprocessing: Option<QueryPreprocessing>,
}

impl RetrieveRequest {
    /// Create a new builder for `RetrieveRequest`.
    pub fn builder(query: impl Into<String>) -> RetrieveRequestBuilder {
        RetrieveRequestBuilder {
            query: query.into(),
            channels: vec![ChannelConfig::semantic(0.5, 10)],
            global_filters: vec![],
            top_k: 10,
            namespace: None,
            include_embeddings: false,
            query_preprocessing: None,
        }
    }
}

/// Builder for `RetrieveRequest`.
pub struct RetrieveRequestBuilder {
    query: String,
    channels: Vec<ChannelConfig>,
    global_filters: Vec<StructuredFilter>,
    top_k: usize,
    namespace: Option<String>,
    include_embeddings: bool,
    query_preprocessing: Option<QueryPreprocessing>,
}

impl RetrieveRequestBuilder {
    /// Set the channels to search.
    pub fn channels(mut self, channels: Vec<ChannelConfig>) -> Self {
        self.channels = channels;
        self
    }

    /// Set the maximum number of results.
    pub fn top_k(mut self, top_k: usize) -> Self {
        self.top_k = top_k;
        self
    }

    /// Set the namespace for scoping.
    pub fn namespace(mut self, ns: impl Into<String>) -> Self {
        self.namespace = Some(ns.into());
        self
    }

    /// Build the `RetrieveRequest`.
    pub fn build(self) -> RetrieveRequest {
        RetrieveRequest {
            query: self.query,
            channels: self.channels,
            global_filters: self.global_filters,
            top_k: self.top_k,
            namespace: self.namespace,
            include_embeddings: self.include_embeddings,
            query_preprocessing: self.query_preprocessing,
        }
    }
}

// === Query Preprocessing ===

/// Query preprocessing strategy.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryPreprocessing {
    /// No preprocessing.
    None,
    /// HYDE: generate a hypothetical answer passage.
    Hyde,
    /// Generate N query variations for better recall.
    QueryExpansion {
        /// Number of variations to generate.
        count: usize,
    },
    /// Translate query to English before retrieval.
    TranslateToEnglish,
}

// === Structured Filter ===

/// Structured filter for metadata-based retrieval.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StructuredFilter {
    /// Field equals value.
    MetadataEq {
        /// Metadata key.
        key: String,
        /// Expected value.
        value: String,
    },
    /// Field is one of the given values.
    MetadataIn {
        /// Metadata key.
        key: String,
        /// Allowed values.
        values: Vec<String>,
    },
    /// Field does not equal value.
    MetadataNe {
        /// Metadata key.
        key: String,
        /// Excluded value.
        value: String,
    },
    /// Field exists (has any value).
    MetadataExists {
        /// Metadata key.
        key: String,
    },
    /// Raw SQL filter expression.
    SqlFilter(String),
    /// Logical AND of two filters.
    And(Box<StructuredFilter>, Box<StructuredFilter>),
    /// Logical OR of two filters.
    Or(Box<StructuredFilter>, Box<StructuredFilter>),
}

// === Retrieve Result ===

/// Result of a multi-channel retrieval operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrieveResult {
    /// Ranked list of retrieved chunks.
    pub hits: Vec<RetrievedChunk>,
    /// Per-channel statistics.
    pub channel_report: HashMap<String, ChannelStats>,
    /// Total latency in milliseconds.
    pub latency_ms: u64,
    /// The query text actually used (after preprocessing).
    pub effective_query: String,
}

/// A single retrieved chunk with score and metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrievedChunk {
    /// Chunk identifier.
    pub chunk_id: String,
    /// Parent document identifier.
    pub document_id: String,
    /// Chunk text content.
    pub content: String,
    /// Relevance score (post-fusion).
    pub score: f32,
    /// Channel that produced this chunk.
    pub channel: String,
    /// Original score from the channel (pre-fusion).
    pub channel_score: f32,
    /// Chunk metadata.
    pub metadata: ChunkMetadata,
    /// Optional embedding vector.
    pub embedding: Option<Vec<f32>>,
}

/// Per-channel retrieval statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelStats {
    /// Channel type identifier.
    pub channel_type: String,
    /// Number of hits from this channel.
    pub hits: usize,
    /// Latency in milliseconds.
    pub latency_ms: u64,
    /// Minimum score among hits.
    pub min_score: f32,
    /// Maximum score among hits.
    pub max_score: f32,
}