notmongo 0.1.7

use bson::Bson;
use std::collections::HashSet;

pub struct Projection {
    keep_by_default: bool,
    exceptions: HashSet<String>,
}

impl Projection {
    pub fn new_keep_all() -> Projection {
        Projection {
            keep_by_default: true,
            exceptions: HashSet::new(),
        }
    }

    pub fn new_drop_id() -> Projection {
        let mut exceptions = HashSet::new();
        exceptions.insert("_id".to_string());

        Projection {
            keep_by_default: true,
            exceptions: exceptions,
        }
    }

    pub fn parse_bson<B: Into<Bson>>(projection: B) -> Result<Projection, String> {
        // From the MongoDB documentation:
        //
        // "A list of field names that should be returned in the result
        // set or a dict specifying the fields to include or exclude. If projection
        // is a list “_id” will always be returned. Use a dict to exclude fields
        // from the result (e.g. projection={‘_id’: False})."

        let projection: &Bson = &projection.into();

        // Case: The projection is an array
        if let Bson::Array(projection) = projection {
            let mut result = Projection {
                keep_by_default: false,
                exceptions: HashSet::new(),
            };

            result.exceptions.insert("_id".to_string());

            for field_name in projection {
                match field_name {
                    Bson::String(string) => result.exceptions.insert(string.clone()),
                    _ => {
                        return Err("The projection array contains a non-string element".to_string())
                    }
                };
            }

            return Ok(result);
        }

        // Case: The projection is an object
        if let Bson::Document(projection) = projection {
            // See which fields are inclusive and which ones exclusive
            let mut inclusive = HashSet::new();
            let mut exclusive = HashSet::new();

            for entry in projection {
                let field_name = entry.0;
                let include = match entry.1 {
                    Bson::Boolean(value) => *value,
                    _ => {
                        return Err(
                            "The values in projection objects need to be either `True` or `False`"
                                .to_string(),
                        );
                    }
                };

                if include {
                    inclusive.insert(field_name.to_string());
                } else {
                    exclusive.insert(field_name.to_string());
                }
            }

            // _id is a special case, keep track
            let explicit_inclusive_id = inclusive.contains("_id");
            let explicit_exclusive_id = exclusive.contains("_id");
            assert!(!(explicit_inclusive_id && explicit_exclusive_id));

            let id_is_explicit = explicit_inclusive_id || explicit_exclusive_id;

            let have_inclusive = inclusive.len() - explicit_inclusive_id as usize > 0;
            let have_exclusive = exclusive.len() - explicit_exclusive_id as usize > 0;

            if have_inclusive && have_exclusive {
                return Err(
                    "Projections cannot have inclusive and exclusive items, except for _id"
                        .to_string(),
                );
            }

            // Order here matters! If no elements are specified, all fields are
            // retained
            return if have_inclusive {
                if !id_is_explicit {
                    inclusive.insert("_id".to_string());
                }

                Ok(Projection {
                    keep_by_default: false,
                    exceptions: inclusive,
                })
            } else {
                Ok(Projection {
                    keep_by_default: true,
                    exceptions: exclusive,
                })
            };
        }

        // Case: The projection is invalid
        Err("Projections need to be either Arrays of field names, or Objects.".to_string())
    }

    pub(crate) fn is_keep_all(&self) -> bool {
        self.keep_by_default && self.exceptions.is_empty()
    }

    fn keep_field(&self, field_name: &str) -> bool {
        let is_exception = self.exceptions.contains(field_name);
        self.keep_by_default ^ is_exception
    }

    #[allow(dead_code)]
    pub(crate) fn apply_to_clone(&self, document: &bson::Document) -> bson::Document {
        let mut result = document.clone();

        for (name, value) in document.iter() {
            if self.keep_field(name) {
                result.insert(name, value);
            }
        }

        result
    }

    pub(crate) fn apply_in_place(&self, document: &mut bson::Document) {
        let keys: Vec<String> = document.keys().cloned().collect();

        for name in keys {
            if !self.keep_field(&name) {
                document.remove(name);
            }
        }
    }
}