use bson::{Bson, Document};

/// Get a nested value from a document using dot-notation path (e.g., "a.b.c").
/// Handles both nested documents and array traversal.
pub fn get_nested_value(doc: &Document, path: &str) -> Option<Bson> {
    let parts: Vec<&str> = path.split('.').collect();
    get_nested_recursive(&Bson::Document(doc.clone()), &parts)
}

fn get_nested_recursive(value: &Bson, parts: &[&str]) -> Option<Bson> {
    if parts.is_empty() {
        return Some(value.clone());
    }

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

    match value {
        Bson::Document(doc) => {
            let child = doc.get(key)?;
            get_nested_recursive(child, rest)
        }
        Bson::Array(arr) => {
            // Try numeric index first
            if let Ok(idx) = key.parse::<usize>() {
                if let Some(elem) = arr.get(idx) {
                    return get_nested_recursive(elem, rest);
                }
            }
            // Otherwise, collect from all elements
            let results: Vec<Bson> = arr
                .iter()
                .filter_map(|elem| get_nested_recursive(elem, parts))
                .collect();
            if results.is_empty() {
                None
            } else if results.len() == 1 {
                Some(results.into_iter().next().unwrap())
            } else {
                Some(Bson::Array(results))
            }
        }
        _ => None,
    }
}

/// Set a nested value in a document using dot-notation path.
pub fn set_nested_value(doc: &mut Document, path: &str, value: Bson) {
    let parts: Vec<&str> = path.split('.').collect();
    set_nested_recursive(doc, &parts, value);
}

fn set_nested_recursive(doc: &mut Document, parts: &[&str], value: Bson) {
    if parts.len() == 1 {
        doc.insert(parts[0].to_string(), value);
        return;
    }

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

    // Get or create nested document
    if !doc.contains_key(key) {
        doc.insert(key.to_string(), Bson::Document(Document::new()));
    }

    if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
        set_nested_recursive(nested, rest, value);
    }
}

/// Remove a nested value from a document using dot-notation path.
pub fn remove_nested_value(doc: &mut Document, path: &str) -> Option<Bson> {
    let parts: Vec<&str> = path.split('.').collect();
    remove_nested_recursive(doc, &parts)
}

fn remove_nested_recursive(doc: &mut Document, parts: &[&str]) -> Option<Bson> {
    if parts.len() == 1 {
        return doc.remove(parts[0]);
    }

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

    if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
        remove_nested_recursive(nested, rest)
    } else {
        None
    }
}

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

    #[test]
    fn test_get_nested_simple() {
        let doc = bson::doc! { "a": { "b": { "c": 42 } } };
        assert_eq!(get_nested_value(&doc, "a.b.c"), Some(Bson::Int32(42)));
    }

    #[test]
    fn test_get_nested_missing() {
        let doc = bson::doc! { "a": { "b": 1 } };
        assert_eq!(get_nested_value(&doc, "a.c"), None);
    }

    #[test]
    fn test_set_nested() {
        let mut doc = bson::doc! {};
        set_nested_value(&mut doc, "a.b.c", Bson::Int32(42));
        assert_eq!(get_nested_value(&doc, "a.b.c"), Some(Bson::Int32(42)));
    }
}