luckdb/

lib.rs

#![allow(warnings)]
#![doc = include_str!("../README.md")]

/// A lightweight JSON document database written in Rust.
///
/// LuckDB provides a document-oriented interface for storing, querying, and managing JSON documents
/// with features like indexing, aggregation, persistence, and AES-256 encryption.
///
/// ## Overview
///
/// The main components of LuckDB are:
///
/// - `Client`: The primary entry point for interacting with databases
/// - `Database`: Manages collections of documents
/// - `Collection`: Stores and manipulates documents
/// - `Document`: A JSON document (alias for `serde_json::Value`)
/// - `Query`: Used to filter and find documents
/// - `Index`: Improves query performance
/// - `config::DatabaseConfig`: Configuration management for storage, encryption, and authentication
/// - `encryption::AesEncryption`: AES-256 encryption for data protection
///
/// ## Basic Usage
///
/// ```no_run
/// use luckdb::{Client, config::DatabaseConfig};
/// use serde_json::json;
///
/// // Create a configuration with storage and encryption
/// let config = DatabaseConfig::with_storage_path("./data")
///     .with_encryption("my_secure_password");
///
/// let mut client = Client::with_config(config).unwrap();
///
/// // Access a database
/// let db = client.db("mydb");
///
/// // Access a collection
/// let coll = db.collection("mycollection");
///
/// // Insert a document
/// let doc = json!({"name": "John", "age": 30});
/// let id = coll.insert(doc).unwrap();
///
/// // Find documents
/// let query = luckdb::Query::new().eq("name", "John");
/// let results = coll.find(query, None).unwrap();
///
/// // Save changes to disk (encrypted)
/// client.save().unwrap();
/// ```
///
/// ## Encryption
///
/// LuckDB supports AES-256 encryption for data persistence:
///
/// ```no_run
/// use luckdb::{Client, config::DatabaseConfig};
///
/// let config = DatabaseConfig::new()
///     .with_encryption("my_password")
///     .with_storage_path("./secure_data");
///
/// let client = Client::with_config(config).unwrap();
/// // All data will be encrypted when saved to disk
/// ```
///
/// ## Server Mode
///
/// LuckDB supports client-server mode with authentication:
///
/// ```no_run
/// use luckdb::{Server, config::DatabaseConfig};
///
/// let config = DatabaseConfig::with_storage_path("./server_data")
///     .with_auth("admin", "password")
///     .with_encryption("server_secret");
///
/// let mut server = Server::with_config("127.0.0.1:27017".parse().unwrap(), config)?;
/// server.start()?;
/// ```
///
/// ## Configuration Files
///
/// LuckDB supports configuration via TOML files for easy deployment and environment management:
///
/// ```no_run
/// use luckdb::{Client, config::DatabaseConfig};
///
/// // Load configuration from file
/// let mut client = Client::with_config_file("config.toml")?;
///
/// // Or load config and modify it
/// let mut config = DatabaseConfig::load_from_file("config.toml")?;
/// config.encryption_enabled = true;
/// config.encryption_password = Some("new_password".to_string());
/// let mut client = Client::with_config(config)?;
/// ```
///
/// A typical `config.toml` file looks like:
/// ```toml
/// [luckdb]
/// storage_path = "./data"
/// encryption_enabled = true
/// encryption_password = "your_password"
/// auth_username = "admin"
/// auth_password = "auth_password"
/// auto_save_interval_seconds = 300
/// max_backup_files = 10
/// ```
use chrono;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value as Document, Value};
use std::clone;
use std::cmp::Ordering;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::{self, File};
use std::io::{BufRead, BufReader, BufWriter};
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::ops::Bound::{Excluded, Included, Unbounded};
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};
use std::thread;
use thiserror::Error;
use uuid::Uuid;

// Import new modules
pub mod config;
pub mod encryption;

// --- Error Types ---
#[derive(Error, Debug)]
/// Error types returned by database operations.
pub enum DbError {
    #[error("Document not found")]
    /// A requested document could not be found.
    NotFound,
    #[error("Invalid query: {0}")]
    /// A query was malformed or contained invalid parameters.
    InvalidQuery(String),
    #[error("Serialization error: {0}")]
    /// Error during JSON serialization or deserialization.
    Serialization(#[from] serde_json::Error),
    #[error("IO error: {0}")]
    /// File system or IO-related error.
    IoError(#[from] std::io::Error),
    #[error("Other error: {0}")]
    /// Generic error not covered by other categories.
    Other(String),
    #[error("Index error: {0}")]
    /// Error during index creation or management.
    IndexError(String),
    #[error("Update error: {0}")]
    /// Error during document update operation.
    UpdateError(String),
    #[error("Validation error: {0}")]
    /// Data validation error.
    ValidationError(String),
    #[error("Connection error: {0}")]
    /// Error establishing or maintaining a connection.
    ConnectionError(String),
    #[error("Protocol error: {0}")]
    /// Error in the communication protocol.
    ProtocolError(String),
    #[error("Encryption error: {0}")]
    /// Error during encryption/decryption operations.
    EncryptionError(String),
    #[error("Key management error: {0}")]
    /// Error in encryption key management.
    KeyManagementError(String),
    #[error("Configuration error: {0}")]
    /// Error in database configuration.
    ConfigurationError(String),
}
pub type Result<T> = std::result::Result<T, DbError>;

// --- Document ID ---
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct DocId(String);
impl DocId {
    pub fn new() -> Self {
        Self(Uuid::new_v4().to_string())
    }

    pub fn from_str(s: &str) -> Self {
        Self(s.to_string())
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}
impl std::fmt::Display for DocId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

// --- Query Operators ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryOperator {
    Eq(Value),
    Ne(Value),
    Gt(Value),
    Gte(Value),
    Lt(Value),
    Lte(Value),
    In(Vec<Value>),
    Nin(Vec<Value>),
    Exists(bool),
    Regex(String),
    And(Vec<Query>),
    Or(Vec<Query>),
    Nor(Vec<Query>),
    Not(Box<QueryOperator>),
    // Array operators
    All(Vec<Value>),
    ElemMatch(Query),
    Size(usize),
    // Geospatial operators (simplified)
    Near {
        point: (f64, f64),
        max_distance: Option<f64>,
    },
    Within {
        shape: GeoShape,
    },
    Intersects {
        shape: GeoShape,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GeoShape {
    Box {
        bottom_left: (f64, f64),
        top_right: (f64, f64),
    },
    Polygon {
        points: Vec<(f64, f64)>,
    },
    Center {
        center: (f64, f64),
        radius: f64,
    },
}

// --- Query Condition ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryCondition {
    field: String,
    operator: QueryOperator,
}

// --- Query ---
#[derive(Debug, Clone, Serialize, Deserialize)]
/// Represents a query used to filter documents in a collection.
///
/// The Query struct allows you to build complex queries using method chaining
/// and various operators. It supports equality checks, comparisons, logical operations,
/// array operations, and geospatial queries.
///
/// # Example
///
/// ```
/// use luckdb::{Query, Value};
///
/// // Create a query for documents where age > 18 and status is "active"
/// let query = Query::new()
///     .gt("age", Value::Number(18.0))
///     .eq("status", Value::String("active".to_string()));
///
/// // Create a complex query with logical operators
/// let complex_query = Query::new()
///     .eq("department", Value::String("engineering".to_string()))
///     .and(vec![
///         Query::new().gte("salary", Value::Number(50000.0)),
///         Query::new().lt("years_experience", Value::Number(5.0))
///     ]);
/// ```
pub struct Query {
    /// The list of conditions that documents must satisfy.
    conditions: Vec<QueryCondition>,
    /// The list of logical operations to apply to nested queries.
    logical_ops: Vec<(LogicalOp, Vec<Query>)>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LogicalOp {
    And,
    Or,
    Nor,
    Not,
}
// Add this helper function after the Query implementation
impl Query {
    /// Creates a query from a Value object, typically used for deserialization.
    ///
    /// This method converts a Value (usually an Object) into a Query structure
    /// by parsing operator syntax like `{ "field": { "$gt": 10 } }`.
    ///
    /// # Arguments
    ///
    /// * `value` - The Value to convert into a Query. Must be an Object.
    ///
    /// # Returns
    ///
    /// A Result containing the parsed Query or an InvalidQuery error.
    pub fn from_value(value: Value) -> Result<Self> {
        match value {
            Value::Object(map) => {
                let mut query = Query::new();

                for (field, op_value) in map {
                    match op_value {
                        Value::Object(op_map) => {
                            for (op, val) in op_map {
                                match op.as_str() {
                                    "$eq" => query = query.eq(&field, val.clone()),
                                    "$ne" => query = query.ne(&field, val.clone()),
                                    "$gt" => query = query.gt(&field, val.clone()),
                                    "$gte" => query = query.gte(&field, val.clone()),
                                    "$lt" => query = query.lt(&field, val.clone()),
                                    "$lte" => query = query.lte(&field, val.clone()),
                                    "$in" => {
                                        if let Value::Array(arr) = val {
                                            query = query.in_(&field, arr.clone());
                                        }
                                    }
                                    "$nin" => {
                                        if let Value::Array(arr) = val {
                                            query = query.nin(&field, arr.clone());
                                        }
                                    }
                                    "$exists" => {
                                        if let Value::Bool(exists) = val {
                                            query = query.exists(&field, exists);
                                        }
                                    }
                                    "$regex" => {
                                        if let Value::String(pattern) = val {
                                            query = query.regex(&field, &pattern);
                                        }
                                    }
                                    _ => {
                                        return Err(DbError::InvalidQuery(format!(
                                            "Unknown operator: {}",
                                            op
                                        )))
                                    }
                                }
                            }
                        }
                        _ => query = query.eq(&field, op_value.clone()),
                    }
                }

                Ok(query)
            }
            _ => Err(DbError::InvalidQuery("Query must be an object".to_string())),
        }
    }
}

impl Query {
    /// Creates a new empty query.
    ///
    /// This is typically used as the starting point for building complex queries using method chaining.
    pub fn new() -> Self {
        Self {
            conditions: Vec::new(),
            logical_ops: Vec::new(),
        }
    }

    /// Adds an equality condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check for equality.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn eq(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Eq(value),
        });
        self
    }

    /// Adds a not equal condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check for inequality.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn ne(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Ne(value),
        });
        self
    }

    /// Adds a greater than condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn gt(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Gt(value),
        });
        self
    }

    /// Adds a greater than or equal condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn gte(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Gte(value),
        });
        self
    }

    /// Adds a less than condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn lt(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Lt(value),
        });
        self
    }

    /// Adds a less than or equal condition to the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `value` - The value to compare against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn lte(mut self, key: &str, value: Value) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Lte(value),
        });
        self
    }

    /// Adds an "in" condition to the query.
    ///
    /// Matches documents where the field's value is in the specified list of values.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `values` - The list of values to check against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn in_(mut self, key: &str, values: Vec<Value>) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::In(values),
        });
        self
    }

    /// Adds a "not in" condition to the query.
    ///
    /// Matches documents where the field's value is not in the specified list of values.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `values` - The list of values to check against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn nin(mut self, key: &str, values: Vec<Value>) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Nin(values),
        });
        self
    }

    /// Adds an "exists" condition to the query.
    ///
    /// Matches documents where the field exists or doesn't exist as specified.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check for existence.
    /// * `exists` - Whether the field should exist (true) or not (false).
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn exists(mut self, key: &str, exists: bool) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Exists(exists),
        });
        self
    }

    /// Adds a regular expression condition to the query.
    ///
    /// Matches documents where the field's value matches the specified regular expression.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check.
    /// * `pattern` - The regular expression pattern to match against.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn regex(mut self, key: &str, pattern: &str) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Regex(pattern.to_string()),
        });
        self
    }

    /// Adds a logical AND operation to the query.
    ///
    /// All specified sub-queries must match for the document to be included.
    ///
    /// # Arguments
    ///
    /// * `queries` - A list of queries that must all match.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn and(mut self, queries: Vec<Query>) -> Self {
        self.logical_ops.push((LogicalOp::And, queries));
        self
    }

    /// Adds a logical OR operation to the query.
    ///
    /// At least one of the specified sub-queries must match for the document to be included.
    ///
    /// # Arguments
    ///
    /// * `queries` - A list of queries where at least one must match.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn or(mut self, queries: Vec<Query>) -> Self {
        self.logical_ops.push((LogicalOp::Or, queries));
        self
    }

    /// Adds a logical NOR operation to the query.
    ///
    /// None of the specified sub-queries must match for the document to be included.
    ///
    /// # Arguments
    ///
    /// * `queries` - A list of queries that must all not match.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn nor(mut self, queries: Vec<Query>) -> Self {
        self.logical_ops.push((LogicalOp::Nor, queries));
        self
    }

    /// Adds a logical NOT operation to the query.
    ///
    /// The specified condition must not match for the document to be included.
    /// Note: This implementation only works with simple single-condition queries.
    ///
    /// # Arguments
    ///
    /// * `query` - The query that must not match.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn not(mut self, query: Query) -> Self {
        // Convert the query to a single condition for Not operator
        if let Some(cond) = query.conditions.first() {
            self.conditions.push(QueryCondition {
                field: cond.field.clone(),
                operator: QueryOperator::Not(Box::new(cond.operator.clone())),
            });
        }
        self
    }

    /// Array operator: Matches arrays that contain all elements specified in the query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must be an array).
    /// * `values` - The list of values that the array must contain.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn all(mut self, key: &str, values: Vec<Value>) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::All(values),
        });
        self
    }

    /// Array operator: Matches arrays that contain at least one element that matches the specified query.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must be an array).
    /// * `query` - The query that at least one array element must match.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn elem_match(mut self, key: &str, query: Query) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::ElemMatch(query),
        });
        self
    }

    /// Array operator: Matches arrays that have the specified size.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must be an array).
    /// * `size` - The exact size the array must have.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn size(mut self, key: &str, size: usize) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Size(size),
        });
        self
    }

    /// Geospatial operator: Matches geospatial points that are within a certain distance of a point.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must contain geospatial data).
    /// * `point` - The reference point as a tuple (longitude, latitude).
    /// * `max_distance` - Optional maximum distance from the point.
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn near(mut self, key: &str, point: (f64, f64), max_distance: Option<f64>) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Near {
                point,
                max_distance,
            },
        });
        self
    }

    /// Geospatial operator: Matches geospatial points that are within the specified shape.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must contain geospatial data).
    /// * `shape` - The shape to check against (Box, Polygon, or Center).
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn within(mut self, key: &str, shape: GeoShape) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Within { shape },
        });
        self
    }

    /// Geospatial operator: Matches geospatial shapes that intersect with the specified shape.
    ///
    /// # Arguments
    ///
    /// * `key` - The field to check (must contain geospatial data).
    /// * `shape` - The shape to check against (Box, Polygon, or Center).
    ///
    /// # Returns
    ///
    /// The modified query for method chaining.
    pub fn intersects(mut self, key: &str, shape: GeoShape) -> Self {
        self.conditions.push(QueryCondition {
            field: key.to_string(),
            operator: QueryOperator::Intersects { shape },
        });
        self
    }

    /// Checks if a document matches all conditions and logical operations of the query.
    ///
    /// This method is used internally by the database to determine if a document should be included
    /// in query results. It checks both the direct conditions and any nested logical operations.
    ///
    /// # Arguments
    ///
    /// * `doc` - The document to check against the query.
    ///
    /// # Returns
    ///
    /// `true` if the document matches the query, `false` otherwise.
    pub fn matches(&self, doc: &Document) -> bool {
        // Check all conditions
        for cond in &self.conditions {
            if !self.matches_condition(doc, cond) {
                return false;
            }
        }

        // Check logical operations
        for (op, queries) in &self.logical_ops {
            match op {
                LogicalOp::And => {
                    if !queries.iter().all(|q| q.matches(doc)) {
                        return false;
                    }
                }
                LogicalOp::Or => {
                    if !queries.iter().any(|q| q.matches(doc)) {
                        return false;
                    }
                }
                LogicalOp::Nor => {
                    if queries.iter().any(|q| q.matches(doc)) {
                        return false;
                    }
                }
                LogicalOp::Not => {
                    if queries.iter().any(|q| q.matches(doc)) {
                        return false;
                    }
                }
            }
        }

        true
    }

    fn matches_condition(&self, doc: &Document, cond: &QueryCondition) -> bool {
        let field_parts: Vec<&str> = cond.field.split('.').collect();
        let value = self.get_nested_value(doc, &field_parts);

        match &cond.operator {
            QueryOperator::Eq(expected) => match value {
                Some(v) => v == expected,
                None => false,
            },
            QueryOperator::Ne(expected) => match value {
                Some(v) => v != expected,
                None => true,
            },
            QueryOperator::Gt(expected) => match (value, expected) {
                (Some(Value::Number(a)), Value::Number(b)) => {
                    a.as_f64().unwrap_or(0.0) > b.as_f64().unwrap_or(0.0)
                }
                _ => false,
            },
            QueryOperator::Gte(expected) => match (value, expected) {
                (Some(Value::Number(a)), Value::Number(b)) => {
                    a.as_f64().unwrap_or(0.0) >= b.as_f64().unwrap_or(0.0)
                }
                _ => false,
            },
            QueryOperator::Lt(expected) => match (value, expected) {
                (Some(Value::Number(a)), Value::Number(b)) => {
                    a.as_f64().unwrap_or(0.0) < b.as_f64().unwrap_or(0.0)
                }
                _ => false,
            },
            QueryOperator::Lte(expected) => match (value, expected) {
                (Some(Value::Number(a)), Value::Number(b)) => {
                    a.as_f64().unwrap_or(0.0) <= b.as_f64().unwrap_or(0.0)
                }
                _ => false,
            },
            QueryOperator::In(values) => match value {
                Some(v) => values.contains(v),
                None => false,
            },
            QueryOperator::Nin(values) => match value {
                Some(v) => !values.contains(v),
                None => true,
            },
            QueryOperator::Exists(exists) => value.is_some() == *exists,
            QueryOperator::Regex(pattern) => {
                match value {
                    Some(Value::String(s)) => {
                        // Simple regex matching (in a real implementation, use regex crate)
                        s.contains(pattern)
                    }
                    _ => false,
                }
            }
            QueryOperator::All(values) => match value {
                Some(Value::Array(arr)) => values.iter().all(|v| arr.contains(v)),
                _ => false,
            },
            QueryOperator::ElemMatch(query) => match value {
                Some(Value::Array(arr)) => arr.iter().any(|elem| query.matches(elem)),
                _ => false,
            },
            QueryOperator::Size(size) => match value {
                Some(Value::Array(arr)) => arr.len() == *size,
                _ => false,
            },
            QueryOperator::Near {
                point,
                max_distance,
            } => {
                // Simplified geospatial query
                match value {
                    Some(Value::Array(arr)) if arr.len() >= 2 => {
                        if let (Some(Value::Number(x)), Some(Value::Number(y))) =
                            (arr.get(0), arr.get(1))
                        {
                            let doc_x = x.as_f64().unwrap_or(0.0);
                            let doc_y = y.as_f64().unwrap_or(0.0);

                            // Calculate distance (simplified Euclidean distance)
                            let distance =
                                ((doc_x - point.0).powi(2) + (doc_y - point.1).powi(2)).sqrt();

                            if let Some(max) = max_distance {
                                distance <= *max
                            } else {
                                true
                            }
                        } else {
                            false
                        }
                    }
                    _ => false,
                }
            }
            QueryOperator::Within { shape } => {
                // Simplified implementation
                match (value, shape) {
                    (
                        Some(Value::Array(arr)),
                        GeoShape::Box {
                            bottom_left,
                            top_right,
                        },
                    ) if arr.len() >= 2 => {
                        if let (Some(Value::Number(x)), Some(Value::Number(y))) =
                            (arr.get(0), arr.get(1))
                        {
                            let doc_x = x.as_f64().unwrap_or(0.0);
                            let doc_y = y.as_f64().unwrap_or(0.0);

                            doc_x >= bottom_left.0
                                && doc_x <= top_right.0
                                && doc_y >= bottom_left.1
                                && doc_y <= top_right.1
                        } else {
                            false
                        }
                    }
                    _ => false,
                }
            }
            QueryOperator::Intersects { shape } => {
                // Simplified implementation
                match (value, shape) {
                    (
                        Some(Value::Array(arr)),
                        GeoShape::Box {
                            bottom_left,
                            top_right,
                        },
                    ) if arr.len() >= 2 => {
                        if let (Some(Value::Number(x)), Some(Value::Number(y))) =
                            (arr.get(0), arr.get(1))
                        {
                            let doc_x = x.as_f64().unwrap_or(0.0);
                            let doc_y = y.as_f64().unwrap_or(0.0);

                            doc_x >= bottom_left.0
                                && doc_x <= top_right.0
                                && doc_y >= bottom_left.1
                                && doc_y <= top_right.1
                        } else {
                            false
                        }
                    }
                    _ => false,
                }
            }
            QueryOperator::And(_)
            | QueryOperator::Or(_)
            | QueryOperator::Nor(_)
            | QueryOperator::Not(_) => {
                // These are handled at the query level
                true
            }
        }
    }

    fn get_nested_value<'a>(&self, doc: &'a Value, path: &[&str]) -> Option<&'a Value> {
        if path.is_empty() {
            return Some(doc);
        }

        let key = path[0];
        let rest = &path[1..];

        match doc {
            Value::Object(map) => {
                if let Some(next_value) = map.get(key) {
                    self.get_nested_value(next_value, rest)
                } else {
                    None
                }
            }
            _ => None,
        }
    }
}

// --- Update Operators ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum UpdateOperator {
    Set(Map<String, Value>),
    Unset(Vec<String>),
    Inc(Map<String, Value>),
    Mul(Map<String, Value>),
    Rename(Vec<(String, String)>),
    SetOnInsert(Map<String, Value>),
    Min(Map<String, Value>),
    Max(Map<String, Value>),
    CurrentDate(Map<String, Value>),
    // Array operators
    Push(Map<String, Value>),
    PushAll(BTreeMap<String, Vec<Value>>),
    AddToSet(Map<String, Value>),
    Pop(BTreeMap<String, i64>),
    Pull(Map<String, Value>),
    PullAll(BTreeMap<String, Vec<Value>>),
    Bit(Map<String, Value>),
}

// --- Update Document ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateDocument {
    operators: Vec<UpdateOperator>,
}

impl UpdateDocument {
    pub fn new() -> Self {
        Self {
            operators: Vec::new(),
        }
    }

    pub fn set(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Set(map));
        self
    }

    pub fn unset(mut self, key: &str) -> Self {
        self.operators
            .push(UpdateOperator::Unset(vec![key.to_string()]));
        self
    }

    pub fn inc(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Inc(map));
        self
    }

    pub fn mul(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Mul(map));
        self
    }

    pub fn rename(mut self, old_key: &str, new_key: &str) -> Self {
        self.operators.push(UpdateOperator::Rename(vec![(
            old_key.to_string(),
            new_key.to_string(),
        )]));
        self
    }

    pub fn set_on_insert(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::SetOnInsert(map));
        self
    }

    pub fn min(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Min(map));
        self
    }

    pub fn max(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Max(map));
        self
    }

    pub fn current_date(mut self, key: &str, type_spec: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), type_spec);
        self.operators.push(UpdateOperator::CurrentDate(map));
        self
    }

    // Array operators
    pub fn push(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::Push(map));
        self
    }

    pub fn push_all(mut self, key: &str, values: Vec<Value>) -> Self {
        let mut btree_map = BTreeMap::new();
        btree_map.insert(key.to_string(), values);
        self.operators.push(UpdateOperator::PushAll(btree_map));
        self
    }

    pub fn add_to_set(mut self, key: &str, value: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), value);
        self.operators.push(UpdateOperator::AddToSet(map));
        self
    }

    pub fn pop(mut self, key: &str, pos: i64) -> Self {
        let mut btree_map = BTreeMap::new();
        btree_map.insert(key.to_string(), pos);
        self.operators.push(UpdateOperator::Pop(btree_map));
        self
    }

    pub fn pull(mut self, key: &str, condition: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), condition);
        self.operators.push(UpdateOperator::Pull(map));
        self
    }

    pub fn pull_all(mut self, key: &str, values: Vec<Value>) -> Self {
        let mut btree_map = BTreeMap::new();
        btree_map.insert(key.to_string(), values);
        self.operators.push(UpdateOperator::PullAll(btree_map));
        self
    }

    // Bitwise operator
    pub fn bit(mut self, key: &str, operation: Value) -> Self {
        let mut map = Map::new();
        map.insert(key.to_string(), operation);
        self.operators.push(UpdateOperator::Bit(map));
        self
    }
}

// --- Index ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum IndexType {
    Ascending,
    Descending,
    Geospatial,
    Text,
    Hashed,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Index {
    name: String,
    key: Vec<(String, IndexType)>,
    unique: bool,
    sparse: bool,
    background: bool,
}

impl Index {
    pub fn new(name: String, key: Vec<(String, IndexType)>) -> Self {
        Self {
            name,
            key,
            unique: false,
            sparse: false,
            background: false,
        }
    }

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

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

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

// --- Index Entry ---
#[derive(Debug, Clone, Serialize, Deserialize)]
struct IndexEntry {
    doc_id: DocId,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct OrdValue(Value);

impl PartialEq for OrdValue {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl Eq for OrdValue {}

impl PartialOrd for OrdValue {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for OrdValue {
    fn cmp(&self, other: &Self) -> Ordering {
        // Simple string-based comparison; adjust as needed
        format!("{:?}", self.0).cmp(&format!("{:?}", other.0))
    }
}

// --- Index Data Structure ---
#[derive(Debug, Serialize, Clone, Deserialize)]
struct IndexData {
    entries: BTreeMap<OrdValue, Vec<IndexEntry>>,
    unique: bool,
    sparse: bool,
}

impl IndexData {
    fn new(unique: bool, sparse: bool) -> Self {
        Self {
            entries: BTreeMap::new(),
            unique,
            sparse,
        }
    }

    fn insert(&mut self, value: Value, doc_id: DocId) -> Result<()> {
        let ord_value = OrdValue(value.clone());

        if self.unique && self.entries.contains_key(&ord_value) {
            return Err(DbError::IndexError(format!(
                "Duplicate key error: {:?}",
                value
            )));
        }

        self.entries
            .entry(ord_value)
            .or_insert_with(Vec::new)
            .push(IndexEntry { doc_id });
        Ok(())
    }

    fn remove(&mut self, value: &Value, doc_id: &DocId) -> Result<()> {
        let ord_value = OrdValue(value.clone());

        if let Some(entries) = self.entries.get_mut(&ord_value) {
            entries.retain(|entry| entry.doc_id != *doc_id);
            if entries.is_empty() {
                self.entries.remove(&ord_value);
            }
        }
        Ok(())
    }

    fn find(&self, value: &Value) -> Vec<&IndexEntry> {
        let ord_value = OrdValue(value.clone());
        self.entries
            .get(&ord_value)
            .map_or(Vec::new(), |entries| entries.iter().collect())
    }

    fn find_range(&self, min: &Value, max: &Value) -> Vec<&IndexEntry> {
        let mut result = Vec::new();
        let ord_min = OrdValue(min.clone());
        let ord_max = OrdValue(max.clone());

        for (key, entries) in self.entries.range((Included(&ord_min), Included(&ord_max))) {
            result.extend(entries.iter());
        }
        result
    }
}

// --- Find Options ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FindOptions {
    pub projection: Option<Map<String, Value>>,
    pub sort: Option<Vec<(String, SortOrder)>>,
    pub skip: Option<usize>,
    pub limit: Option<usize>,
    pub max_time_ms: Option<u64>,
    pub allow_partial_results: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SortOrder {
    Ascending,
    Descending,
}

impl Default for FindOptions {
    fn default() -> Self {
        Self {
            projection: None,
            sort: None,
            skip: None,
            limit: None,
            max_time_ms: None,
            allow_partial_results: false,
        }
    }
}

// --- Collection ---
#[derive(Debug, Clone, Serialize, Deserialize)]
/// A collection of documents in the database.
///
/// Collections are analogous to tables in relational databases and store JSON documents.
/// They provide methods for inserting, querying, updating, and deleting documents.
///
/// # Example
///
/// ```no_run
/// use luckdb::{Client, Document};
/// use serde_json::json;
///
/// let mut client = Client::new("mongodb://localhost");
/// let db = client.db("mydb");
/// let coll = db.collection("users");
///
/// // Insert a document
/// let user = json!({"name": "Alice", "email": "alice@example.com"});
/// coll.insert(user).unwrap();
/// ```
pub struct Collection {
    /// The name of the collection.
    name: String,
    /// The documents stored in the collection, indexed by their document ID.
    docs: HashMap<DocId, Document>,
    /// Internal index data structures for efficient querying.
    indexes: HashMap<String, IndexData>,
    /// Definitions of indexes created on the collection.
    index_definitions: HashMap<String, Index>,
}

// --- Persistent Storage ---
#[derive(Debug, Serialize, Deserialize)]
pub struct PersistentCollection {
    name: String,
    docs: HashMap<DocId, Document>,
    indexes: HashMap<String, IndexData>,
    index_definitions: HashMap<String, Index>,
}

impl From<Collection> for PersistentCollection {
    fn from(collection: Collection) -> Self {
        Self {
            name: collection.name,
            docs: collection.docs,
            indexes: collection.indexes,
            index_definitions: collection.index_definitions,
        }
    }
}

impl From<PersistentCollection> for Collection {
    fn from(persistent: PersistentCollection) -> Self {
        Self {
            name: persistent.name,
            docs: persistent.docs,
            indexes: persistent.indexes,
            index_definitions: persistent.index_definitions,
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct PersistentDatabase {
    name: String,
    collections: HashMap<String, PersistentCollection>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct PersistentStorage {
    databases: HashMap<String, PersistentDatabase>,
}

impl PersistentStorage {
    pub fn new() -> Self {
        Self {
            databases: HashMap::new(),
        }
    }

    pub fn save_to_file(
        &self,
        path: &Path,
        encryption_key: Option<&crate::encryption::EncryptionKey>,
    ) -> Result<()> {
        let data = serde_json::to_string_pretty(self)?;

        let final_data = if let Some(key) = encryption_key {
            crate::encryption::AesEncryption::encrypt_string(&data, key)?
        } else {
            data
        };

        std::fs::write(path, final_data)?;
        Ok(())
    }

    pub fn load_from_file(
        path: &Path,
        encryption_key: Option<&crate::encryption::EncryptionKey>,
    ) -> Result<Self> {
        if !path.exists() {
            return Ok(Self::new());
        }

        let mut file = File::open(path)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;

        let final_data = if let Some(key) = encryption_key {
            // Try to decrypt as encrypted data first
            match crate::encryption::AesEncryption::decrypt_string(&contents, key) {
                Ok(decrypted) => decrypted,
                Err(_) => {
                    // If decryption fails, try loading as plain text (for backward compatibility)
                    contents
                }
            }
        } else {
            contents
        };

        let storage = serde_json::from_str(&final_data)?;
        Ok(storage)
    }
}

// --- Collection ---
impl Collection {
    /// Creates a new empty collection with the given name.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the new collection.
    pub fn new(name: String) -> Self {
        Self {
            name,
            docs: HashMap::new(),
            indexes: HashMap::new(),
            index_definitions: HashMap::new(),
        }
    }

    /// Inserts a single document into the collection, auto-generating an `_id` field if not present.
    ///
    /// # Arguments
    ///
    /// * `doc` - The document to insert. Must be a JSON object.
    ///
    /// # Returns
    ///
    /// The ID of the inserted document.
    ///
    /// # Errors
    ///
    /// Returns an error if the document is not a JSON object.
    pub fn insert(&mut self, mut doc: Document) -> Result<DocId> {
        let id = DocId::new();

        // Inject _id field
        if let Some(obj) = doc.as_object_mut() {
            obj.insert("_id".to_string(), Value::String(id.to_string()));
        } else {
            return Err(DbError::Other("Document must be an object".into()));
        }

        self.docs.insert(id.clone(), doc.clone());

        // Update indexes
        self.update_indexes_on_insert(&id, &doc)?;

        Ok(id)
    }

    /// Inserts multiple documents into the collection, auto-generating `_id` fields for each.
    ///
    /// # Arguments
    ///
    /// * `docs` - A list of documents to insert. Each must be a JSON object.
    ///
    /// # Returns
    ///
    /// A list of IDs of the inserted documents.
    ///
    /// # Errors
    ///
    /// Returns an error if any document is not a JSON object.
    pub fn insert_many(&mut self, docs: Vec<Document>) -> Result<Vec<DocId>> {
        let mut ids = Vec::new();

        for mut doc in docs {
            let id = DocId::new();

            // Inject _id field
            if let Some(obj) = doc.as_object_mut() {
                obj.insert("_id".to_string(), Value::String(id.to_string()));
            } else {
                return Err(DbError::Other("Document must be an object".into()));
            }

            self.docs.insert(id.clone(), doc.clone());
            ids.push(id.clone());

            // Update indexes
            self.update_indexes_on_insert(&id, &doc)?;
        }

        Ok(ids)
    }

    /// Finds all documents matching the given query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    /// * `options` - Optional parameters for sorting, pagination, and projection.
    ///
    /// # Returns
    ///
    /// A list of tuples containing document IDs and the matching documents.
    pub fn find(
        &self,
        query: Query,
        options: Option<FindOptions>,
    ) -> Result<Vec<(DocId, Document)>> {
        let mut results = Vec::new();

        // Try to use indexes if possible
        let indexed_results = self.try_indexed_query(&query)?;

        if !indexed_results.is_empty() {
            // Use indexed results
            for (id, doc) in indexed_results {
                if query.matches(&doc) {
                    results.push((id, doc));
                }
            }
        } else {
            // Full collection scan
            for (id, doc) in self.docs.iter() {
                if query.matches(doc) {
                    results.push((id.clone(), doc.clone()));
                }
            }
        }

        // Apply options
        if let Some(opts) = options {
            // Sort
            if let Some(sort_spec) = opts.sort {
                self.sort_results(&mut results, &sort_spec);
            }

            // Skip
            if let Some(skip) = opts.skip {
                if skip < results.len() {
                    results.drain(0..skip);
                } else {
                    results.clear();
                }
            }

            // Limit
            if let Some(limit) = opts.limit {
                results.truncate(limit);
            }

            // Projection
            if let Some(proj) = opts.projection {
                self.apply_projection(&mut results, &proj);
            }
        }

        Ok(results)
    }

    /// Finds the first document matching the given query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    /// * `options` - Optional parameters for sorting and projection.
    ///
    /// # Returns
    ///
    /// A tuple containing the document ID and the matching document.
    ///
    /// # Errors
    ///
    /// Returns `DbError::NotFound` if no document matches the query.
    pub fn find_one(
        &self,
        query: Query,
        options: Option<FindOptions>,
    ) -> Result<(DocId, Document)> {
        let mut opts = options.unwrap_or_default();
        opts.limit = Some(1);

        let mut results = self.find(query, Some(opts))?;
        results.pop().ok_or(DbError::NotFound)
    }

    /// Finds a document, updates it, and returns the document (either before or after the update).
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    /// * `update` - The update operations to apply.
    /// * `options` - Optional parameters including upsert behavior and which document version to return.
    ///
    /// # Returns
    ///
    /// The document before or after the update, depending on options.
    pub fn find_one_and_update(
        &mut self,
        query: Query,
        update: UpdateDocument,
        options: Option<FindOneAndUpdateOptions>,
    ) -> Result<Option<(DocId, Document)>> {
        let opts = options.unwrap_or_default();

        // Find the document
        let result = self.find_one(query.clone(), None);

        match result {
            Ok((id, mut doc)) => {
                let original_doc = doc.clone();

                // Apply update
                self.apply_update(&mut doc, &update, opts.upsert)?;

                // Update the document
                self.docs.insert(id.clone(), doc.clone());

                // Update indexes
                self.update_indexes_on_update(&id, &original_doc, &doc)?;

                if opts.return_document == ReturnDocument::After {
                    Ok(Some((id, doc)))
                } else {
                    Ok(Some((id, original_doc)))
                }
            }
            Err(DbError::NotFound) if opts.upsert => {
                // Upsert: create a new document
                let mut new_doc = Map::new();

                // Extract fields from query
                for cond in &query.conditions {
                    if let QueryOperator::Eq(value) = &cond.operator {
                        new_doc.insert(cond.field.clone(), value.clone());
                    }
                }

                // Apply update
                let mut doc = Value::Object(new_doc);
                self.apply_update(&mut doc, &update, true)?;

                // Insert the document
                let id = self.insert(doc)?;

                // Return the new document if requested
                if opts.return_document == ReturnDocument::After {
                    self.find_one(Query::new().eq("_id", Value::String(id.to_string())), None)
                        .map(Some)
                } else {
                    Ok(None)
                }
            }
            Err(_) => Ok(None),
        }
    }

    // Find and replace a document
    pub fn find_one_and_replace(
        &mut self,
        query: Query,
        replacement: Document,
        options: Option<FindOneAndReplaceOptions>,
    ) -> Result<Option<(DocId, Document)>> {
        let opts = options.unwrap_or_default();

        // Find the document
        let result = self.find_one(query.clone(), None);

        match result {
            Ok((id, original_doc)) => {
                // Ensure replacement has _id
                let mut replacement = replacement;
                if let Some(obj) = replacement.as_object_mut() {
                    obj.insert("_id".to_string(), Value::String(id.to_string()));
                } else {
                    return Err(DbError::Other("Replacement must be an object".into()));
                }

                // Update the document
                self.docs.insert(id.clone(), replacement.clone());

                // Update indexes
                self.update_indexes_on_update(&id, &original_doc, &replacement)?;

                if opts.return_document == ReturnDocument::After {
                    Ok(Some((id, replacement)))
                } else {
                    Ok(Some((id, original_doc)))
                }
            }
            Err(DbError::NotFound) if opts.upsert => {
                // Upsert: create a new document
                let id = self.insert(replacement)?;

                // Return the new document if requested
                if opts.return_document == ReturnDocument::After {
                    self.find_one(Query::new().eq("_id", Value::String(id.to_string())), None)
                        .map(Some)
                } else {
                    Ok(None)
                }
            }
            Err(_) => Ok(None),
        }
    }

    // Find and delete a document
    pub fn find_one_and_delete(
        &mut self,
        query: Query,
        options: Option<FindOneAndDeleteOptions>,
    ) -> Result<Option<(DocId, Document)>> {
        let opts = options.unwrap_or_default();

        // Find the document
        let result = self.find_one(query.clone(), None);

        match result {
            Ok((id, doc)) => {
                // Delete the document
                self.docs.remove(&id);

                // Update indexes
                self.update_indexes_on_delete(&id, &doc)?;

                Ok(Some((id, doc)))
            }
            Err(_) => Ok(None),
        }
    }

    /// Updates the first document matching the query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    /// * `update` - The update operations to apply.
    /// * `upsert` - If true, creates a new document if no match is found.
    ///
    /// # Returns
    ///
    /// The number of documents updated (either 0 or 1).
    pub fn update_one(
        &mut self,
        query: Query,
        update: UpdateDocument,
        upsert: bool,
    ) -> Result<usize> {
        let result = self.find_one(query.clone(), None);

        match result {
            Ok((id, mut doc)) => {
                let original_doc = doc.clone();

                // Apply update
                self.apply_update(&mut doc, &update, upsert)?;

                // Update the document
                self.docs.insert(id.clone(), doc.clone());

                // Update indexes
                self.update_indexes_on_update(&id, &original_doc, &doc)?;

                Ok(1)
            }
            Err(DbError::NotFound) if upsert => {
                // Upsert: create a new document
                let mut new_doc = Map::new();

                // Extract fields from query
                for cond in &query.conditions {
                    if let QueryOperator::Eq(value) = &cond.operator {
                        new_doc.insert(cond.field.clone(), value.clone());
                    }
                }

                // Apply update
                let mut doc = Value::Object(new_doc);
                self.apply_update(&mut doc, &update, true)?;

                // Insert the document
                self.insert(doc)?;

                Ok(1)
            }
            Err(_) => Ok(0),
        }
    }

    /// Updates all documents matching the query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    /// * `update` - The update operations to apply.
    ///
    /// # Returns
    ///
    /// The number of documents updated.
    pub fn update_many(&mut self, query: Query, update: UpdateDocument) -> Result<usize> {
        let docs = self.find(query, None)?;
        let mut count = 0;

        for (id, doc) in docs {
            let mut doc = doc;
            let original_doc = doc.clone();

            // Apply update
            self.apply_update(&mut doc, &update, false)?;

            // Update the document
            self.docs.insert(id.clone(), doc.clone());

            // Update indexes
            self.update_indexes_on_update(&id, &original_doc, &doc)?;

            count += 1;
        }

        Ok(count)
    }

    // Replace the first matching document
    pub fn replace_one(
        &mut self,
        query: Query,
        replacement: Document,
        upsert: bool,
    ) -> Result<usize> {
        let result = self.find_one(query.clone(), None);

        match result {
            Ok((id, original_doc)) => {
                // Ensure replacement has _id
                let mut replacement = replacement;
                if let Some(obj) = replacement.as_object_mut() {
                    obj.insert("_id".to_string(), Value::String(id.to_string()));
                } else {
                    return Err(DbError::Other("Replacement must be an object".into()));
                }

                // Update the document
                self.docs.insert(id.clone(), replacement.clone());

                // Update indexes
                self.update_indexes_on_update(&id, &original_doc, &replacement)?;

                Ok(1)
            }
            Err(DbError::NotFound) if upsert => {
                // Upsert: create a new document
                self.insert(replacement)?;
                Ok(1)
            }
            Err(_) => Ok(0),
        }
    }

    /// Deletes the first document matching the query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    ///
    /// # Returns
    ///
    /// The number of documents deleted (either 0 or 1).
    pub fn delete_one(&mut self, query: Query) -> Result<usize> {
        let result = self.find_one(query, None);

        match result {
            Ok((id, doc)) => {
                // Delete the document
                self.docs.remove(&id);

                // Update indexes
                self.update_indexes_on_delete(&id, &doc)?;

                Ok(1)
            }
            Err(_) => Ok(0),
        }
    }

    /// Deletes all documents matching the query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    ///
    /// # Returns
    ///
    /// The number of documents deleted.
    pub fn delete_many(&mut self, query: Query) -> Result<usize> {
        let docs = self.find(query, None)?;
        let mut count = 0;

        for (id, doc) in docs {
            // Delete the document
            self.docs.remove(&id);

            // Update indexes
            self.update_indexes_on_delete(&id, &doc)?;

            count += 1;
        }

        Ok(count)
    }

    /// Counts the number of documents matching the query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query to filter documents.
    ///
    /// # Returns
    ///
    /// The number of matching documents.
    pub fn count_documents(&self, query: Query) -> Result<usize> {
        // Try to use indexes if possible
        let indexed_results = self.try_indexed_query(&query)?;

        if !indexed_results.is_empty() {
            // Use indexed results
            let mut count = 0;
            for (_, doc) in indexed_results {
                if query.matches(&doc) {
                    count += 1;
                }
            }
            Ok(count)
        } else {
            // Full collection scan
            let mut count = 0;
            for doc in self.docs.values() {
                if query.matches(doc) {
                    count += 1;
                }
            }
            Ok(count)
        }
    }

    /// Returns the estimated count of documents in the collection.
    ///
    /// # Returns
    ///
    /// The total number of documents in the collection.
    pub fn estimated_document_count(&self) -> Result<usize> {
        Ok(self.docs.len())
    }

    // Create an index
    /// Creates an index on the collection.
    ///
    /// # Arguments
    ///
    /// * `index` - The index definition containing field, type, uniqueness, and other properties.
    ///
    /// # Errors
    ///
    /// Returns an error if an index with the same name already exists or if there's a problem creating the index.
    pub fn create_index(&mut self, index: Index) -> Result<()> {
        let index_name = index.name.clone();

        // Check if index already exists
        if self.index_definitions.contains_key(&index_name) {
            return Err(DbError::IndexError(format!(
                "Index {} already exists",
                index_name
            )));
        }

        // Create index data
        let index_data = IndexData::new(index.unique, index.sparse);

        // Add to indexes
        self.indexes.insert(index_name.clone(), index_data);

        // Add to definitions
        self.index_definitions
            .insert(index_name.clone(), index.clone());

        // Collect document IDs and contents for indexing
        let docs_to_index: Vec<(DocId, Document)> = self
            .docs
            .iter()
            .map(|(id, doc)| (id.clone(), doc.clone()))
            .collect();

        // Build index from existing documents
        for (id, doc) in docs_to_index {
            self.add_to_indexes(&index_name, &id, &doc)?;
        }

        Ok(())
    }

    /// Drops an existing index from the collection.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the index to drop.
    pub fn drop_index(&mut self, name: &str) -> Result<()> {
        self.indexes.remove(name);
        self.index_definitions.remove(name);
        Ok(())
    }

    /// Lists all indexes created on the collection.
    ///
    /// # Returns
    ///
    /// A list of index definitions.
    pub fn list_indexes(&self) -> Result<Vec<Index>> {
        Ok(self.index_definitions.values().cloned().collect())
    }

    /// Drops all indexes from the collection.
    pub fn drop_indexes(&mut self) -> Result<()> {
        self.indexes.clear();
        self.index_definitions.clear();
        Ok(())
    }

    // Get collection statistics
    pub fn stats(&self) -> Result<CollectionStats> {
        Ok(CollectionStats {
            count: self.docs.len(),
            size: self.calculate_size(),
            avg_obj_size: if self.docs.is_empty() {
                0
            } else {
                self.calculate_size() / self.docs.len()
            },
            index_count: self.index_definitions.len(),
            index_size: self.calculate_index_size(),
        })
    }

    // Aggregate pipeline
    pub fn aggregate(&self, pipeline: Vec<AggregationStage>) -> Result<Vec<Document>> {
        let mut results = Vec::new();

        // Start with all documents
        let mut docs: Vec<Document> = self.docs.values().cloned().collect();

        // Apply pipeline stages
        for stage in pipeline {
            docs = self.apply_aggregation_stage(docs, stage)?;
        }

        results.extend(docs);
        Ok(results)
    }

    // Distinct values for a field
    pub fn distinct(&self, field: &str, query: Option<Query>) -> Result<Vec<Value>> {
        let mut values = HashSet::new();

        for doc in self.docs.values() {
            if let Some(q) = &query {
                if !q.matches(doc) {
                    continue;
                }
            }

            let field_parts: Vec<&str> = field.split('.').collect();
            if let Some(value) = self.get_nested_value(doc, &field_parts) {
                values.insert(value.clone());
            }
        }

        Ok(values.into_iter().collect())
    }

    // Get the number of documents in the collection
    pub fn count(&self) -> usize {
        self.docs.len()
    }

    // --- Helper methods ---

    fn apply_update(
        &self,
        doc: &mut Document,
        update: &UpdateDocument,
        is_upsert: bool,
    ) -> Result<()> {
        for op in &update.operators {
            match op {
                UpdateOperator::Set(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            obj.insert(key.clone(), value.clone());
                        }
                    }
                }
                UpdateOperator::Unset(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for key in fields {
                            obj.remove(key);
                        }
                    }
                }
                UpdateOperator::Inc(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(Value::Number(n)) = obj.get(key) {
                                if let (Some(current), Some(inc)) = (n.as_f64(), value.as_f64()) {
                                    obj.insert(
                                        key.clone(),
                                        Value::Number(
                                            serde_json::Number::from_f64(current + inc).unwrap(),
                                        ),
                                    );
                                }
                            } else {
                                // Field doesn't exist or isn't a number
                                if let Some(inc) = value.as_f64() {
                                    obj.insert(
                                        key.clone(),
                                        Value::Number(serde_json::Number::from_f64(inc).unwrap()),
                                    );
                                }
                            }
                        }
                    }
                }
                UpdateOperator::Mul(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(Value::Number(n)) = obj.get(key) {
                                if let (Some(current), Some(mul)) = (n.as_f64(), value.as_f64()) {
                                    obj.insert(
                                        key.clone(),
                                        Value::Number(
                                            serde_json::Number::from_f64(current * mul).unwrap(),
                                        ),
                                    );
                                }
                            }
                        }
                    }
                }
                UpdateOperator::Rename(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (old_key, new_key) in fields {
                            if let Some(value) = obj.remove(old_key) {
                                obj.insert(new_key.clone(), value);
                            }
                        }
                    }
                }
                UpdateOperator::SetOnInsert(fields) => {
                    if is_upsert {
                        if let Some(obj) = doc.as_object_mut() {
                            for (key, value) in fields {
                                obj.insert(key.clone(), value.clone());
                            }
                        }
                    }
                }
                UpdateOperator::Min(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(current) = obj.get(key) {
                                // Compare values (simplified)
                                if self.compare_values(value, current) == Ordering::Less {
                                    obj.insert(key.clone(), value.clone());
                                }
                            } else {
                                obj.insert(key.clone(), value.clone());
                            }
                        }
                    }
                }
                UpdateOperator::Max(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(current) = obj.get(key) {
                                // Compare values (simplified)
                                if self.compare_values(value, current) == Ordering::Greater {
                                    obj.insert(key.clone(), value.clone());
                                }
                            } else {
                                obj.insert(key.clone(), value.clone());
                            }
                        }
                    }
                }
                UpdateOperator::CurrentDate(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, type_spec) in fields {
                            let now = chrono::Utc::now();

                            match type_spec {
                                Value::String(type_str) if type_str == "date" => {
                                    obj.insert(key.clone(), Value::String(now.to_rfc3339()));
                                }
                                Value::String(type_str) if type_str == "timestamp" => {
                                    obj.insert(
                                        key.clone(),
                                        Value::Number(serde_json::Number::from(now.timestamp())),
                                    );
                                }
                                Value::Object(spec) => {
                                    if let Some(Value::String(type_str)) = spec.get("$type") {
                                        match type_str.as_str() {
                                            "date" => {
                                                obj.insert(
                                                    key.clone(),
                                                    Value::String(now.to_rfc3339()),
                                                );
                                            }
                                            "timestamp" => {
                                                obj.insert(
                                                    key.clone(),
                                                    Value::Number(serde_json::Number::from(
                                                        now.timestamp(),
                                                    )),
                                                );
                                            }
                                            _ => {
                                                return Err(DbError::UpdateError(format!(
                                                    "Invalid date type specification: {}",
                                                    type_str
                                                )));
                                            }
                                        }
                                    }
                                }
                                _ => {
                                    return Err(DbError::UpdateError(
                                        "Invalid date type specification".into(),
                                    ));
                                }
                            }
                        }
                    }
                }
                UpdateOperator::Push(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                arr.push(value.clone());
                            } else {
                                // Field doesn't exist or isn't an array
                                obj.insert(key.clone(), Value::Array(vec![value.clone()]));
                            }
                        }
                    }
                }
                UpdateOperator::PushAll(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, values) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                arr.extend(values.iter().cloned());
                            } else {
                                // Field doesn't exist or isn't an array
                                obj.insert(key.clone(), Value::Array(values.clone()));
                            }
                        }
                    }
                }
                UpdateOperator::AddToSet(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, value) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                if !arr.contains(value) {
                                    arr.push(value.clone());
                                }
                            } else {
                                // Field doesn't exist or isn't an array
                                obj.insert(key.clone(), Value::Array(vec![value.clone()]));
                            }
                        }
                    }
                }
                UpdateOperator::Pop(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, pos) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                if *pos == 1 {
                                    arr.pop();
                                } else if *pos == -1 {
                                    arr.remove(0);
                                }
                            }
                        }
                    }
                }
                UpdateOperator::Pull(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, condition) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                // Simplified: remove all elements matching the condition
                                arr.retain(|elem| {
                                    match condition {
                                        Value::Object(cond) => {
                                            // Create a query from the condition
                                            let query = Query::new();
                                            // This is a simplified version
                                            !query.matches(elem)
                                        }
                                        _ => elem != condition,
                                    }
                                });
                            }
                        }
                    }
                }
                UpdateOperator::PullAll(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, values) in fields {
                            if let Some(Value::Array(arr)) = obj.get_mut(key) {
                                arr.retain(|elem| !values.contains(elem));
                            }
                        }
                    }
                }
                UpdateOperator::Bit(fields) => {
                    if let Some(obj) = doc.as_object_mut() {
                        for (key, operation) in fields {
                            if let Value::Object(op) = operation {
                                if let Some(Value::Number(current)) = obj.get(key) {
                                    let mut current_num = current.as_i64().unwrap_or(0);

                                    if let Some(Value::Number(and)) = op.get("and") {
                                        current_num &= and.as_i64().unwrap_or(0);
                                    }
                                    if let Some(Value::Number(or)) = op.get("or") {
                                        current_num |= or.as_i64().unwrap_or(0);
                                    }
                                    if let Some(Value::Number(xor)) = op.get("xor") {
                                        current_num ^= xor.as_i64().unwrap_or(0);
                                    }

                                    obj.insert(
                                        key.clone(),
                                        Value::Number(serde_json::Number::from(current_num)),
                                    );
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    fn compare_values(&self, a: &Value, b: &Value) -> Ordering {
        match (a, b) {
            (Value::Number(a_num), Value::Number(b_num)) => a_num
                .as_f64()
                .unwrap_or(0.0)
                .partial_cmp(&b_num.as_f64().unwrap_or(0.0))
                .unwrap_or(Ordering::Equal),
            (Value::String(a_str), Value::String(b_str)) => a_str.cmp(b_str),
            (Value::Bool(a_bool), Value::Bool(b_bool)) => a_bool.cmp(b_bool),
            (Value::Array(a_arr), Value::Array(b_arr)) => a_arr.len().cmp(&b_arr.len()),
            (Value::Object(_), Value::Object(_)) => {
                Ordering::Equal // Simplified
            }
            _ => Ordering::Equal,
        }
    }

    fn update_indexes_on_insert(&mut self, id: &DocId, doc: &Document) -> Result<()> {
        for (name, index_def) in self.index_definitions.iter() {
            // Pre-calculate all field values before mutable borrow
            let mut field_values = Vec::new();
            for (field, _) in &index_def.key {
                let field_parts: Vec<&str> = field.split('.').collect();
                let value = self.get_nested_value(doc, &field_parts);
                field_values.push((field, value));
            }

            if let Some(index_data) = self.indexes.get_mut(name) {
                for (field, value) in field_values {
                    if let Some(value) = value {
                        index_data.insert(value.clone(), id.clone())?;
                    } else if index_data.sparse {
                        // Skip missing fields for sparse indexes
                        continue;
                    } else {
                        return Err(DbError::IndexError(format!(
                            "Field {} not found in document for non-sparse index",
                            field
                        )));
                    }
                }
            }
        }

        Ok(())
    }

    fn update_indexes_on_update(
        &mut self,
        id: &DocId,
        old_doc: &Document,
        new_doc: &Document,
    ) -> Result<()> {
        // Remove from indexes
        self.update_indexes_on_delete(id, old_doc)?;

        // Add to indexes
        self.update_indexes_on_insert(id, new_doc)?;

        Ok(())
    }
    // Fix update_indexes_on_delete method
    fn update_indexes_on_delete(&mut self, id: &DocId, doc: &Document) -> Result<()> {
        let index_names: Vec<String> = self.index_definitions.keys().cloned().collect();

        for name in index_names {
            if let Some(index_def) = self.index_definitions.get(&name) {
                // Extract all field values first
                let mut field_values = Vec::new();
                for (field, _) in &index_def.key {
                    let field_parts: Vec<&str> = field.split('.').collect();
                    let value = self.get_nested_value(doc, &field_parts);
                    field_values.push(value.cloned());
                }

                // Now get mutable borrow and process
                if let Some(index_data) = self.indexes.get_mut(&name) {
                    for (i, (_, _)) in index_def.key.iter().enumerate() {
                        if let Some(val) = &field_values[i] {
                            index_data.remove(val, id)?;
                        }
                    }
                }
            }
        }

        Ok(())
    }

    // Fix add_to_indexes method
    fn add_to_indexes(&mut self, index_name: &str, id: &DocId, doc: &Document) -> Result<()> {
        if let Some(index_def) = self.index_definitions.get(index_name).cloned() {
            // Extract all field values first
            let mut field_values = Vec::new();
            let field_names: Vec<String> = index_def
                .key
                .iter()
                .map(|(field, _)| field.clone())
                .collect();

            for field in &field_names {
                let field_parts: Vec<&str> = field.split('.').collect();
                let value = self.get_nested_value(doc, &field_parts);
                field_values.push(value.cloned());
            }

            // Now get mutable borrow and process
            if let Some(index_data) = self.indexes.get_mut(index_name) {
                for (i, field) in field_names.iter().enumerate() {
                    if let Some(val) = &field_values[i] {
                        index_data.insert(val.clone(), id.clone())?;
                    } else if index_data.sparse {
                        // Skip missing fields for sparse indexes
                        continue;
                    } else {
                        return Err(DbError::IndexError(format!(
                            "Field {} not found in document for non-sparse index",
                            field
                        )));
                    }
                }
            }
        }

        Ok(())
    }

    fn try_indexed_query(&self, query: &Query) -> Result<Vec<(DocId, Document)>> {
        // Try to find an index that can be used for the query
        for (index_name, index_def) in self.index_definitions.iter() {
            // Check if the first field in the index is used in the query
            if let Some((field, _)) = index_def.key.first() {
                for cond in &query.conditions {
                    if cond.field == *field {
                        // We can use this index
                        let mut results = Vec::new();

                        if let Some(index_data) = self.indexes.get(index_name) {
                            match &cond.operator {
                                QueryOperator::Eq(value) => {
                                    let entries = index_data.find(value);
                                    for entry in entries {
                                        if let Some(doc) = self.docs.get(&entry.doc_id) {
                                            results.push((entry.doc_id.clone(), doc.clone()));
                                        }
                                    }
                                }
                                QueryOperator::Gt(value) => {
                                    // Find all entries greater than the value
                                    let ord_value = OrdValue(value.clone());
                                    for (key, entries) in
                                        index_data.entries.range((Excluded(&ord_value), Unbounded))
                                    {
                                        for entry in entries {
                                            if let Some(doc) = self.docs.get(&entry.doc_id) {
                                                results.push((entry.doc_id.clone(), doc.clone()));
                                            }
                                        }
                                    }
                                }
                                QueryOperator::Gte(value) => {
                                    // Find all entries greater than or equal to the value
                                    let ord_value = OrdValue(value.clone());
                                    for (key, entries) in
                                        index_data.entries.range((Included(&ord_value), Unbounded))
                                    {
                                        for entry in entries {
                                            if let Some(doc) = self.docs.get(&entry.doc_id) {
                                                results.push((entry.doc_id.clone(), doc.clone()));
                                            }
                                        }
                                    }
                                }
                                QueryOperator::Lt(value) => {
                                    // Find all entries less than the value
                                    let ord_value = OrdValue(value.clone());
                                    for (key, entries) in
                                        index_data.entries.range((Unbounded, Excluded(&ord_value)))
                                    {
                                        for entry in entries {
                                            if let Some(doc) = self.docs.get(&entry.doc_id) {
                                                results.push((entry.doc_id.clone(), doc.clone()));
                                            }
                                        }
                                    }
                                }
                                QueryOperator::Lte(value) => {
                                    // Find all entries less than or equal to the value
                                    let ord_value = OrdValue(value.clone());
                                    for (key, entries) in
                                        index_data.entries.range((Unbounded, Included(&ord_value)))
                                    {
                                        for entry in entries {
                                            if let Some(doc) = self.docs.get(&entry.doc_id) {
                                                results.push((entry.doc_id.clone(), doc.clone()));
                                            }
                                        }
                                    }
                                }
                                _ => {
                                    // Can't use index for this operator
                                    continue;
                                }
                            }
                        }

                        return Ok(results);
                    }
                }
            }
        }

        // No suitable index found
        Ok(Vec::new())
    }

    fn sort_results(
        &self,
        results: &mut Vec<(DocId, Document)>,
        sort_spec: &[(String, SortOrder)],
    ) {
        results.sort_by(|a, b| {
            for (field, order) in sort_spec {
                let field_parts: Vec<&str> = field.split('.').collect();
                let a_val = self.get_nested_value(&a.1, &field_parts);
                let b_val = self.get_nested_value(&b.1, &field_parts);

                let cmp = match (a_val, b_val) {
                    (Some(Value::Number(a_num)), Some(Value::Number(b_num))) => a_num
                        .as_f64()
                        .unwrap_or(0.0)
                        .partial_cmp(&b_num.as_f64().unwrap_or(0.0))
                        .unwrap_or(Ordering::Equal),
                    (Some(Value::String(a_str)), Some(Value::String(b_str))) => a_str.cmp(b_str),
                    (Some(Value::Bool(a_bool)), Some(Value::Bool(b_bool))) => a_bool.cmp(b_bool),
                    (Some(Value::Array(a_arr)), Some(Value::Array(b_arr))) => {
                        a_arr.len().cmp(&b_arr.len())
                    }
                    (Some(_), None) => Ordering::Greater,
                    (None, Some(_)) => Ordering::Less,
                    (None, None) => Ordering::Equal,
                    _ => Ordering::Equal,
                };

                if cmp != Ordering::Equal {
                    return match order {
                        SortOrder::Ascending => cmp,
                        SortOrder::Descending => cmp.reverse(),
                    };
                }
            }

            Ordering::Equal
        });
    }

    fn apply_projection(
        &self,
        results: &mut Vec<(DocId, Document)>,
        projection: &Map<String, Value>,
    ) {
        let is_inclusive = projection.values().any(|v| match v {
            Value::Bool(b) => *b,
            Value::Number(n) => n.as_i64().unwrap_or(0) == 1,
            _ => false,
        });

        for (_, doc) in results.iter_mut() {
            if let Some(obj) = doc.as_object_mut() {
                let mut new_obj = Map::new();

                if is_inclusive {
                    // Include only specified fields
                    for (key, value) in projection {
                        if obj.contains_key(key) {
                            new_obj.insert(key.clone(), obj.get(key).unwrap().clone());
                        }
                    }

                    // Always include _id
                    if obj.contains_key("_id") && !projection.contains_key("_id") {
                        new_obj.insert("_id".to_string(), obj.get("_id").unwrap().clone());
                    }
                } else {
                    // Exclude specified fields
                    for (key, value) in obj {
                        if !projection.contains_key(key) || key == "_id" {
                            new_obj.insert(key.clone(), value.clone());
                        }
                    }
                }

                *doc = Value::Object(new_obj);
            }
        }
    }

    fn get_nested_value<'a>(&self, doc: &'a Document, path: &[&str]) -> Option<&'a Value> {
        if path.is_empty() {
            return Some(doc);
        }

        let current = path[0];
        let rest = &path[1..];

        match doc.get(current) {
            Some(value) => {
                if rest.is_empty() {
                    Some(value)
                } else {
                    match value {
                        Value::Object(obj) => self.get_nested_value(value, rest),
                        _ => None,
                    }
                }
            }
            None => None,
        }
    }

    fn calculate_size(&self) -> usize {
        // Simplified size calculation
        self.docs
            .values()
            .map(|doc| serde_json::to_string(doc).unwrap_or_default().len())
            .sum()
    }

    fn calculate_index_size(&self) -> usize {
        // Simplified index size calculation
        self.indexes
            .values()
            .map(|index_data| {
                index_data.entries.len() * 16 // Approximate size per entry
            })
            .sum()
    }

    fn apply_aggregation_stage(
        &self,
        docs: Vec<Document>,
        stage: AggregationStage,
    ) -> Result<Vec<Document>> {
        match stage {
            AggregationStage::Match(query) => {
                let mut result = Vec::new();
                for doc in docs {
                    if query.matches(&doc) {
                        result.push(doc);
                    }
                }
                Ok(result)
            }
            AggregationStage::Project(projection) => {
                let mut result = Vec::new();
                for doc in docs {
                    if let Some(obj) = doc.as_object() {
                        let mut new_obj = Map::new();

                        for (key, value) in &projection {
                            if value.as_object().and_then(|o| o.get("$")).is_some() {
                                // Field projection with expression
                                if let Some(expr) = value.as_object().and_then(|o| o.get("$")) {
                                    // Simplified expression evaluation
                                    if let Some(field_name) = expr.as_str() {
                                        if let Some(field_value) =
                                            obj.get(field_name.trim_start_matches('$'))
                                        {
                                            new_obj.insert(key.clone(), field_value.clone());
                                        }
                                    }
                                }
                            } else if value.as_object().and_then(|o| o.get("$concat")).is_some() {
                                // String concatenation
                                if let Some(Value::Array(fields)) =
                                    value.as_object().and_then(|o| o.get("$concat"))
                                {
                                    let mut result_str = String::new();
                                    for field in fields {
                                        if let Some(field_name) = field.as_str() {
                                            if let Some(field_value) =
                                                obj.get(field_name.trim_start_matches('$'))
                                            {
                                                if let Some(s) = field_value.as_str() {
                                                    result_str.push_str(s);
                                                }
                                            }
                                        }
                                    }
                                    new_obj.insert(key.clone(), Value::String(result_str));
                                }
                            } else {
                                // Simple field inclusion/exclusion
                                let include = match value {
                                    Value::Bool(b) => b,
                                    Value::Number(n) => &(n.as_i64().unwrap_or(0) == 1),
                                    _ => &false,
                                };

                                if *include && obj.contains_key(key.as_str()) {
                                    new_obj.insert(
                                        key.clone(),
                                        obj.get(key.as_str()).unwrap().clone(),
                                    );
                                } else if !include && !new_obj.contains_key(key.as_str()) {
                                    new_obj.insert(
                                        key.clone(),
                                        obj.get(key.as_str()).unwrap().clone(),
                                    );
                                }
                            }
                        }

                        // Always include _id if not explicitly excluded
                        if obj.contains_key("_id") && !projection.contains_key("_id") {
                            new_obj.insert("_id".to_string(), obj.get("_id").unwrap().clone());
                        }

                        result.push(Value::Object(new_obj));
                    }
                }
                Ok(result)
            }
            AggregationStage::Sort(sort_spec) => {
                let mut result = docs;

                result.sort_by(|a, b| {
                    for (field, order) in &sort_spec {
                        let field_parts: Vec<&str> = field.split('.').collect();
                        let a_val = self.get_nested_value(a, &field_parts);
                        let b_val = self.get_nested_value(b, &field_parts);

                        let cmp = match (a_val, b_val) {
                            (Some(Value::Number(a_num)), Some(Value::Number(b_num))) => a_num
                                .as_f64()
                                .unwrap_or(0.0)
                                .partial_cmp(&b_num.as_f64().unwrap_or(0.0))
                                .unwrap_or(Ordering::Equal),
                            (Some(Value::String(a_str)), Some(Value::String(b_str))) => {
                                a_str.cmp(b_str)
                            }
                            (Some(Value::Bool(a_bool)), Some(Value::Bool(b_bool))) => {
                                a_bool.cmp(b_bool)
                            }
                            (Some(Value::Array(a_arr)), Some(Value::Array(b_arr))) => {
                                a_arr.len().cmp(&b_arr.len())
                            }
                            (Some(_), None) => Ordering::Greater,
                            (None, Some(_)) => Ordering::Less,
                            (None, None) => Ordering::Equal,
                            _ => Ordering::Equal,
                        };

                        if cmp != Ordering::Equal {
                            return match order {
                                SortOrder::Ascending => cmp,
                                SortOrder::Descending => cmp.reverse(),
                            };
                        }
                    }

                    Ordering::Equal
                });

                Ok(result)
            }
            AggregationStage::Skip(n) => {
                let mut result = docs;
                if n < result.len() {
                    result.drain(0..n);
                } else {
                    result.clear();
                }
                Ok(result)
            }
            AggregationStage::Limit(n) => {
                let mut result = docs;
                result.truncate(n);
                Ok(result)
            }
            AggregationStage::Group(group_spec) => {
                let mut groups: HashMap<String, Vec<Document>> = HashMap::new();

                // Group documents
                for doc in docs {
                    let group_key = self.calculate_group_key(&doc, &group_spec.id)?;
                    groups.entry(group_key).or_insert_with(Vec::new).push(doc);
                }

                // Apply group operations
                let mut result = Vec::new();
                for (key, group_docs) in groups {
                    let mut group_obj = Map::new();

                    // Add _id field
                    group_obj.insert("_id".to_string(), Value::String(key));

                    // Apply group operations
                    for (field, op) in &group_spec.operations {
                        match op {
                            GroupOperation::Sum(expr) => {
                                let sum = self.calculate_sum(&group_docs, expr)?;
                                group_obj.insert(field.clone(), sum);
                            }
                            GroupOperation::Avg(expr) => {
                                let avg = self.calculate_avg(&group_docs, expr)?;
                                group_obj.insert(field.clone(), avg);
                            }
                            GroupOperation::Min(expr) => {
                                let min = self.calculate_min(&group_docs, expr)?;
                                group_obj.insert(field.clone(), min);
                            }
                            GroupOperation::Max(expr) => {
                                let max = self.calculate_max(&group_docs, expr)?;
                                group_obj.insert(field.clone(), max);
                            }
                            GroupOperation::First(expr) => {
                                if let Some(doc) = group_docs.first() {
                                    if let Some(value) = self.evaluate_expression(doc, expr) {
                                        group_obj.insert(field.clone(), value);
                                    }
                                }
                            }
                            GroupOperation::Last(expr) => {
                                if let Some(doc) = group_docs.last() {
                                    if let Some(value) = self.evaluate_expression(doc, expr) {
                                        group_obj.insert(field.clone(), value);
                                    }
                                }
                            }
                            GroupOperation::Push(expr) => {
                                let mut values = Vec::new();
                                for doc in &group_docs {
                                    if let Some(value) = self.evaluate_expression(doc, expr) {
                                        values.push(value);
                                    }
                                }
                                group_obj.insert(field.clone(), Value::Array(values));
                            }
                            GroupOperation::AddToSet(expr) => {
                                let mut values = HashSet::new();
                                for doc in &group_docs {
                                    if let Some(value) = self.evaluate_expression(doc, expr) {
                                        values.insert(value);
                                    }
                                }
                                group_obj.insert(
                                    field.clone(),
                                    Value::Array(values.into_iter().collect()),
                                );
                            }
                            GroupOperation::StdDevPop(expr) => {
                                let std_dev = self.calculate_std_dev_pop(&group_docs, expr)?;
                                group_obj.insert(field.clone(), std_dev);
                            }
                            GroupOperation::StdDevSamp(expr) => {
                                let std_dev = self.calculate_std_dev_samp(&group_docs, expr)?;
                                group_obj.insert(field.clone(), std_dev);
                            }
                        }
                    }

                    result.push(Value::Object(group_obj));
                }

                Ok(result)
            }
            AggregationStage::Unwind(field) => {
                let mut result = Vec::new();

                for doc in docs {
                    if let Some(obj) = doc.as_object() {
                        if let Some(value) = obj.get(&field) {
                            if let Value::Array(arr) = value {
                                for item in arr {
                                    let mut new_obj = obj.clone();
                                    new_obj.insert(field.clone(), item.clone());
                                    result.push(Value::Object(new_obj));
                                }
                            } else {
                                // Not an array, just include the document as is
                                result.push(doc.clone());
                            }
                        } else {
                            // Field doesn't exist, include the document as is
                            result.push(doc.clone());
                        }
                    }
                }

                Ok(result)
            }
            AggregationStage::Lookup(lookup_spec) => {
                // Simplified implementation of $lookup
                let mut result = Vec::new();

                for doc in docs {
                    if let Some(obj) = doc.as_object() {
                        let mut new_obj = obj.clone();

                        // Get the local field value
                        let local_field_value = obj.get(&lookup_spec.local_field);

                        if let Some(local_value) = local_field_value {
                            // Find matching documents in the foreign collection
                            let query =
                                Query::new().eq(&lookup_spec.foreign_field, local_value.clone());

                            // This is a simplified implementation - in a real DB, we would access the foreign collection
                            new_obj.insert(lookup_spec.as_field.clone(), Value::Array(Vec::new()));
                        } else {
                            // Local field doesn't exist
                            new_obj.insert(lookup_spec.as_field.clone(), Value::Array(Vec::new()));
                        }

                        result.push(Value::Object(new_obj));
                    }
                }

                Ok(result)
            }
            AggregationStage::Out(collection_name) => {
                // Simplified implementation of $out
                // In a real DB, this would write to another collection
                Ok(docs)
            }
        }
    }

    fn calculate_group_key(&self, doc: &Document, group_id: &GroupId) -> Result<String> {
        match group_id {
            GroupId::Field(field) => {
                let field_parts: Vec<&str> = field.split('.').collect();
                if let Some(value) = self.get_nested_value(doc, &field_parts) {
                    Ok(value.to_string())
                } else {
                    Ok("null".to_string())
                }
            }
            GroupId::Expression(expr) => {
                if let Some(value) = self.evaluate_expression(doc, expr) {
                    Ok(value.to_string())
                } else {
                    Ok("null".to_string())
                }
            }
            GroupId::Null => Ok("null".to_string()),
        }
    }

    fn evaluate_expression(&self, doc: &Document, expr: &Value) -> Option<Value> {
        match expr {
            Value::String(s) if s.starts_with('$') => {
                let field = &s[1..];
                let field_parts: Vec<&str> = field.split('.').collect();
                self.get_nested_value(doc, &field_parts).cloned()
            }
            Value::Object(obj) => {
                // Handle expression operators
                if let Some(sum_expr) = obj.get("$sum") {
                    if let Some(field) = sum_expr.as_str() {
                        if field.starts_with('$') {
                            let field_name = &field[1..];
                            let field_parts: Vec<&str> = field_name.split('.').collect();
                            return self.get_nested_value(doc, &field_parts).cloned();
                        }
                    }
                }
                None
            }
            _ => Some(expr.clone()),
        }
    }

    fn calculate_sum(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        let mut sum = 0.0;

        for doc in docs {
            if let Some(value) = self.evaluate_expression(doc, expr) {
                if let Some(num) = value.as_f64() {
                    sum += num;
                }
            }
        }

        Ok(Value::Number(serde_json::Number::from_f64(sum).unwrap()))
    }

    fn calculate_avg(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        if docs.is_empty() {
            return Ok(Value::Number(serde_json::Number::from(0)));
        }

        let sum = self.calculate_sum(docs, expr)?;
        if let Some(sum_num) = sum.as_f64() {
            let avg = sum_num / docs.len() as f64;
            Ok(Value::Number(serde_json::Number::from_f64(avg).unwrap()))
        } else {
            Ok(Value::Number(serde_json::Number::from(0)))
        }
    }

    fn calculate_min(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        let mut min_value: Option<Value> = None;

        for doc in docs {
            if let Some(value) = self.evaluate_expression(doc, expr) {
                if let Some(ref mut min) = min_value {
                    if self.compare_values(&value, min) == Ordering::Less {
                        *min = value;
                    }
                } else {
                    min_value = Some(value);
                }
            }
        }

        Ok(min_value.unwrap_or(Value::Null))
    }

    fn calculate_max(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        let mut max_value: Option<Value> = None;

        for doc in docs {
            if let Some(value) = self.evaluate_expression(doc, expr) {
                if let Some(ref mut max) = max_value {
                    if self.compare_values(&value, max) == Ordering::Greater {
                        *max = value;
                    }
                } else {
                    max_value = Some(value);
                }
            }
        }

        Ok(max_value.unwrap_or(Value::Null))
    }

    fn calculate_std_dev_pop(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        if docs.is_empty() {
            return Ok(Value::Number(serde_json::Number::from(0)));
        }

        let avg = self.calculate_avg(docs, expr)?;
        let avg_num = avg.as_f64().unwrap_or(0.0);

        let mut sum_sq_diff = 0.0;
        for doc in docs {
            if let Some(value) = self.evaluate_expression(doc, expr) {
                if let Some(num) = value.as_f64() {
                    let diff = num - avg_num;
                    sum_sq_diff += diff * diff;
                }
            }
        }

        let variance = sum_sq_diff / docs.len() as f64;
        let std_dev = variance.sqrt();

        Ok(Value::Number(
            serde_json::Number::from_f64(std_dev).unwrap(),
        ))
    }

    fn calculate_std_dev_samp(&self, docs: &[Document], expr: &Value) -> Result<Value> {
        if docs.len() <= 1 {
            return Ok(Value::Number(serde_json::Number::from(0)));
        }

        let avg = self.calculate_avg(docs, expr)?;
        let avg_num = avg.as_f64().unwrap_or(0.0);

        let mut sum_sq_diff = 0.0;
        for doc in docs {
            if let Some(value) = self.evaluate_expression(doc, expr) {
                if let Some(num) = value.as_f64() {
                    let diff = num - avg_num;
                    sum_sq_diff += diff * diff;
                }
            }
        }

        let variance = sum_sq_diff / (docs.len() - 1) as f64;
        let std_dev = variance.sqrt();

        Ok(Value::Number(
            serde_json::Number::from_f64(std_dev).unwrap(),
        ))
    }
}

// --- Aggregation Stages ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AggregationStage {
    Match(Query),
    Project(Map<String, Value>),
    Sort(Vec<(String, SortOrder)>),
    Skip(usize),
    Limit(usize),
    Group(GroupSpecification),
    Unwind(String),
    Lookup(LookupSpecification),
    Out(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GroupSpecification {
    pub id: GroupId,
    pub operations: HashMap<String, GroupOperation>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GroupId {
    Field(String),
    Expression(Value),
    Null,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GroupOperation {
    Sum(Value),
    Avg(Value),
    Min(Value),
    Max(Value),
    First(Value),
    Last(Value),
    Push(Value),
    AddToSet(Value),
    StdDevPop(Value),
    StdDevSamp(Value),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LookupSpecification {
    pub from: String,
    pub local_field: String,
    pub foreign_field: String,
    pub as_field: String, // Renamed from 'as' to avoid using a reserved keyword
}

// --- Find and Modify Options ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FindOneAndUpdateOptions {
    pub upsert: bool,
    pub return_document: ReturnDocument,
    pub projection: Option<Map<String, Value>>,
    pub sort: Option<Vec<(String, SortOrder)>>,
    pub max_time_ms: Option<u64>,
}

impl Default for FindOneAndUpdateOptions {
    fn default() -> Self {
        Self {
            upsert: false,
            return_document: ReturnDocument::Before,
            projection: None,
            sort: None,
            max_time_ms: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FindOneAndReplaceOptions {
    pub upsert: bool,
    pub return_document: ReturnDocument,
    pub projection: Option<Map<String, Value>>,
    pub sort: Option<Vec<(String, SortOrder)>>,
    pub max_time_ms: Option<u64>,
}

impl Default for FindOneAndReplaceOptions {
    fn default() -> Self {
        Self {
            upsert: false,
            return_document: ReturnDocument::Before,
            projection: None,
            sort: None,
            max_time_ms: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FindOneAndDeleteOptions {
    pub projection: Option<Map<String, Value>>,
    pub sort: Option<Vec<(String, SortOrder)>>,
    pub max_time_ms: Option<u64>,
}

impl Default for FindOneAndDeleteOptions {
    fn default() -> Self {
        Self {
            projection: None,
            sort: None,
            max_time_ms: None,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ReturnDocument {
    Before,
    After,
}

// --- Collection Statistics ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectionStats {
    pub count: usize,
    pub size: usize,
    pub avg_obj_size: usize,
    pub index_count: usize,
    pub index_size: usize,
}

// --- Database ---
#[derive(Debug, Clone, Serialize, Deserialize)]
/// A database containing collections of documents.
///
/// A database is a container for collections, which in turn contain documents.
/// It provides methods for creating, accessing, and managing collections.
///
/// # Example
///
/// ```no_run
/// use luckdb::Client;
///
/// let mut client = Client::new("mongodb://localhost");
/// let db = client.db("mydatabase");
///
/// // Access or create a collection
/// let coll = db.collection("mycollection");
/// ```
pub struct Database {
    /// The name of the database.
    name: String,
    /// The collections contained in the database, indexed by name.
    collections: HashMap<String, Collection>,
}

impl Database {
    /// Creates a new empty database with the given name.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the new database.
    pub fn new(name: String) -> Self {
        Self {
            name,
            collections: HashMap::new(),
        }
    }

    /// Gets a collection by name, creating it if it doesn't exist.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the collection to access.
    ///
    /// # Returns
    ///
    /// A mutable reference to the collection.
    pub fn collection(&mut self, name: &str) -> &mut Collection {
        self.collections
            .entry(name.to_string())
            .or_insert_with(|| Collection::new(name.to_string()))
    }

    /// Lists the names of all collections in the database.
    ///
    /// # Returns
    ///
    /// A list of collection names.
    pub fn list_collection_names(&self) -> Vec<String> {
        self.collections.keys().cloned().collect()
    }

    /// Creates a new collection with the specified options.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the new collection.
    /// * `options` - Optional parameters for creating the collection, including indexes.
    ///
    /// # Errors
    ///
    /// Returns an error if a collection with the same name already exists.
    pub fn create_collection(
        &mut self,
        name: &str,
        options: Option<CreateCollectionOptions>,
    ) -> Result<()> {
        if self.collections.contains_key(name) {
            return Err(DbError::Other(format!(
                "Collection {} already exists",
                name
            )));
        }

        let mut collection = Collection::new(name.to_string());

        // Apply options if provided
        if let Some(opts) = options {
            // Create indexes if specified
            for index in opts.indexes {
                collection.create_index(index)?;
            }
        }

        self.collections.insert(name.to_string(), collection);
        Ok(())
    }

    /// Drops a collection from the database.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the collection to drop.
    ///
    /// # Errors
    ///
    /// Returns an error if the collection does not exist.
    pub fn drop_collection(&mut self, name: &str) -> Result<()> {
        if self.collections.remove(name).is_none() {
            return Err(DbError::Other(format!(
                "Collection {} does not exist",
                name
            )));
        }

        Ok(())
    }

    /// Gets statistics about the database.
    ///
    /// # Returns
    ///
    /// Database statistics including collection counts and sizes.
    pub fn stats(&self) -> Result<DatabaseStats> {
        let mut collections = Vec::new();
        let mut total_size = 0;
        let mut total_index_size = 0;

        for (name, collection) in self.collections.iter() {
            let coll_stats = collection.stats()?;
            collections.push(CollectionStatsInfo {
                name: name.clone(),
                count: coll_stats.count,
                size: coll_stats.size,
                index_count: coll_stats.index_count,
                index_size: coll_stats.index_size,
            });

            total_size += coll_stats.size;
            total_index_size += coll_stats.index_size;
        }

        Ok(DatabaseStats {
            collections,
            total_size,
            total_index_size,
        })
    }

    /// Runs a database command.
    ///
    /// Supports commands like "create", "drop", "listCollections", and "dbStats".
    ///
    /// # Arguments
    ///
    /// * `command` - The command to execute, represented as a document.
    ///
    /// # Returns
    ///
    /// The result of the command execution.
    pub fn run_command(&mut self, command: &Document) -> Result<Document> {
        if let Some(obj) = command.as_object() {
            if let Some(cmd_name) = obj.keys().next() {
                match cmd_name.as_str() {
                    "create" => {
                        if let Some(Value::String(coll_name)) = obj.get(cmd_name) {
                            self.create_collection(coll_name, None)?;
                            Ok(serde_json::json!({ "ok": 1 }))
                        } else {
                            Err(DbError::Other("Invalid create command".into()))
                        }
                    }
                    "drop" => {
                        if let Some(Value::String(coll_name)) = obj.get(cmd_name) {
                            self.drop_collection(coll_name)?;
                            Ok(serde_json::json!({ "ok": 1 }))
                        } else {
                            Err(DbError::Other("Invalid drop command".into()))
                        }
                    }
                    "listCollections" => {
                        let coll_names = self.list_collection_names();
                        let collections: Vec<_> = coll_names
                            .into_iter()
                            .map(|name| {
                                serde_json::json!({
                                    "name": name,
                                    "type": "collection"
                                })
                            })
                            .collect();

                        Ok(serde_json::json!({
                            "cursor": {
                                "id": 0,
                                "ns": format!("{}.collections", self.name),
                                "firstBatch": collections
                            },
                            "ok": 1
                        }))
                    }
                    "dbStats" => {
                        let stats = self.stats()?;
                        Ok(serde_json::json!({
                            "db": self.name,
                            "collections": stats.collections.len(),
                            "objects": stats.collections.iter().map(|c| c.count).sum::<usize>(),
                            "avgObjSize": if stats.collections.iter().map(|c| c.count).sum::<usize>() > 0 {
                                stats.total_size / stats.collections.iter().map(|c| c.count).sum::<usize>()
                            } else {
                                0
                            },
                            "dataSize": stats.total_size,
                            "indexSize": stats.total_index_size,
                            "ok": 1
                        }))
                    }
                    _ => Err(DbError::Other(format!("Unknown command: {}", cmd_name))),
                }
            } else {
                Err(DbError::Other("Empty command".into()))
            }
        } else {
            Err(DbError::Other("Command must be an object".into()))
        }
    }
}

// --- Database Statistics ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseStats {
    pub collections: Vec<CollectionStatsInfo>,
    pub total_size: usize,
    pub total_index_size: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectionStatsInfo {
    pub name: String,
    pub count: usize,
    pub size: usize,
    pub index_count: usize,
    pub index_size: usize,
}

// --- Create Collection Options ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateCollectionOptions {
    pub capped: bool,
    pub size: Option<usize>,
    pub max: Option<usize>,
    pub storage_engine: Option<Map<String, Value>>,
    pub validator: Option<Document>,
    pub validation_level: Option<String>,
    pub validation_action: Option<String>,
    pub index_option_defaults: Option<Map<String, Value>>,
    pub view_on: Option<String>,
    pub pipeline: Option<Vec<Document>>,
    pub collation: Option<Map<String, Value>>,
    pub write_concern: Option<Map<String, Value>>,
    pub indexes: Vec<Index>,
}

impl Default for CreateCollectionOptions {
    fn default() -> Self {
        Self {
            capped: false,
            size: None,
            max: None,
            storage_engine: None,
            validator: None,
            validation_level: None,
            validation_action: None,
            index_option_defaults: None,
            view_on: None,
            pipeline: None,
            collation: None,
            write_concern: None,
            indexes: Vec::new(),
        }
    }
}

// --- Client ---
#[derive(Debug, Clone)]
/// The main entry point for interacting with LuckDB databases.
///
/// The Client struct is responsible for managing databases and providing access to them.
/// It also handles persistence operations like saving and loading data to disk.
///
/// # Example
///
/// ```no_run
/// use luckdb::Client;
///
/// // Create a new client with a storage path
/// let mut client = Client::with_storage_path("mongodb://localhost", "/path/to/data");
///
/// // Access or create a database
/// let db = client.db("mydatabase");
///
/// // Save changes to disk
/// client.save().unwrap();
/// ```
pub struct Client {
    /// The databases managed by this client, indexed by name.
    pub databases: HashMap<String, Database>,
    /// The connection URI for the database server.
    pub uri: String,
    /// Database configuration
    pub config: config::DatabaseConfig,
    /// Encryption key for data protection
    pub encryption_key: Option<encryption::EncryptionKey>,
}

impl Client {
    /// Creates a new client with default settings.
    ///
    /// Uses the default configuration with in-memory storage and no encryption.
    /// This is suitable for testing and temporary data storage.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use luckdb::Client;
    ///
    /// let mut client = Client::new();
    /// let db = client.db("test");
    /// ```
    ///
    /// # Note
    ///
    /// Data stored in-memory will be lost when the client is dropped.
    /// Use `with_config()` or `with_storage_path()` for persistent storage.
    pub fn new() -> Self {
        let config = config::DatabaseConfig::default();
        Self {
            databases: HashMap::new(),
            uri: "luckdb://localhost:27017".to_string(), // Updated from mongodb://
            encryption_key: None,
            config,
        }
    }

    /// Creates a new client with a custom connection URI.
    ///
    /// # Arguments
    ///
    /// * `uri` - The connection URI for the database server.
    pub fn with_uri(uri: &str) -> Self {
        let config = config::DatabaseConfig::default();
        Self {
            databases: HashMap::new(),
            uri: uri.to_string(),
            encryption_key: None,
            config,
        }
    }

    /// Creates a new client with a custom connection URI and storage path.
    ///
    /// # Arguments
    ///
    /// * `uri` - The connection URI for the database server.
    /// * `path` - The file path where data will be persisted to disk.
    pub fn with_storage_path<P: AsRef<Path>>(uri: &str, path: P) -> Self {
        let mut config = config::DatabaseConfig::default();
        config.storage_path = Some(path.as_ref().to_path_buf());
        Self {
            databases: HashMap::new(),
            uri: uri.to_string(),
            encryption_key: None,
            config,
        }
    }

    /// Creates a new client with the specified configuration.
    ///
    /// This method allows you to configure storage, encryption, authentication,
    /// and other settings through a `DatabaseConfig` object.
    ///
    /// # Arguments
    ///
    /// * `config` - Database configuration object
    ///
    /// # Returns
    ///
    /// Returns a `Result<Client>` which will be an error if the configuration
    /// validation fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use luckdb::{Client, config::DatabaseConfig};
    ///
    /// // Create configuration with storage and encryption
    /// let config = DatabaseConfig::with_storage_path("./data")
    ///     .with_encryption("secure_password");
    ///
    /// let mut client = Client::with_config(config)?;
    /// let db = client.db("mydb");
    /// ```
    ///
    /// # Errors
    ///
    /// Returns `DbError::ValidationError` if the configuration is invalid.
    pub fn with_config(config: config::DatabaseConfig) -> Result<Self> {
        config.validate()?;

        let encryption_key = if config.is_encryption_configured() {
            Some(encryption::EncryptionKey::from_password(
                config.encryption_password.as_ref().unwrap(),
                b"luckdb_salt_2024", // Default salt for key derivation
            ))
        } else {
            None
        };

        Ok(Self {
            databases: HashMap::new(),
            uri: "luckdb://localhost:27017".to_string(),
            config,
            encryption_key,
        })
    }

    /// Creates a new client by loading configuration from a TOML file.
    ///
    /// This is a convenience method that loads configuration from a TOML file
    /// and creates a client with that configuration. The file must contain
    /// a `[luckdb]` section with valid configuration parameters.
    ///
    /// # Arguments
    ///
    /// * `config_path` - Path to the TOML configuration file
    ///
    /// # Returns
    ///
    /// Returns a `Result<Client>` which will be an error if the file cannot be
    /// read, is invalid TOML, or contains invalid configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use luckdb::Client;
    ///
    /// // Load from config.toml file
    /// let mut client = Client::with_config_file("config.toml")?;
    /// let db = client.db("mydb");
    /// ```
    ///
    /// # Configuration File Format
    ///
    /// ```toml
    /// [luckdb]
    /// storage_path = "./data"
    /// encryption_enabled = true
    /// encryption_password = "your_password"
    /// auth_username = "admin"
    /// auth_password = "auth_password"
    /// ```
    ///
    /// # Errors
    ///
    /// - `DbError::IoError` - If the file cannot be read
    /// - `DbError::TomlError` - If the file contains invalid TOML
    /// - `DbError::ValidationError` - If the configuration is invalid
    pub fn with_config_file<P: AsRef<Path>>(config_path: P) -> Result<Self> {
        let config = config::DatabaseConfig::load_from_file(config_path)?;
        Self::with_config(config)
    }

    /// Enable encryption with a password.
    ///
    /// # Arguments
    ///
    /// * `password` - Password for encryption key derivation.
    pub fn with_encryption<S: Into<String>>(mut self, password: S) -> Result<Self> {
        self.config.encryption_enabled = true;
        self.config.encryption_password = Some(password.into());

        let key = encryption::EncryptionKey::from_password(
            self.config.encryption_password.as_ref().unwrap(),
            b"luckdb_salt_2024",
        );

        self.encryption_key = Some(key);
        self.config.validate()?;
        Ok(self)
    }

    /// Gets a database by name, creating it if it doesn't exist.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the database to access.
    ///
    /// # Returns
    ///
    /// A mutable reference to the database.
    pub fn db(&mut self, name: &str) -> &mut Database {
        self.databases
            .entry(name.to_string())
            .or_insert_with(|| Database::new(name.to_string()))
    }

    /// Lists the names of all databases managed by this client.
    ///
    /// # Returns
    ///
    /// A list of database names.
    pub fn list_database_names(&self) -> Vec<String> {
        self.databases.keys().cloned().collect()
    }

    /// Drops a database managed by this client.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the database to drop.
    ///
    /// # Errors
    ///
    /// Returns an error if the database does not exist.
    pub fn drop_database(&mut self, name: &str) -> Result<()> {
        if self.databases.remove(name).is_none() {
            return Err(DbError::Other(format!("Database {} does not exist", name)));
        }

        Ok(())
    }

    /// Gets the connection URI used by this client.
    ///
    /// # Returns
    ///
    /// A reference to the connection URI string.
    pub fn uri(&self) -> &str {
        &self.uri
    }

    /// Saves all data to disk at the configured storage path.
    ///
    /// # Errors
    ///
    /// Returns an error if no storage path is configured or if there's an IO error.
    pub fn save(&self) -> Result<()> {
        let storage_path = self.config.get_or_create_storage_path()?;

        // Create directory if it doesn't exist
        if let Some(parent) = storage_path.parent() {
            fs::create_dir_all(parent)?;
        }

        let mut persistent_storage = PersistentStorage::new();

        for (name, db) in &self.databases {
            let mut persistent_db = PersistentDatabase {
                name: name.clone(),
                collections: HashMap::new(),
            };

            for (coll_name, coll) in &db.collections {
                let persistent_coll = PersistentCollection::from(coll.clone());
                persistent_db
                    .collections
                    .insert(coll_name.clone(), persistent_coll);
            }

            persistent_storage
                .databases
                .insert(name.clone(), persistent_db);
        }

        // Save with encryption if configured
        persistent_storage.save_to_file(&storage_path, self.encryption_key.as_ref())?;
        Ok(())
    }

    /// Loads all data from disk at the configured storage path.
    ///
    /// # Errors
    ///
    /// Returns an error if no storage path is configured or if there's an IO or deserialization error.
    pub fn load(&mut self) -> Result<()> {
        let storage_path = self.config.get_or_create_storage_path()?;

        if storage_path.exists() {
            let persistent_storage =
                PersistentStorage::load_from_file(&storage_path, self.encryption_key.as_ref())?;

            for (name, persistent_db) in persistent_storage.databases {
                let mut db = Database::new(name.clone());

                for (coll_name, persistent_coll) in persistent_db.collections {
                    let coll = Collection::from(persistent_coll);
                    db.collections.insert(coll_name, coll);
                }

                self.databases.insert(name, db);
            }
        }
        Ok(())
    }

    /// Get configuration reference
    pub fn config(&self) -> &config::DatabaseConfig {
        &self.config
    }

    /// Update configuration
    pub fn update_config(&mut self, config: config::DatabaseConfig) -> Result<()> {
        config.validate()?;
        self.config = config;

        // Update encryption key if needed
        if self.config.is_encryption_configured() {
            self.encryption_key = Some(encryption::EncryptionKey::from_password(
                self.config.encryption_password.as_ref().unwrap(),
                b"luckdb_salt_2024",
            ));
        } else {
            self.encryption_key = None;
        }

        Ok(())
    }
}

// --- Bulk Write Operations ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BulkWriteOperation {
    InsertOne {
        document: Document,
    },
    UpdateOne {
        filter: Query,
        update: UpdateDocument,
        upsert: bool,
    },
    UpdateMany {
        filter: Query,
        update: UpdateDocument,
        upsert: bool,
    },
    ReplaceOne {
        filter: Query,
        replacement: Document,
        upsert: bool,
    },
    DeleteOne {
        filter: Query,
    },
    DeleteMany {
        filter: Query,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BulkWriteOptions {
    pub ordered: bool,
    pub bypass_document_validation: bool,
    pub write_concern: Option<Map<String, Value>>,
}

impl Default for BulkWriteOptions {
    fn default() -> Self {
        Self {
            ordered: true,
            bypass_document_validation: false,
            write_concern: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BulkWriteResult {
    pub inserted_count: usize,
    pub matched_count: usize,
    pub modified_count: usize,
    pub deleted_count: usize,
    pub upserted_count: usize,
    pub upserted_ids: HashMap<usize, DocId>,
}

impl Collection {
    // Execute bulk write operations
    pub fn bulk_write(
        &mut self,
        operations: Vec<BulkWriteOperation>,
        options: Option<BulkWriteOptions>,
    ) -> Result<BulkWriteResult> {
        let opts = options.unwrap_or_default();
        let mut result = BulkWriteResult {
            inserted_count: 0,
            matched_count: 0,
            modified_count: 0,
            deleted_count: 0,
            upserted_count: 0,
            upserted_ids: HashMap::new(),
        };

        let mut index = 0;

        for op in operations {
            match op {
                BulkWriteOperation::InsertOne { document } => {
                    let id = self.insert(document)?;
                    result.inserted_count += 1;
                    result.upserted_ids.insert(index, id);
                }
                BulkWriteOperation::UpdateOne {
                    filter,
                    update,
                    upsert,
                } => {
                    let count = self.update_one(filter, update, upsert)?;
                    if count > 0 {
                        result.matched_count += 1;
                        result.modified_count += 1;
                    }
                    if upsert && count > 0 {
                        result.upserted_count += 1;
                    }
                }
                BulkWriteOperation::UpdateMany {
                    filter,
                    update,
                    upsert,
                } => {
                    let count = self.update_many(filter, update)?;
                    result.matched_count += count;
                    result.modified_count += count;
                    if upsert && count > 0 {
                        result.upserted_count += count;
                    }
                }
                BulkWriteOperation::ReplaceOne {
                    filter,
                    replacement,
                    upsert,
                } => {
                    let count = self.replace_one(filter, replacement, upsert)?;
                    if count > 0 {
                        result.matched_count += 1;
                        result.modified_count += 1;
                    }
                    if upsert && count > 0 {
                        result.upserted_count += 1;
                    }
                }
                BulkWriteOperation::DeleteOne { filter } => {
                    let count = self.delete_one(filter)?;
                    result.deleted_count += count;
                }
                BulkWriteOperation::DeleteMany { filter } => {
                    let count = self.delete_many(filter)?;
                    result.deleted_count += count;
                }
            }

            index += 1;

            // If ordered and an error occurs, we would stop here
            // For simplicity, we're not handling errors in this example
        }

        Ok(result)
    }
}

// --- Remote Client ---
#[derive(Debug, Clone)]
pub struct RemoteClient {
    addr: SocketAddr,
}

impl RemoteClient {
    pub fn new(addr: SocketAddr) -> Self {
        Self { addr }
    }

    pub fn connect(&self) -> Result<RemoteConnection> {
        let stream = TcpStream::connect(self.addr)?;
        Ok(RemoteConnection { stream })
    }
}

#[derive(Debug)]
pub struct RemoteConnection {
    stream: TcpStream,
}

impl RemoteConnection {
    pub fn send_command(&mut self, command: &str) -> Result<String> {
        let mut writer = BufWriter::new(&self.stream);
        writer.write_all(command.as_bytes())?;
        writer.write_all(b"\n")?;
        writer.flush()?;

        let mut reader = BufReader::new(&self.stream);
        let mut response = String::new();
        reader.read_line(&mut response)?;

        Ok(response)
    }

    pub fn close(self) -> Result<()> {
        // The connection will be closed when RemoteConnection is dropped
        Ok(())
    }
}

// --- Server ---
#[derive(Debug)]
pub struct Server {
    addr: SocketAddr,
    client: Client,
}

impl Server {
    /// Creates a new server with the given address and storage configuration
    pub fn new(addr: SocketAddr, storage_path: Option<PathBuf>) -> Self {
        let mut config = config::DatabaseConfig::default();
        config.storage_path = storage_path;

        let client = Client::with_config(config).unwrap();
        Self { addr, client }
    }

    /// Creates a new server with custom configuration
    pub fn with_config(addr: SocketAddr, config: config::DatabaseConfig) -> Result<Self> {
        let client = Client::with_config(config)?;
        Ok(Self { addr, client })
    }

    /// Creates a new server with configuration from file
    pub fn with_config_file<P: AsRef<Path>>(addr: SocketAddr, config_path: P) -> Result<Self> {
        let client = Client::with_config_file(config_path)?;
        Ok(Self { addr, client })
    }

    /// Enable authentication (deprecated - use config instead)
    #[deprecated(note = "Use with_config() instead")]
    pub fn with_auth(mut self, username: String, password: String) -> Self {
        let mut config = self.client.config().clone();
        config.auth_username = Some(username);
        config.auth_password = Some(password);
        let _ = self.client.update_config(config);
        self
    }

    /// Enable encryption on the server
    pub fn with_encryption(mut self, password: String) -> Result<Self> {
        self.client = self.client.with_encryption(password)?;
        Ok(self)
    }

    pub fn start(&mut self) -> Result<()> {
        // Load data from disk if storage path is configured
        if self.client.config().storage_path.is_some() {
            self.client.load()?;
        }
        let listener = TcpListener::bind(self.addr)?;
        println!("Server listening on {}", self.addr);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    let config = self.client.config().clone();
                    let encryption_key = self.client.encryption_key.clone();
                    thread::spawn(move || {
                        if let Err(e) = Self::handle_client(stream, config, encryption_key) {
                            eprintln!("Error handling client: {}", e);
                        }
                    });
                }
                Err(e) => {
                    eprintln!("Failed to accept connection: {}", e);
                }
            }
        }
        Ok(())
    }

    fn handle_client(
        mut stream: TcpStream,
        config: config::DatabaseConfig,
        encryption_key: Option<encryption::EncryptionKey>,
    ) -> Result<()> {
        let mut client = Client::with_config(config)?;
        client.encryption_key = encryption_key;

        let mut reader = BufReader::new(&stream);
        let mut writer = BufWriter::new(&stream);
        let mut authenticated = !client.config().is_auth_configured(); // If no auth is configured, client is authenticated by default

        loop {
            let mut command = String::new();
            match reader.read_line(&mut command) {
                Ok(0) => break, // Connection closed
                Ok(_) => {
                    // Trim whitespace and newline
                    let command = command.trim();
                    if command.is_empty() {
                        continue;
                    }

                    // Handle authentication if required
                    if !authenticated {
                        if command.starts_with("AUTH") {
                            let parts: Vec<&str> = command.split_whitespace().collect();
                            if parts.len() != 3 {
                                let response =
                                    "ERROR: Usage: AUTH <username> <password>".to_string();
                                writer.write_all(response.as_bytes())?;
                                writer.write_all(b"\n")?;
                                writer.flush()?;
                                continue;
                            }

                            let username = parts[1];
                            let password = parts[2];

                            if let (Some(auth_user), Some(auth_pass)) = (
                                &client.config().auth_username,
                                &client.config().auth_password,
                            ) {
                                if username == auth_user && password == auth_pass {
                                    authenticated = true;
                                    let response = "OK: Authenticated".to_string();
                                    writer.write_all(response.as_bytes())?;
                                    writer.write_all(b"\n")?;
                                    writer.flush()?;
                                    continue;
                                }
                            }

                            let response = "ERROR: Authentication failed".to_string();
                            writer.write_all(response.as_bytes())?;
                            writer.write_all(b"\n")?;
                            writer.flush()?;
                            continue;
                        } else {
                            let response =
                                "ERROR: Not authenticated. Use AUTH <username> <password>"
                                    .to_string();
                            writer.write_all(response.as_bytes())?;
                            writer.write_all(b"\n")?;
                            writer.flush()?;
                            continue;
                        }
                    }

                    // Parse and execute the command
                    let response = match Self::process_command(command, &mut client) {
                        Ok(response) => response,
                        Err(e) => format!("ERROR: {}", e),
                    };

                    // Send response
                    writer.write_all(response.as_bytes())?;
                    writer.write_all(b"\n")?;
                    writer.flush()?;

                    // If the command was EXIT, break the loop
                    if command == "EXIT" {
                        break;
                    }
                }
                Err(e) => {
                    eprintln!("Error reading from client: {}", e);
                    break;
                }
            }
        }
        Ok(())
    }

    fn process_command(command: &str, client: &mut Client) -> Result<String> {
        let parts: Vec<&str> = command.split_whitespace().collect();
        if parts.is_empty() {
            return Ok("ERROR: Empty command".to_string());
        }
        match parts[0] {
            "INSERT" => {
                if parts.len() < 3 {
                    return Ok(
                        "ERROR: Usage: INSERT <database> <collection> <document>".to_string()
                    );
                }
                let db_name = parts[1];
                let coll_name = parts[2];
                let doc_json = parts[3..].join(" ");
                let doc: Document = serde_json::from_str(&doc_json)?;
                let db = client.db(db_name);
                let coll = db.collection(coll_name);
                let id = coll.insert(doc)?;
                Ok(format!("OK: Document inserted with ID: {}", id))
            }
            "FIND" => {
                if parts.len() < 3 {
                    return Ok("ERROR: Usage: FIND <database> <collection> [query]".to_string());
                }
                let db_name = parts[1];
                let coll_name = parts[2];
                let query = if parts.len() > 3 {
                    let query_json = parts[3..].join(" ");
                    let query_value: Value = serde_json::from_str(&query_json)?;
                    Query::from_value(query_value)?
                } else {
                    Query::new()
                };
                let db = client.db(db_name);
                let coll = db.collection(coll_name);
                let results = coll.find(query, None)?;
                let results_json = serde_json::to_string(&results)?;
                Ok(format!("OK: {}", results_json))
            }
            "UPDATE" => {
                if parts.len() < 4 {
                    return Ok(
                        "ERROR: Usage: UPDATE <database> <collection> <query> <update>".to_string(),
                    );
                }
                let db_name = parts[1];
                let coll_name = parts[2];
                // Find the end of the query
                let mut query_end = 3;
                let mut brace_count = 0;
                for (i, c) in command.char_indices() {
                    if i < parts[0].len() + parts[1].len() + parts[2].len() + 2 {
                        continue;
                    }
                    if c == '{' {
                        brace_count += 1;
                    } else if c == '}' {
                        brace_count -= 1;
                        if brace_count == 0 {
                            query_end = i + 1;
                            break;
                        }
                    }
                }
                let query_json =
                    command[parts[0].len() + parts[1].len() + parts[2].len() + 3..query_end].trim();
                let update_json = command[query_end..].trim();
                let query: Query = serde_json::from_str(query_json)?;
                let update: UpdateDocument = serde_json::from_str(update_json)?;
                let db = client.db(db_name);
                let coll = db.collection(coll_name);
                let count = coll.update_one(query, update, false)?;
                Ok(format!("OK: Updated {} documents", count))
            }
            "DELETE" => {
                if parts.len() < 3 {
                    return Ok("ERROR: Usage: DELETE <database> <collection> [query]".to_string());
                }
                let db_name = parts[1];
                let coll_name = parts[2];
                let query = if parts.len() > 3 {
                    let query_json = parts[3..].join(" ");
                    let query_value: Value = serde_json::from_str(&query_json)?;
                    Query::from_value(query_value)?
                } else {
                    Query::new()
                };
                let db = client.db(db_name);
                let coll = db.collection(coll_name);
                let count = coll.delete_one(query)?;
                Ok(format!("OK: Deleted {} documents", count))
            }
            "SAVE" => {
                client.save()?;
                Ok("OK: Data saved to disk".to_string())
            }
            "LOAD" => {
                client.load()?;
                Ok("OK: Data loaded from disk".to_string())
            }
            "LIST_DB" => {
                let dbs = client.list_database_names();
                Ok(format!("OK: {}", serde_json::to_string(&dbs)?))
            }
            "LIST_COLL" => {
                if parts.len() < 2 {
                    return Ok("ERROR: Usage: LIST_COLL <database>".to_string());
                }
                let db_name = parts[1];
                let db = client.db(db_name);
                let colls = db.list_collection_names();
                Ok(format!("OK: {}", serde_json::to_string(&colls)?))
            }
            "STATS" => {
                if parts.len() < 2 {
                    return Ok("ERROR: Usage: STATS <database>".to_string());
                }
                let db_name = parts[1];
                let db = client.db(db_name);
                let stats = db.stats()?;
                Ok(format!("OK: {}", serde_json::to_string(&stats)?))
            }
            "EXIT" => {
                // Save data before exiting
                if client.config().storage_path.is_some() {
                    if let Err(e) = client.save() {
                        return Ok(format!("ERROR: Failed to save data: {}", e));
                    }
                }
                Ok("OK: Connection closing".to_string())
            }
            _ => Ok("ERROR: Unknown command".to_string()),
        }
    }
}