//! Query Engine - Unified query layer.
//!
//! Supports SQL-like queries, document filters, and vector similarity search.

use crate::document::{Document, DocumentFilter, DocumentStore};
use crate::error::DatabaseError;
use crate::index::IndexManager;
use crate::vector::VectorStore;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use std::sync::Arc;

/// Sort order for query results.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SortOrder {
    Ascending,
    Descending,
}

/// Filter operators.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Filter {
    /// Equality: field = value
    Eq(String, JsonValue),
    /// Not equal: field != value
    Ne(String, JsonValue),
    /// Greater than: field > value
    Gt(String, JsonValue),
    /// Greater than or equal: field >= value
    Gte(String, JsonValue),
    /// Less than: field < value
    Lt(String, JsonValue),
    /// Less than or equal: field <= value
    Lte(String, JsonValue),
    /// In array: field IN [values]
    In(String, Vec<JsonValue>),
    /// Not in array: field NOT IN [values]
    NotIn(String, Vec<JsonValue>),
    /// String contains: field LIKE %value%
    Contains(String, String),
    /// String starts with: field LIKE value%
    StartsWith(String, String),
    /// String ends with: field LIKE %value
    EndsWith(String, String),
    /// Field exists
    Exists(String),
    /// Field is null
    IsNull(String),
    /// Logical AND
    And(Vec<Filter>),
    /// Logical OR
    Or(Vec<Filter>),
    /// Logical NOT
    Not(Box<Filter>),
    /// Vector similarity search
    VectorSimilar {
        field: String,
        vector: Vec<f32>,
        limit: usize,
        threshold: f32,
    },
}

impl Filter {
    /// Converts to DocumentFilter.
    pub fn to_document_filter(&self) -> DocumentFilter {
        match self {
            Filter::Eq(field, value) => DocumentFilter::new().eq(field, value.clone()),
            Filter::Ne(field, value) => DocumentFilter::new().ne(field, value.clone()),
            Filter::Gt(field, value) => DocumentFilter::new().gt(field, value.clone()),
            Filter::Gte(field, value) => DocumentFilter::new().gte(field, value.clone()),
            Filter::Lt(field, value) => DocumentFilter::new().lt(field, value.clone()),
            Filter::Lte(field, value) => DocumentFilter::new().lte(field, value.clone()),
            Filter::In(field, values) => DocumentFilter::new().in_array(field, values.clone()),
            Filter::Contains(field, substr) => DocumentFilter::new().contains(field, substr),
            Filter::Exists(field) => DocumentFilter::new().exists(field, true),
            Filter::IsNull(field) => DocumentFilter::new().eq(field, JsonValue::Null),
            Filter::And(filters) => {
                let doc_filters: Vec<_> = filters.iter().map(|f| f.to_document_filter()).collect();
                DocumentFilter::new().and(doc_filters)
            }
            Filter::Or(filters) => {
                let doc_filters: Vec<_> = filters.iter().map(|f| f.to_document_filter()).collect();
                DocumentFilter::new().or(doc_filters)
            }
            _ => DocumentFilter::new(),
        }
    }
}

/// Query builder.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Query {
    /// Collection to query.
    pub collection: String,
    /// Filter conditions.
    pub filter: Option<Filter>,
    /// Fields to select (empty = all).
    pub select: Vec<String>,
    /// Sort fields.
    pub sort: Vec<(String, SortOrder)>,
    /// Skip count.
    pub skip: usize,
    /// Limit count (0 = no limit).
    pub limit: usize,
    /// Include vector similarity.
    pub vector_query: Option<VectorQuery>,
}

/// Vector similarity query.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct VectorQuery {
    /// Field containing vectors.
    pub field: String,
    /// Query vector.
    pub vector: Vec<f32>,
    /// Number of results.
    pub limit: usize,
    /// Minimum similarity threshold.
    pub threshold: f32,
}

impl Query {
    /// Creates a new query for a collection.
    pub fn new(collection: impl Into<String>) -> Self {
        Self {
            collection: collection.into(),
            ..Default::default()
        }
    }

    /// Adds a filter.
    pub fn filter(mut self, filter: Filter) -> Self {
        self.filter = Some(filter);
        self
    }

    /// Adds an equality filter.
    pub fn eq(mut self, field: impl Into<String>, value: JsonValue) -> Self {
        let new_filter = Filter::Eq(field.into(), value);
        self.filter = Some(match self.filter {
            Some(f) => Filter::And(vec![f, new_filter]),
            None => new_filter,
        });
        self
    }

    /// Selects specific fields.
    pub fn select(mut self, fields: Vec<String>) -> Self {
        self.select = fields;
        self
    }

    /// Adds a sort field.
    pub fn sort(mut self, field: impl Into<String>, order: SortOrder) -> Self {
        self.sort.push((field.into(), order));
        self
    }

    /// Sets skip count.
    pub fn skip(mut self, n: usize) -> Self {
        self.skip = n;
        self
    }

    /// Sets limit count.
    pub fn limit(mut self, n: usize) -> Self {
        self.limit = n;
        self
    }

    /// Adds vector similarity search.
    pub fn vector_similar(
        mut self,
        field: impl Into<String>,
        vector: Vec<f32>,
        limit: usize,
        threshold: f32,
    ) -> Self {
        self.vector_query = Some(VectorQuery {
            field: field.into(),
            vector,
            limit,
            threshold,
        });
        self
    }
}

/// Query result.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct QueryResult {
    /// Matched documents.
    pub documents: Vec<Document>,
    /// Total count (before skip/limit).
    pub total: u64,
    /// Query execution time (ms).
    pub execution_time_ms: u64,
    /// Whether more results exist.
    pub has_more: bool,
}

impl QueryResult {
    /// Creates an empty result.
    pub fn empty() -> Self {
        Self {
            documents: Vec::new(),
            total: 0,
            execution_time_ms: 0,
            has_more: false,
        }
    }

    /// Returns the first document.
    pub fn first(&self) -> Option<&Document> {
        self.documents.first()
    }

    /// Maps documents to JSON values.
    pub fn to_json(&self) -> Vec<JsonValue> {
        self.documents.iter().map(|d| d.data.clone()).collect()
    }
}

/// Query execution engine.
pub struct QueryEngine {
    /// Document store reference.
    docs: Arc<DocumentStore>,
    /// Vector store reference.
    vectors: Arc<VectorStore>,
    /// Index manager reference.
    indexes: Arc<IndexManager>,
}

impl QueryEngine {
    /// Creates a new query engine.
    pub fn new(
        docs: Arc<DocumentStore>,
        vectors: Arc<VectorStore>,
        indexes: Arc<IndexManager>,
    ) -> Self {
        Self {
            docs,
            vectors,
            indexes,
        }
    }

    /// Executes a query.
    pub fn execute(&self, query: &Query) -> Result<QueryResult, DatabaseError> {
        let start = std::time::Instant::now();

        // Get documents from collection
        let filter = query.filter.as_ref().map(|f| f.to_document_filter());
        let default_filter = DocumentFilter::default();
        let filter_ref = filter.as_ref().unwrap_or(&default_filter);

        let mut documents = self.docs.find(&query.collection, filter_ref)?;

        // Apply vector similarity if specified
        if let Some(vq) = &query.vector_query {
            documents = self.apply_vector_query(documents, vq)?;
        }

        let total = documents.len() as u64;

        // Apply sorting
        if !query.sort.is_empty() {
            self.sort_documents(&mut documents, &query.sort);
        }

        // Apply skip and limit
        let mut has_more = false;
        if query.skip > 0 {
            if query.skip >= documents.len() {
                documents.clear();
            } else {
                documents = documents.split_off(query.skip);
            }
        }

        if query.limit > 0 && documents.len() > query.limit {
            documents.truncate(query.limit);
            has_more = true;
        }

        // Apply field selection
        if !query.select.is_empty() {
            documents = self.select_fields(documents, &query.select);
        }

        let execution_time_ms = start.elapsed().as_millis() as u64;

        Ok(QueryResult {
            documents,
            total,
            execution_time_ms,
            has_more,
        })
    }

    fn apply_vector_query(
        &self,
        documents: Vec<Document>,
        vq: &VectorQuery,
    ) -> Result<Vec<Document>, DatabaseError> {
        // For now, filter documents in-memory
        // Real implementation would use vector index
        let mut scored: Vec<(Document, f32)> = documents
            .into_iter()
            .filter_map(|doc| {
                let field_value = doc.get(&vq.field)?;
                let arr = field_value.as_array()?;

                let vec: Vec<f32> = arr
                    .iter()
                    .filter_map(|v| v.as_f64().map(|f| f as f32))
                    .collect();

                if vec.len() != vq.vector.len() {
                    return None;
                }

                let similarity = cosine_similarity(&vec, &vq.vector);
                if similarity >= vq.threshold {
                    Some((doc, similarity))
                } else {
                    None
                }
            })
            .collect();

        // Sort by similarity descending
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Apply limit
        scored.truncate(vq.limit);

        Ok(scored.into_iter().map(|(doc, _)| doc).collect())
    }

    fn sort_documents(&self, documents: &mut [Document], sort: &[(String, SortOrder)]) {
        documents.sort_by(|a, b| {
            for (field, order) in sort {
                let a_val = a.get(field);
                let b_val = b.get(field);

                let cmp = match (a_val, b_val) {
                    (Some(av), Some(bv)) => self.compare_json(av, bv),
                    (Some(_), None) => std::cmp::Ordering::Less,
                    (None, Some(_)) => std::cmp::Ordering::Greater,
                    (None, None) => std::cmp::Ordering::Equal,
                };

                let cmp = match order {
                    SortOrder::Ascending => cmp,
                    SortOrder::Descending => cmp.reverse(),
                };

                if cmp != std::cmp::Ordering::Equal {
                    return cmp;
                }
            }
            std::cmp::Ordering::Equal
        });
    }

    fn compare_json(&self, a: &JsonValue, b: &JsonValue) -> std::cmp::Ordering {
        match (a, b) {
            (JsonValue::Number(an), JsonValue::Number(bn)) => {
                let af = an.as_f64().unwrap_or(0.0);
                let bf = bn.as_f64().unwrap_or(0.0);
                af.partial_cmp(&bf).unwrap_or(std::cmp::Ordering::Equal)
            }
            (JsonValue::String(as_), JsonValue::String(bs)) => as_.cmp(bs),
            (JsonValue::Bool(ab), JsonValue::Bool(bb)) => ab.cmp(bb),
            _ => std::cmp::Ordering::Equal,
        }
    }

    fn select_fields(&self, documents: Vec<Document>, fields: &[String]) -> Vec<Document> {
        documents
            .into_iter()
            .map(|mut doc| {
                if let Some(obj) = doc.data.as_object() {
                    let filtered: serde_json::Map<String, JsonValue> = obj
                        .iter()
                        .filter(|(k, _)| fields.contains(k))
                        .map(|(k, v)| (k.clone(), v.clone()))
                        .collect();
                    doc.data = JsonValue::Object(filtered);
                }
                doc
            })
            .collect()
    }

    /// Counts documents matching a filter.
    pub fn count(&self, collection: &str, filter: Option<&Filter>) -> Result<u64, DatabaseError> {
        let doc_filter = filter.map(|f| f.to_document_filter());
        let default_filter = DocumentFilter::default();
        let filter_ref = doc_filter.as_ref().unwrap_or(&default_filter);
        let docs = self.docs.find(collection, filter_ref)?;
        Ok(docs.len() as u64)
    }

    /// Aggregates values.
    pub fn aggregate(
        &self,
        collection: &str,
        field: &str,
        op: AggregateOp,
        filter: Option<&Filter>,
    ) -> Result<JsonValue, DatabaseError> {
        let doc_filter = filter.map(|f| f.to_document_filter());
        let default_filter = DocumentFilter::default();
        let filter_ref = doc_filter.as_ref().unwrap_or(&default_filter);
        let docs = self.docs.find(collection, filter_ref)?;

        let values: Vec<f64> = docs
            .iter()
            .filter_map(|doc| {
                doc.get(field)
                    .and_then(|v| v.as_f64())
            })
            .collect();

        let result = match op {
            AggregateOp::Count => JsonValue::Number(serde_json::Number::from(values.len())),
            AggregateOp::Sum => {
                let sum: f64 = values.iter().sum();
                serde_json::to_value(sum).unwrap_or(JsonValue::Null)
            }
            AggregateOp::Avg => {
                if values.is_empty() {
                    JsonValue::Null
                } else {
                    let avg = values.iter().sum::<f64>() / values.len() as f64;
                    serde_json::to_value(avg).unwrap_or(JsonValue::Null)
                }
            }
            AggregateOp::Min => {
                values
                    .iter()
                    .copied()
                    .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|v| serde_json::to_value(v).unwrap_or(JsonValue::Null))
                    .unwrap_or(JsonValue::Null)
            }
            AggregateOp::Max => {
                values
                    .iter()
                    .copied()
                    .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|v| serde_json::to_value(v).unwrap_or(JsonValue::Null))
                    .unwrap_or(JsonValue::Null)
            }
        };

        Ok(result)
    }
}

/// Aggregation operations.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum AggregateOp {
    Count,
    Sum,
    Avg,
    Min,
    Max,
}

/// Cosine similarity between two vectors.
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() {
        return 0.0;
    }

    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let mag_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let mag_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if mag_a == 0.0 || mag_b == 0.0 {
        0.0
    } else {
        dot / (mag_a * mag_b)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::DocumentStore;
    use crate::index::IndexManager;
    use crate::vector::VectorStore;
    use serde_json::json;

    fn setup_engine() -> QueryEngine {
        let docs = Arc::new(DocumentStore::new());
        let vectors = Arc::new(VectorStore::new(3));
        let indexes = Arc::new(IndexManager::new());
        QueryEngine::new(docs, vectors, indexes)
    }

    #[test]
    fn test_simple_query() {
        let docs = Arc::new(DocumentStore::new());
        docs.create_collection("users").unwrap();
        docs.insert("users", json!({"name": "Alice", "age": 30})).unwrap();
        docs.insert("users", json!({"name": "Bob", "age": 25})).unwrap();

        let vectors = Arc::new(VectorStore::new(3));
        let indexes = Arc::new(IndexManager::new());
        let engine = QueryEngine::new(docs, vectors, indexes);

        let query = Query::new("users").limit(10);
        let result = engine.execute(&query).unwrap();

        assert_eq!(result.documents.len(), 2);
        assert_eq!(result.total, 2);
    }

    #[test]
    fn test_filter_query() {
        let docs = Arc::new(DocumentStore::new());
        docs.create_collection("users").unwrap();
        docs.insert("users", json!({"name": "Alice", "age": 30})).unwrap();
        docs.insert("users", json!({"name": "Bob", "age": 25})).unwrap();

        let vectors = Arc::new(VectorStore::new(3));
        let indexes = Arc::new(IndexManager::new());
        let engine = QueryEngine::new(docs, vectors, indexes);

        let query = Query::new("users").filter(Filter::Eq("name".to_string(), json!("Alice")));
        let result = engine.execute(&query).unwrap();

        assert_eq!(result.documents.len(), 1);
        assert_eq!(result.documents[0].get("name"), Some(&json!("Alice")));
    }

    #[test]
    fn test_sorted_query() {
        let docs = Arc::new(DocumentStore::new());
        docs.create_collection("users").unwrap();
        docs.insert("users", json!({"name": "Alice", "age": 30})).unwrap();
        docs.insert("users", json!({"name": "Bob", "age": 25})).unwrap();
        docs.insert("users", json!({"name": "Charlie", "age": 35})).unwrap();

        let vectors = Arc::new(VectorStore::new(3));
        let indexes = Arc::new(IndexManager::new());
        let engine = QueryEngine::new(docs, vectors, indexes);

        let query = Query::new("users").sort("age", SortOrder::Ascending);
        let result = engine.execute(&query).unwrap();

        assert_eq!(result.documents[0].get("name"), Some(&json!("Bob")));
        assert_eq!(result.documents[2].get("name"), Some(&json!("Charlie")));
    }

    #[test]
    fn test_cosine_similarity() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);

        let c = vec![0.0, 1.0, 0.0];
        assert!((cosine_similarity(&a, &c) - 0.0).abs() < 0.001);
    }
}
