bigrag 0.1.1

use serde::{Deserialize, Serialize};

/// Search mode for query operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SearchMode {
    /// Vector similarity search.
    Semantic,
    /// Keyword/BM25 search.
    Keyword,
    /// Combined semantic + keyword search.
    Hybrid,
}

/// Body for a single-collection query.
#[derive(Debug, Clone, Default, Serialize)]
pub struct QueryBody {
    /// Search query text (required).
    pub query: String,
    /// Maximum number of results.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_k: Option<u32>,
    /// Metadata filters.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub filters: Option<serde_json::Value>,
    /// Minimum similarity score.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_score: Option<f64>,
    /// Search mode.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_mode: Option<SearchMode>,
    /// Whether to apply reranking.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rerank: Option<bool>,
}

/// A single query result.
#[derive(Debug, Clone, Deserialize)]
pub struct QueryResult {
    /// Result ID.
    pub id: String,
    /// Matched text content.
    pub text: String,
    /// Similarity score (0–1).
    pub score: f64,
    /// Source document ID.
    pub document_id: Option<String>,
    /// Chunk index within the source document.
    pub chunk_index: Option<u32>,
    /// Result metadata.
    pub metadata: serde_json::Value,
}

/// Response from a single-collection query.
#[derive(Debug, Clone, Deserialize)]
pub struct QueryResponse {
    /// Matched results.
    pub results: Vec<QueryResult>,
    /// The query that was executed.
    pub query: String,
    /// Collection that was searched.
    pub collection: String,
    /// Total number of results.
    pub total: u32,
}

/// Body for a multi-collection query.
#[derive(Debug, Clone, Serialize)]
pub struct MultiQueryBody {
    /// Search query text (required).
    pub query: String,
    /// Collections to search across (required).
    pub collections: Vec<String>,
    /// Maximum number of results.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_k: Option<u32>,
    /// Metadata filters.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub filters: Option<serde_json::Value>,
    /// Minimum similarity score.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_score: Option<f64>,
    /// Search mode.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_mode: Option<SearchMode>,
    /// Whether to apply reranking.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rerank: Option<bool>,
}

/// A single result from a multi-collection query, includes source collection.
#[derive(Debug, Clone, Deserialize)]
pub struct MultiQueryResult {
    /// Result ID.
    pub id: String,
    /// Matched text content.
    pub text: String,
    /// Similarity score.
    pub score: f64,
    /// Source document ID.
    pub document_id: Option<String>,
    /// Chunk index within the source document.
    pub chunk_index: Option<u32>,
    /// Collection this result came from.
    pub collection: String,
    /// Result metadata.
    pub metadata: serde_json::Value,
}

/// Response from a multi-collection query.
#[derive(Debug, Clone, Deserialize)]
pub struct MultiQueryResponse {
    /// Matched results across collections.
    pub results: Vec<MultiQueryResult>,
    /// The query that was executed.
    pub query: String,
    /// Collections that were searched.
    pub collections: Vec<String>,
    /// Total number of results.
    pub total: u32,
}

/// A single query within a batch query request.
#[derive(Debug, Clone, Serialize)]
pub struct BatchQueryItem {
    /// Collection to query.
    pub collection: String,
    /// Search query text.
    pub query: String,
    /// Maximum number of results.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_k: Option<u32>,
    /// Metadata filters.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub filters: Option<serde_json::Value>,
    /// Minimum similarity score.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_score: Option<f64>,
    /// Search mode.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_mode: Option<SearchMode>,
    /// Whether to apply reranking.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rerank: Option<bool>,
}

/// Body for a batch query request.
#[derive(Debug, Clone, Serialize)]
pub struct BatchQueryBody {
    /// Queries to execute (1–20).
    pub queries: Vec<BatchQueryItem>,
}

/// Response from a batch query.
#[derive(Debug, Clone, Deserialize)]
pub struct BatchQueryResponse {
    /// One `QueryResponse` per query in the batch.
    pub results: Vec<QueryResponse>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_search_mode_serializes_lowercase() {
        let json = serde_json::to_value(SearchMode::Hybrid).unwrap();
        assert_eq!(json, "hybrid");
    }

    #[test]
    fn test_serialize_query_body_skips_none() {
        let body = QueryBody {
            query: "hello world".into(),
            ..Default::default()
        };
        let json = serde_json::to_value(&body).unwrap();
        assert_eq!(json["query"], "hello world");
        assert!(json.get("top_k").is_none());
        assert!(json.get("search_mode").is_none());
    }

    #[test]
    fn test_deserialize_query_response() {
        let json = r#"{
            "results": [{"id":"r1","text":"hello","score":0.95,"document_id":"d1","chunk_index":0,"metadata":{}}],
            "query": "hello",
            "collection": "docs",
            "total": 1
        }"#;
        let resp: QueryResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.total, 1);
        assert_eq!(resp.results[0].score, 0.95);
        assert_eq!(resp.results[0].document_id, Some("d1".into()));
    }

    #[test]
    fn test_deserialize_multi_query_response() {
        let json = r#"{
            "results": [{"id":"r1","text":"hi","score":0.9,"document_id":null,"chunk_index":null,"collection":"docs","metadata":{}}],
            "query": "hi",
            "collections": ["docs"],
            "total": 1
        }"#;
        let resp: MultiQueryResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.results[0].collection, "docs");
        assert_eq!(resp.results[0].document_id, None);
    }

    #[test]
    fn test_deserialize_batch_query_response() {
        let json = r#"{"results":[{"results":[],"query":"test","collection":"a","total":0}]}"#;
        let resp: BatchQueryResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.results.len(), 1);
        assert_eq!(resp.results[0].collection, "a");
    }
}