notmongo 0.1.5

use crate::{helpers, parsing, FindQuery, Projection, Sort, UpdateQuery};
use bson;
use bson::Bson;
use rand::Rng;
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fs;
use std::path::Path;

pub struct DeleteResult {
    pub n_deleted: usize,
}

pub struct UpdateResult {
    pub matched_count: usize,
    pub modified_count: usize,
    pub upserted_id: Option<Bson>,
}

pub struct InsertOneResult {
    pub inserted_id: Bson,
}

pub struct InsertManyResult {
    pub inserted_ids: Vec<Bson>,
}

pub struct Collection {
    pub documents: Vec<bson::Document>,
}

#[derive(Clone, Copy, PartialEq)]
pub enum Compression {
    None,
    Zip,
}

impl Collection {
    pub fn new_empty() -> Collection {
        Collection {
            documents: Vec::new(),
        }
    }

    fn find_by_id(&self, id: &Bson) -> Option<&bson::Document> {
        self.documents
            .iter()
            .find(|doc| doc.get("_id").unwrap() == id)
    }

    fn new_from_bson(db: &Bson) -> Result<Collection, String> {
        let db = parsing::ensure_is_object(db, "collection root")?;
        let documents = parsing::get_key(db, "documents", "collection documents")?;
        let documents = parsing::ensure_is_array(documents, "collection documents")?;

        let mut result = Collection::new_empty();
        // let mut visited_ids = HashSet::new();

        for document in documents.iter() {
            // Is an object?
            let document = parsing::ensure_is_object(document, "document")?;

            // Make sure the document has an `_id`
            match document.get("_id") {
                Some(id) => id,
                None => {
                    return Err(format!(
                        "Document is missing the `_id` field: {:?}",
                        document
                    ))
                }
            };

            // TODO: Ensure the id is unique. This is nontrivial because BSON
            // values aren't comparable/hashable. Looking up the object via
            // `find_by_id` is possible, but incurs quadratic cost. This is how
            // this function used to be implemented, but real world databases
            // were slowed down FAR too much.
            // if !visited_ids.insert(doc_id) {
            //     return Err(format!(
            //         "Multiple documents share the same `_id` field {:?}",
            //         document
            //     ));
            // }

            // TOOD: Make sure the values are valid ids? Which types does
            // mongodb allow as ids?

            result.documents.push(document.clone());
        }

        Ok(result)
    }

    pub fn new_from_json_string(json: &str) -> Result<Collection, String> {
        let as_json: serde_json::Value = match serde_json::from_str(json) {
            Ok(value) => value,
            Err(_) => {
                return Err("The string is not valid JSON".to_string());
            }
        };

        let as_bson = match bson::Bson::try_from(as_json) {
            Ok(value) => value,
            Err(_) => {
                return Err("The string is valid JSON, but not valid BSON".to_owned());
            }
        };

        Collection::new_from_bson(&as_bson)
    }

    pub fn new_from_json_file<P: AsRef<Path>>(path: P) -> Result<Collection, String> {
        let json_string = match fs::read_to_string(path) {
            std::io::Result::Ok(value) => value,
            std::io::Result::Err(value) => {
                return Err(format!("Failed to read file: {}", value));
            }
        };

        Collection::new_from_json_string(&json_string)
    }

    pub fn dump_to_bson_value(&self) -> bson::Document {
        // Gather all documents
        let mut documents_field = Vec::<Bson>::new();

        for entry in self.documents.iter() {
            documents_field.push(Bson::Document(entry.clone()));
        }

        // Construct the result and return
        bson::doc! {
            "documents": documents_field
        }
    }

    pub fn len(&self) -> usize {
        self.documents.len()
    }

    pub fn is_empty(&self) -> bool {
        self.documents.is_empty()
    }

    /// Recently accessed documents are often soon accessed again. This
    /// function moves them forward in the array so they are found quickly
    /// during linear scans. The exact position is undefined.
    fn optimize_move_document_forward(&mut self, current_index: usize) {
        // Just moving the document to position 0 would make it slow to find the
        // document previously there. This isn't good, because that document was
        // likely put there just recently by this exact optimizaion. So instead,
        // move the document to a random location close to the start, but not
        // actually at the start.
        let mut rng = rand::thread_rng();
        let target_pos = std::cmp::min(rng.gen_range(0..100), self.documents.len() - 1);
        self.documents.swap(current_index, target_pos)
    }

    pub fn count_documents(
        &self,
        filter: &FindQuery,
        skip: usize,
        limit: usize,
    ) -> Result<usize, String> {
        // Never visit more documents than this many
        let max_visit_counter = match limit {
            0 => self.documents.len(),
            value => skip + value,
        };

        // Count the matching documents, respecting the limit
        let mut counter = 0;

        for (ii, document) in self.documents.iter().enumerate() {
            if ii > max_visit_counter {
                break;
            }

            if filter.matches(document) {
                counter += 1;
            }
        }

        // Apply the skip value
        if counter > skip {
            Ok(counter - skip)
        } else {
            Ok(0)
        }
    }

    pub fn delete_many(&mut self, filter: &FindQuery) -> Result<DeleteResult, String> {
        // Remove mismatching documents

        // FIXME there has got to be a better way of doing this.
        // Filter + Collect doesn't work, because the resulting values are
        // references, not proper documents.

        // This method needs to repeatedly move all remaining documents in the
        // array forward.
        let len_before = self.documents.len();
        let mut ii = 0;
        while ii < self.documents.len() {
            if !filter.matches(&self.documents[ii]) {
                ii += 1;
                continue;
            }

            self.documents.remove(ii);
        }

        // Construct the result and return
        return Ok(DeleteResult {
            n_deleted: len_before - self.documents.len(),
        });
    }

    pub fn delete_one(&mut self, filter: &FindQuery) -> Result<DeleteResult, String> {
        // Try to find a matching document
        for (ii, document) in self.documents.iter().enumerate() {
            // Does it match the query?
            if !filter.matches(document) {
                continue;
            }

            // Delete
            self.documents.remove(ii);

            return Ok(DeleteResult { n_deleted: 1 });
        }

        // No matching document found, return
        Ok(DeleteResult { n_deleted: 0 })
    }

    pub fn find(
        &mut self,
        filter: &FindQuery,
        projection: &Projection,
        skip: usize,
        limit: usize,
        sort: &Sort,
    ) -> impl IntoIterator<Item = bson::Document> {
        // Don't copy documents that will be skipped and don't visit more
        // documents than required. Take care not to miss any documents if
        // sorting is used.
        let mut pre_skip;
        let post_skip;
        let max_keep;

        if sort.is_no_op() {
            pre_skip = skip;
            post_skip = 0;
            max_keep = limit;
        } else {
            pre_skip = 0;
            post_skip = skip;
            max_keep = self.documents.len();
        }

        // Find all matching documents
        let mut results = Vec::new();
        for doc_ii in 0..self.documents.len() {
            let doc = &self.documents[doc_ii];

            // Does the document match the query?
            if !filter.matches(doc) {
                continue;
            }

            // Will this document be skipped anyway?
            if pre_skip > 0 {
                pre_skip -= 1;
                continue;
            }

            // Keep the document
            results.push(doc.clone());

            // Move the document forward to accelerate future queries involving
            // it
            self.documents.swap(doc_ii, results.len() - 1);

            // Found the maximum required number of documents
            if results.len() == max_keep {
                break;
            }
        }

        // Sort if required
        if !sort.is_no_op() {
            results.sort_unstable_by(|a, b| sort.compare(&a, &b));
        }

        // Apply the skip value
        results.drain(0..post_skip);

        // Apply the limit
        if limit > 0 {
            results.truncate(limit);
        }

        // Apply the projection
        if !projection.is_keep_all() {
            for document in results.iter_mut() {
                projection.apply_in_place(document);
            }
        }

        // Done
        results
    }

    pub fn find_one(
        &mut self,
        filter: &FindQuery,
        projection: &Projection,
        skip: usize,
        sort: &Sort,
    ) -> Option<bson::Document> {
        let mut iter = self.find(filter, projection, skip, 1, sort).into_iter();

        match iter.next() {
            Some(value) => Some(value.clone()),
            None => None,
        }
    }

    pub fn insert_one(&mut self, mut document: bson::Document) -> Result<InsertOneResult, String> {
        // Add an _id if there is none
        let entry_id = match document.get("_id") {
            None => {
                let id = Bson::ObjectId(bson::oid::ObjectId::new());
                document.insert("_id", id.clone());
                id
            }
            Some(id) => id.clone(),
        };

        // Make sure there are no duplicate IDs
        match self.find_by_id(&entry_id) {
            None => {}
            Some(_) => {
                return Err(format!("Duplicate ID found: {:?}", entry_id));
            }
        }

        // Insert & return
        self.documents.push(document);
        self.optimize_move_document_forward(self.documents.len() - 1);

        // Construct the reply and return
        return Ok(InsertOneResult {
            inserted_id: entry_id,
        });
    }

    pub fn insert_many<I: IntoIterator<Item = bson::Document>>(
        &mut self,
        documents: I,
    ) -> Result<InsertManyResult, String> {
        // TODO: The current implementation is inefficient. It iterates over the
        //       entire collection for every inserted element just to make sure
        //       the ids are unique. This could be done in just one iteration
        //       for all documents.

        // Insert all values, remembering the IDs
        let mut ids = Vec::new();

        for document in documents.into_iter() {
            let insert_result = self.insert_one(document)?;
            ids.push(insert_result.inserted_id);
        }

        // Construct the result and return
        return Ok(InsertManyResult { inserted_ids: ids });
    }

    pub fn replace_one(
        &mut self,
        filter: &FindQuery,
        replacement: bson::Document,
        upsert: bool,
    ) -> Result<UpdateResult, String> {
        // Try to find a matching document
        for (ii, document) in self.documents.iter().enumerate() {
            // Does it match the query?
            if !filter.matches(document) {
                continue;
            }

            // Replace it
            // TODO: Verify the document? What if the resulting document causes
            // duplicate ids for example?
            self.documents[ii] = replacement;
            self.optimize_move_document_forward(ii);

            return Ok(UpdateResult {
                matched_count: 1,
                modified_count: 1,
                upserted_id: None,
            });
        }

        // No matching document found, upsert?
        if upsert {
            let insert_result = self.insert_one(replacement)?;
            return Ok(UpdateResult {
                matched_count: 0,
                modified_count: 1,
                upserted_id: Some(insert_result.inserted_id),
            });
        }

        // No matching document found, return
        return Ok(UpdateResult {
            matched_count: 0,
            modified_count: 0,
            upserted_id: None,
        });
    }

    pub fn update_one(
        &mut self,
        filter: &FindQuery,
        update: &UpdateQuery,
        upsert: bool,
    ) -> Result<UpdateResult, String> {
        for (ii, document) in self.documents.iter_mut().enumerate() {
            // Does it match the query?
            if !filter.matches(document) {
                continue;
            }

            // Apply the update
            // TODO: Verify the resulting document?
            // - What if it causes duplicate ids for example
            // - What if the document doesn't have an id at all?
            update.apply_inplace(document);
            self.optimize_move_document_forward(ii);

            return Ok(UpdateResult {
                matched_count: 1,
                modified_count: 1,
                upserted_id: None,
            });
        }

        // No matching document found, upsert?
        if upsert {
            let insert_result = self.insert_one(update.apply_to_empty())?;
            return Ok(UpdateResult {
                matched_count: 0,
                modified_count: 1,
                upserted_id: Some(insert_result.inserted_id),
            });
        }

        // No matching document found, return
        Ok(UpdateResult {
            matched_count: 0,
            modified_count: 0,
            upserted_id: None,
        })
    }

    pub fn update_many(
        &mut self,
        filter: &FindQuery,
        update: &UpdateQuery,
        upsert: bool,
    ) -> Result<UpdateResult, String> {
        let mut result = UpdateResult {
            matched_count: 0,
            modified_count: 0,
            upserted_id: None,
        };

        for document in self.documents.iter_mut() {
            // Does it match the query?
            if !filter.matches(document) {
                continue;
            }
            result.matched_count += 1;

            // Apply the update
            // TODO: Verify the resulting document?
            // - What if it causes duplicate ids for example
            // - What if the document doesn't have an id at all?
            update.apply_inplace(document);
            result.modified_count += 1;
        }

        // No matching document found, upsert?
        if result.matched_count == 0 && upsert {
            let insert_result = self.insert_one(update.apply_to_empty())?;
            return Ok(UpdateResult {
                matched_count: 0,
                modified_count: 1,
                upserted_id: Some(insert_result.inserted_id),
            });
        }

        Ok(result)
    }
}

pub struct Database {
    all_collections: HashMap<String, Collection>,
}

impl Database {
    pub fn new_empty() -> Database {
        Database {
            all_collections: HashMap::new(),
        }
    }

    pub fn new_from_bson_value(db: &Bson) -> Result<Database, String> {
        let db = parsing::ensure_is_object(db, "database root")?;
        let collections = parsing::get_key(db, "collections", "database collections")?;
        let collections = parsing::ensure_is_object(collections, "database collections")?;

        let mut result = Database::new_empty();

        for (name, collection) in collections.iter() {
            result
                .all_collections
                .insert(name.to_string(), Collection::new_from_bson(collection)?);
        }

        Ok(result)
    }

    pub fn new_from_json_string(json: &str) -> Result<Database, String> {
        let as_json: serde_json::Value = match serde_json::from_str(json) {
            Ok(value) => value,
            Err(_) => {
                return Err("The string is not valid JSON".to_string());
            }
        };

        let as_bson = match bson::Bson::try_from(as_json) {
            Ok(value) => value,
            Err(_) => {
                return Err("The string is valid JSON, but not valid BSON".to_owned());
            }
        };

        let res = Database::new_from_bson_value(&as_bson);
        res
    }

    pub fn new_from_bson_file<P: AsRef<Path>>(
        path: P,
        compression: Compression,
    ) -> Result<Database, String> {
        // Read the file into a buffer
        let buffer = helpers::read_compressed_file(path, compression)?;

        // Parse the buffer as BSON
        let as_document = bson::Document::from_reader(buffer.as_slice())
            .map_err(|_| "The BSON file is not valid BSON".to_string())?;

        let as_bson = bson::Bson::try_from(as_document).unwrap();

        // Load the database from the BSON
        Database::new_from_bson_value(&as_bson)
    }

    pub fn new_from_json_file<P: AsRef<Path>>(
        path: P,
        compression: Compression,
    ) -> Result<Database, String> {
        // Read the file into a string
        let utf8_string = helpers::read_compressed_file(path, compression)?;
        let json_string = String::from_utf8(utf8_string)
            .map_err(|_| "The JSON file is not valid UTF-8..".to_string())?;

        // Load the database from the string
        Database::new_from_json_string(&json_string)
    }

    pub fn dump_to_bson_value(&self) -> bson::Document {
        // Convert all collections
        let mut collections_field = bson::Document::new();

        for entry in self.all_collections.iter() {
            // No point in serializing empty collections. These shouldn't even
            // exist as far as MongoDB is concerned.
            if entry.1.is_empty() {
                continue;
            }

            collections_field.insert(entry.0.clone(), entry.1.dump_to_bson_value());
        }

        // Construct the result and return
        bson::doc! {
            "collections": collections_field
        }
    }

    pub fn dump_to_bson_file<P: AsRef<Path>>(
        &self,
        path: &P,
        n_backups: u32,
        compression: Compression,
    ) -> Result<(), String> {
        helpers::dump_to_bson_file(
            &self.dump_to_bson_value().into(),
            path,
            n_backups,
            compression,
        )
    }

    pub fn dump_to_json_file<P: AsRef<Path>>(
        &self,
        path: &P,
        relaxed: bool,
        format: bool,
        n_backups: u32,
        compression: Compression,
    ) -> Result<(), String> {
        helpers::dump_to_json_file(
            self.dump_to_bson_value().into(),
            relaxed,
            format,
            path,
            n_backups,
            compression,
        )
    }

    pub fn dump_to_json_string(&self, relaxed: bool, format: bool) -> Result<String, String> {
        helpers::dump_to_json_string(self.dump_to_bson_value().into(), relaxed, format)
    }

    pub fn collection_mut(&mut self, name: &str) -> &mut Collection {
        // Reuse any existing collection, or create a new one of none is
        // present.
        return self
            .all_collections
            .entry(name.to_string())
            .or_insert(Collection::new_empty());
    }

    pub fn drop_collection(&mut self, name: &str) -> bool {
        // FIXME: Note that this will return `true` for empty collections. This
        //        is a bit weird, because saving, then reloading will wipe those
        //        collections. Empty collections aren't really supposed to exist
        //        in MongoDB...
        //
        // -> Check what the real mongo does in this case
        return self.all_collections.remove(name).is_some();
    }

    pub fn n_collections(&self) -> usize {
        self.all_collections.len()
    }

    pub fn collection_names(&self) -> impl IntoIterator<Item = String> {
        let as_vector: Vec<String> = self.all_collections.keys().cloned().collect();
        as_vector
    }
}