radix_immutable 0.1.0

//! The main trie implementation.
//!
//! This module contains the `Trie` type, which provides the primary API for working
//! with the radix trie data structure.

use std::collections::{BTreeMap, VecDeque};
use std::hash::Hash;
use std::marker::PhantomData;
use std::sync::Arc;

use crate::key_converter::{BytesKeyConverter, KeyToBytes, StrKeyConverter};
use crate::node::TrieNode;
use crate::prefix_view::{PrefixView, PrefixViewIter};
use crate::util::prefix_match;

/// An immutable radix trie with structural sharing.
///
/// This Radix Trie (also known as a Patricia Trie) is an ordered tree data structure
/// that efficiently stores and retrieves key-value pairs while preserving
/// the key ordering.
///
/// This implementation is immutable - all operations that would modify the trie
/// return a new trie instance that shares unchanged parts of the structure with
/// the original via `Arc`.
///
/// `K` is the key type, `V` is the value type, and `KC` is the KeyConverter strategy.
///
/// **Note**: For common key types, you can use type aliases like `StringTrie<K, V>` (for string keys) or `BytesTrie<K, V>` (for byte keys).
#[derive(Debug)]
pub struct Trie<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> {
    /// The root node of the trie
    pub(crate) root: Arc<TrieNode<K, V>>,

    /// The number of values stored in the trie
    size: usize,

    /// Phantom data for the key converter type
    _phantom_kc: PhantomData<KC>,
}

impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> Clone for Trie<K, V, KC> {
    fn clone(&self) -> Self {
        Trie {
            root: Arc::clone(&self.root),
            size: self.size,
            _phantom_kc: PhantomData,
        }
    }
}

impl<K: Clone + Hash + Eq, V> Trie<K, V, StrKeyConverter<K>>
where
    K: AsRef<str>, // StrKeyConverter requires K: AsRef<str>
{
    /// Creates a new, empty trie configured for string-like keys (e.g., `String`, `&str`).
    ///
    /// Keys must implement `AsRef<str>`, `Clone`, `Hash`, and `Eq`.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::Trie;
    ///
    /// let trie = Trie::<String, i32, _>::new_str_key();
    /// assert!(trie.is_empty());
    /// ```
    pub fn new_str_key() -> Self {
        Trie {
            root: Arc::new(TrieNode::new(Vec::new())),
            size: 0,
            _phantom_kc: PhantomData,
        }
    }
}

impl<K: Clone + Hash + Eq, V> Trie<K, V, BytesKeyConverter<K>>
where
    K: AsRef<[u8]>, // BytesKeyConverter requires K: AsRef<[u8]>
{
    /// Creates a new, empty trie configured for byte-slice-like keys (e.g., `Vec<u8>`, `&[u8]`).
    ///
    /// Keys must implement `AsRef<[u8]>`, `Clone`, `Hash`, and `Eq`.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::Trie;
    /// use radix_immutable::BytesKeyConverter; // Usually StrKeyConverter is the default or inferred
    ///
    /// let trie = Trie::<Vec<u8>, i32, BytesKeyConverter<Vec<u8>>>::new_bytes_key();
    /// assert!(trie.is_empty());
    /// ```
    pub fn new_bytes_key() -> Self {
        Trie {
            root: Arc::new(TrieNode::new(Vec::new())),
            size: 0,
            _phantom_kc: PhantomData,
        }
    }
}

impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> Trie<K, V, KC> {
    /// Creates a new, empty trie with the specified key converter.
    ///
    /// `K` must implement `Clone`, `Hash`, and `Eq`.
    /// `KC` must implement `KeyToBytes<K>`.
    pub fn new_with_converter(_converter: KC) -> Self {
        Trie {
            root: Arc::new(TrieNode::new(Vec::new())),
            size: 0,
            _phantom_kc: PhantomData,
        }
    }
}

// General methods applicable to any Trie configuration
impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> Trie<K, V, KC> {
    /// Returns the number of values stored in the trie.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new();
    /// assert_eq!(trie.len(), 0);
    ///
    /// let trie2 = trie.insert("hello".to_string(), 42);
    /// assert_eq!(trie2.len(), 1);
    /// ```
    pub fn len(&self) -> usize {
        self.size
    }

    /// Returns `true` if the trie contains no values.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new();
    /// assert!(trie.is_empty());
    ///
    /// let trie2 = trie.insert("hello".to_string(), 42);
    /// assert!(!trie2.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    /// Creates a view of the subtrie at the given key prefix.
    ///
    /// This method returns a lightweight view into a subtrie defined by a key prefix.
    /// The view supports efficient comparison and lookup operations on the subtrie.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::{StringTrie, StringPrefixView};
    ///
    /// let trie = StringTrie::<String, i32>::new()
    ///     .insert("hello".to_string(), 1)
    ///     .insert("help".to_string(), 2);
    ///
    /// // Create a view of the subtrie at "hel" prefix
    /// let view = trie.view_subtrie("hel".to_string());
    ///
    /// // Check if keys exist in the view
    /// assert!(view.contains_key(&"hello".to_string()));
    ///
    /// // Get values from the view
    /// assert_eq!(view.get(&"hello".to_string()), Some(&1));
    /// ```
    /// Creates a prefix view over this trie with the given prefix.
    ///
    /// This allows for efficient querying and iteration of keys that start with the prefix.
    pub fn view_subtrie(&self, prefix: K) -> PrefixView<K, V, KC>
    where
        V: Clone,  // PrefixView might need V: Clone, K: Clone is from K: KeyToBytes
        KC: Clone, // KC from Trie implies KC: KeyToBytes which implies KC: Clone
    {
        PrefixView::new(self.clone(), prefix) // self.clone requires K,V,KC all Clone
    }

    /// Returns an iterator over all key-value pairs whose keys start with the
    /// given prefix.
    ///
    /// This is a convenience for `trie.view_subtrie(prefix).iter()`.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new()
    ///     .insert("hello".to_string(), 1)
    ///     .insert("help".to_string(), 2)
    ///     .insert("world".to_string(), 3);
    ///
    /// let hel_entries: Vec<_> = trie.iter_prefix("hel".to_string()).collect();
    /// assert_eq!(hel_entries.len(), 2);
    /// ```
    pub fn iter_prefix(&self, prefix: K) -> PrefixViewIter<K, V, KC>
    where
        V: Clone,
        KC: Clone,
    {
        self.view_subtrie(prefix).into_iter()
    }

    /// Returns an iterator over all the key-value pairs in the trie.
    ///
    /// The iterator yields pairs of `(K, V)` (cloned values) in depth-first order.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::{Trie, StringTrie};
    /// use std::collections::HashSet;
    ///
    /// // Use explicit type annotation with StringTrie
    /// let mut trie = StringTrie::<String, i32>::new();
    /// trie = trie.insert("hello".to_string(), 1);
    /// trie = trie.insert("world".to_string(), 2);
    ///
    /// let entries: HashSet<_> = trie.iter().collect();
    /// assert_eq!(entries.len(), 2);
    /// ```
    pub fn iter(&self) -> TrieIter<K, V, KC> where V: Clone {
        let mut stack = VecDeque::new();

        // Start with the root node
        stack.push_back(Arc::clone(&self.root));

        TrieIter {
            stack,
            _phantom: PhantomData,
        }
    }
    
    /// Returns an iterator over Arc references to the key-value pairs in the trie.
    ///
    /// The iterator yields pairs of `(Arc<K>, Arc<V>)` in depth-first order.
    /// This avoids cloning the actual values but changes the return type.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::{Trie, StringTrie};
    /// use std::sync::Arc;
    ///
    /// // Use explicit type annotation with StringTrie
    /// let mut trie = StringTrie::<String, i32>::new();
    /// trie = trie.insert("hello".to_string(), 1);
    /// trie = trie.insert("world".to_string(), 2);
    ///
    /// for (key, value) in trie.iter_arc() {
    ///     println!("{}: {}", key, value);
    /// }
    /// ```
    pub fn iter_arc(&self) -> TrieArcIter<K, V, KC> {
        let mut stack = VecDeque::new();

        // Start with the root node
        stack.push_back(Arc::clone(&self.root));

        TrieArcIter {
            stack,
            _phantom: PhantomData,
        }
    }
}

/// An iterator over the entries of a Trie.
///
/// This iterator performs a depth-first traversal of the trie and yields
/// cloned keys and values.
pub struct TrieIter<K: Clone + Hash + Eq, V: Clone, KC: KeyToBytes<K>> {
    /// Stack of nodes to visit
    stack: VecDeque<Arc<TrieNode<K, V>>>,

    /// Phantom data for the key converter type
    _phantom: PhantomData<KC>,
}

impl<K: Clone + Hash + Eq, V: Clone, KC: KeyToBytes<K>> Iterator for TrieIter<K, V, KC> {
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(node) = self.stack.pop_front() {
            // Add children in reverse sorted order for depth-first traversal.
            // BTreeMap iterates in key order; rev() ensures the smallest byte
            // ends up at the front of the deque after push_front.
            for (_, child) in node.children.iter().rev() {
                self.stack.push_front(Arc::clone(child));
            }

            // If this node has a key-value pair, yield it (cloned)
            if let Some(kvp) = &node.data {
                // Clone the key and value
                let key = (*kvp.key).clone();
                let value = (*kvp.value).clone();

                return Some((key, value));
            }
        }

        None
    }
}

/// An iterator over the entries of a Trie that returns Arc references.
///
/// This iterator performs a depth-first traversal of the trie and yields
/// Arc references to the keys and values, avoiding cloning.
pub struct TrieArcIter<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> {
    /// Stack of nodes to visit
    stack: VecDeque<Arc<TrieNode<K, V>>>,

    /// Phantom data for the key converter type
    _phantom: PhantomData<KC>,
}

impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> Iterator for TrieArcIter<K, V, KC> {
    type Item = (Arc<K>, Arc<V>);

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(node) = self.stack.pop_front() {
            for (_, child) in node.children.iter().rev() {
                self.stack.push_front(Arc::clone(child));
            }

            // If this node has a key-value pair, yield Arc references
            if let Some(kvp) = &node.data {
                return Some((Arc::clone(&kvp.key), Arc::clone(&kvp.value)));
            }
        }

        None
    }
}

impl<K: Clone + Hash + Eq, V: Clone, KC: KeyToBytes<K>> IntoIterator for &Trie<K, V, KC> {
    type Item = (K, V);
    type IntoIter = TrieIter<K, V, KC>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<K: Clone + Hash + Eq, V: Clone, KC: KeyToBytes<K> + Default> std::iter::FromIterator<(K, V)>
    for Trie<K, V, KC>
{
    /// Builds a trie from an iterator of key-value pairs.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie: StringTrie<String, i32> = vec![
    ///     ("hello".to_string(), 1),
    ///     ("world".to_string(), 2),
    /// ].into_iter().collect();
    ///
    /// assert_eq!(trie.get(&"hello".to_string()), Some(&1));
    /// assert_eq!(trie.len(), 2);
    /// ```
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let mut trie = Trie::new_with_converter(KC::default());
        for (k, v) in iter {
            trie = trie.insert(k, v);
        }
        trie
    }
}

// Core methods using the KeyConverter KC
impl<K: Clone + Hash + Eq, V: Clone, KC: KeyToBytes<K>> Trie<K, V, KC> {
    /// Retrieves a reference to the value stored for the given key, if any.
    /// The query key `Q` must be convertible to bytes by the Trie's configured `KC` converter,
    /// typically meaning `Q` is compatible with `K` (e.g. `Q = str` when `K = String`).
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, u32>::new()
    ///     .insert("hello".to_string(), 42);
    ///
    /// assert_eq!(trie.get(&"hello".to_string()), Some(&42));
    /// assert_eq!(trie.get(&"world".to_string()), None);
    /// ```
    // Q needs to be convertible by KC. This is simpler if Q is K.
    // For ergonomic querying (e.g. get(&str) when K is String),
    // K must Borrow<Q> and KC::convert must work on K.
    // The KeyToBytes<K>::convert takes &'a K.
    // So, Q must be K, or K must be Borrow<Q> where Q can be turned into &K.
    // Let's assume Q is K for now in the main implementation.
    // Wrappers can provide AsRef<str> or AsRef<[u8]> ergonomics.
    //
    // More general Q version:
    // pub fn get<Q_K, Q_Bytes>(&self, key: &Q_K) -> Option<&V>
    // where
    //     K: Borrow<Q_K>, // K can be borrowed as Q_K
    //     Q_K: Hash + Eq + ?Sized, // Q_K is the type K is borrowed as for hashing
    //     KC: KeyToBytes<K, QueryType = Q_Bytes>, // KC can convert Q_Bytes
    //     Q_K: AsRef<Q_Bytes>, // Q_K can be seen as Q_Bytes
    //     Q_Bytes: ?Sized + AsRef<[u8]>
    // This is too complex.
    //
    // Let's simplify: get takes `&K`.
    // If K=String, you pass &String. Querying with &str needs `get_by_str_ref`.
    // Or, rely on K: Borrow<Q> and assume KC::convert can be called on `key.borrow()`.
    // KeyToBytes<K> operates on &'a K.
    // If K=String, Q=str, K:Borrow<str>. Then key.borrow() is &str.
    // Can't call KC::convert on &str if KC is KeyToBytes<String>.
    //
    // Final decision for this iteration: The primary `get` takes `&K`.
    // Ergonomic wrappers for common query types can be added later if needed.
    pub fn get(&self, key: &K) -> Option<&V> {
        let key_bytes_cow = KC::convert(key);
        let key_bytes_slice = key_bytes_cow.as_ref();

        // Navigate from the root
        let mut current = &self.root;
        let mut remaining = key_bytes_slice;

        while !remaining.is_empty() {
            // Find the common prefix between remaining key and current node's fragment
            let common_len = prefix_match(remaining, &current.key_fragment);

            // If current node's fragment isn't completely matched, key doesn't exist
            if common_len < current.key_fragment.len() {
                return None;
            }

            // Consume the matched part of the key
            remaining = &remaining[common_len..];

            // If there's nothing left, check if this node has a value
            if remaining.is_empty() {
                return current.data.as_ref().map(|kvp| kvp.value.as_ref());
            }

            // Otherwise, look for a child matching the next byte
            let next_byte = remaining[0];
            match current.children.get(&next_byte) {
                Some(child) => {
                    current = child;
                    // Move past the branching byte
                    remaining = &remaining[1..];
                }
                None => return None,
            }
        }

        // If we've consumed the entire key, return the value at this node only
        // if the node's key_fragment was also fully matched (i.e. is empty).
        // A non-empty fragment means the key is shorter than this node's path.
        if !current.key_fragment.is_empty() {
            return None;
        }
        current.data.as_ref().map(|kvp| kvp.value.as_ref())
    }

    /// Returns `true` if the trie contains a value for the given key.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new()
    ///     .insert("hello".to_string(), 42);
    ///
    /// assert_eq!(trie.get(&"hello".to_string()), Some(&42));
    /// assert!(trie.contains_key(&"hello".to_string()));
    /// ```
    // Making contains_key also take &K for consistency with get.
    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }

    /// Returns the key and value of the longest stored key that is a prefix
    /// of the given lookup key (or equals it).
    ///
    /// This is useful for routing tables, URL matching, filesystem path lookups,
    /// and similar problems where you want the most specific matching prefix.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new()
    ///     .insert("/".to_string(), 0)
    ///     .insert("/api".to_string(), 1)
    ///     .insert("/api/users".to_string(), 2);
    ///
    /// // Exact match
    /// assert_eq!(trie.longest_prefix_match(&"/api/users".to_string()), Some((&"/api/users".to_string(), &2)));
    ///
    /// // Longest prefix: "/api" is the longest stored prefix of "/api/posts"
    /// assert_eq!(trie.longest_prefix_match(&"/api/posts".to_string()), Some((&"/api".to_string(), &1)));
    ///
    /// // No stored key is a prefix of "/other"
    /// // (but "/" is a prefix of everything starting with "/")
    /// assert_eq!(trie.longest_prefix_match(&"/other".to_string()), Some((&"/".to_string(), &0)));
    ///
    /// // No match at all
    /// assert_eq!(trie.longest_prefix_match(&"nope".to_string()), None);
    /// ```
    pub fn longest_prefix_match(&self, key: &K) -> Option<(&K, &V)> {
        let key_bytes_cow = KC::convert(key);
        let key_bytes_slice = key_bytes_cow.as_ref();

        let mut current = &self.root;
        let mut remaining = key_bytes_slice;
        let mut best_match: Option<(&K, &V)> = None;

        while !remaining.is_empty() {
            let common_len = prefix_match(remaining, &current.key_fragment);

            if common_len < current.key_fragment.len() {
                break;
            }

            remaining = &remaining[common_len..];

            // Fragment fully matched — if this node has data, it's a prefix match
            if let Some(kvp) = &current.data {
                best_match = Some((kvp.key.as_ref(), kvp.value.as_ref()));
            }

            if remaining.is_empty() {
                break;
            }

            let next_byte = remaining[0];
            match current.children.get(&next_byte) {
                Some(child) => {
                    current = child;
                    remaining = &remaining[1..];
                }
                None => break,
            }
        }

        // If we landed on a node via dispatch and remaining is empty,
        // check if the node's fragment is also empty (meaning we fully
        // reached this node and its data is a valid match).
        if remaining.is_empty() && current.key_fragment.is_empty() {
            if let Some(kvp) = &current.data {
                best_match = Some((kvp.key.as_ref(), kvp.value.as_ref()));
            }
        }

        best_match
    }

    /// Inserts a key-value pair into the trie, returning a new trie.
    ///
    /// If the key already exists, the value is replaced.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new();
    /// let trie2 = trie.insert("hello".to_string(), 42);
    ///
    /// assert!(trie.is_empty());
    /// assert_eq!(trie2.len(), 1);
    /// ```
    pub fn insert(&self, key: K, value: V) -> Self {
        // Convert the key K to bytes using KeyIterBytes trait
        // Get an owned Vec<u8> from Cow to avoid lifetime issues if key is moved.
        let key_bytes_vec: Vec<u8> = KC::convert(&key).into_owned();

        // Call the recursive helper to perform the insertion with path copying
        // Pass a slice of the owned Vec<u8>, and move the original `key`.
        let (new_root, value_replaced) =
            self.insert_recursive(&self.root, &key_bytes_vec, key, value);

        // Calculate the new size
        let new_size = if value_replaced {
            self.size
        } else {
            self.size + 1
        };

        // Create a new trie with the updated root and size
        Trie {
            root: new_root,
            size: new_size,
            _phantom_kc: PhantomData,
        }
    }

    // Recursive helper for insert that handles path copying to allow for structural sharing
    fn insert_recursive(
        &self,
        node: &Arc<TrieNode<K, V>>,
        key_bytes: &[u8],
        original_key: K,
        value: V,
    ) -> (Arc<TrieNode<K, V>>, bool) {
        // Find how much of the key matches the current node's key fragment
        let common_len = prefix_match(key_bytes, &node.key_fragment);

        if common_len < node.key_fragment.len() {
            // The key and node's fragment share a prefix, but they diverge
            // We need to split the current node

            // Create a child with the unmatched part of the node's fragment
            let remaining_fragment = node.key_fragment[common_len + 1..].to_vec();
            let mut child = TrieNode::new(remaining_fragment);
            child.children = node.children.clone();
            child.data = node.data.clone();

            // Create a new node map with the child
            let mut children = BTreeMap::new();
            let branch_byte = node.key_fragment[common_len];
            children.insert(branch_byte, Arc::new(child));

            // If there's more of the key, create a new leaf node
            if common_len < key_bytes.len() {
                let key_fragment = key_bytes[common_len + 1..].to_vec();
                let mut leaf = TrieNode::new(key_fragment);
                leaf.data = Some(crate::node::KeyValuePair {
                    key: Arc::new(original_key.clone()),
                    value: Arc::new(value.clone()),
                });
                let new_branch_byte = key_bytes[common_len];
                children.insert(new_branch_byte, Arc::new(leaf));
            }

            // Create the new split node
            let mut split_node = TrieNode::new(node.key_fragment[..common_len].to_vec());

            // If key is fully consumed at the split point, store the value
            if common_len == key_bytes.len() {
                split_node.data = Some(crate::node::KeyValuePair {
                    key: Arc::new(original_key.clone()),
                    value: Arc::new(value.clone()),
                });
            }

            split_node.children = children;

            // Return the new node (never replacing a value when splitting)
            return (Arc::new(split_node), false);
        }

        // The key fragment was fully matched, now we need to handle the rest of the key
        let remaining = &key_bytes[common_len..];

        if remaining.is_empty() {
            // We've matched the entire key, update the value at this node
            let mut new_node = TrieNode::new(node.key_fragment.clone());
            new_node.children = node.children.clone();
            new_node.data = Some(crate::node::KeyValuePair {
                key: Arc::new(original_key),
                value: Arc::new(value),
            });

            return (Arc::new(new_node), node.data.is_some());
        }

        // We have more key to process, go down the appropriate child
        let next_byte = remaining[0];
        let mut new_children = node.children.clone();

        match node.children.get(&next_byte) {
            Some(child) => {
                // Recursively insert into the child
                let (new_child, value_replaced) =
                    self.insert_recursive(child, &remaining[1..], original_key, value);

                // Update the child in our children map
                new_children.insert(next_byte, new_child);

                // Create a new node with the updated children
                let mut new_node = TrieNode::new(node.key_fragment.clone());
                new_node.data = node.data.clone();
                new_node.children = new_children;

                (Arc::new(new_node), value_replaced)
            }
            None => {
                // No matching child, add a new leaf for the remaining key
                let leaf_fragment = remaining[1..].to_vec();
                let new_leaf = TrieNode::with_key_value(leaf_fragment, original_key, value);

                // Add the new leaf to the children map
                new_children.insert(next_byte, Arc::new(new_leaf));

                // Create a new node with the updated children
                let mut new_node = TrieNode::new(node.key_fragment.clone());
                new_node.data = node.data.clone();
                new_node.children = new_children;

                (Arc::new(new_node), false)
            }
        }
    }

    /// Removes a key-value pair from the trie, returning a new trie.
    ///
    /// If the key exists, the removed value is returned along with the new trie.
    ///
    /// # Examples
    ///
    /// ```
    /// use radix_immutable::StringTrie;
    ///
    /// let trie = StringTrie::<String, i32>::new()
    ///     .insert("hello".to_string(), 42);
    ///
    /// let (trie2, removed_value) = trie.remove(&"hello".to_string());
    ///
    /// assert!(trie2.is_empty());
    /// assert_eq!(removed_value, Some(42));
    /// ```
    // Making remove also take &K for consistency.
    pub fn remove(&self, key: &K) -> (Self, Option<V>) {
        let key_bytes_cow = KC::convert(key);
        let key_bytes_slice = key_bytes_cow.as_ref();
        let mut removed_value = None;

        // Call the recursive helper to perform the removal with path copying
        let new_root = self.remove_recursive(&self.root, key_bytes_slice, &mut removed_value);

        // Calculate the new size
        let new_size = if removed_value.is_some() {
            self.size - 1
        } else {
            self.size
        };

        // Create a new trie with the updated root and size
        (
            Trie {
                root: new_root,
                size: new_size,
                _phantom_kc: PhantomData,
            },
            removed_value,
        )
    }

    // Recursive helper for remove that handles path copying
    fn remove_recursive(
        &self,
        node: &Arc<TrieNode<K, V>>,
        key: &[u8],
        removed_value: &mut Option<V>,
    ) -> Arc<TrieNode<K, V>> {
        // Helper: after removing data from a node, merge with its sole child
        // if the node now has no data and exactly one child (path compression).
        // Skip if this node is the root to keep canonical form.
        fn maybe_merge<K, V>(
            node_fragment: Vec<u8>,
            children: BTreeMap<u8, Arc<TrieNode<K, V>>>,
            is_root: bool,
        ) -> Arc<TrieNode<K, V>> {
            if !is_root && children.len() == 1 {
                let (only_byte, only_child) = children.iter().next().unwrap();
                let mut merged_fragment = node_fragment;
                merged_fragment.push(*only_byte);
                merged_fragment.extend_from_slice(&only_child.key_fragment);
                let mut merged = TrieNode::new(merged_fragment);
                merged.children = only_child.children.clone();
                merged.data = only_child.data.clone();
                return Arc::new(merged);
            }
            let mut new_node = TrieNode::new(node_fragment);
            new_node.children = children;
            Arc::new(new_node)
        }

        let is_root = Arc::ptr_eq(node, &self.root);

        if key.is_empty() {
            // Key is fully consumed but we haven't matched node's fragment yet.
            // Only match if the fragment is also empty (i.e. key truly ends here).
            if !node.key_fragment.is_empty() || node.data.is_none() {
                return Arc::clone(node);
            }
            if let Some(kvp) = &node.data {
                *removed_value = Some((*kvp.value).clone());
            }
            if node.children.is_empty() {
                return Arc::new(TrieNode::new(Vec::new()));
            }
            return maybe_merge(node.key_fragment.clone(), node.children.clone(), is_root);
        }

        let common_len = prefix_match(key, &node.key_fragment);

        if common_len < node.key_fragment.len() {
            return Arc::clone(node);
        }

        let remaining = &key[common_len..];

        if remaining.is_empty() {
            if node.data.is_none() {
                return Arc::clone(node);
            }
            if let Some(kvp) = &node.data {
                *removed_value = Some((*kvp.value).clone());
            }
            if node.children.is_empty() {
                return Arc::new(TrieNode::new(Vec::new()));
            }
            return maybe_merge(node.key_fragment.clone(), node.children.clone(), is_root);
        }

        // We need to go deeper into the trie
        let next_byte = remaining[0];

        // Check if there's a child matching the next byte
        if let Some(child) = node.children.get(&next_byte) {
            // Recursively remove from the child
            let new_child = self.remove_recursive(child, &remaining[1..], removed_value);

            // If nothing was removed, return the original node
            if removed_value.is_none() {
                return Arc::clone(node);
            }

            // Create a new children map for path copying
            let mut new_children = node.children.clone();

            // Check if the child should be completely removed or replaced
            if new_child.key_fragment.is_empty()
                && new_child.children.is_empty()
                && new_child.data.is_none()
            {
                // Child is empty, remove it
                new_children.remove(&next_byte);
            } else {
                // Replace the child
                new_children.insert(next_byte, new_child);
            }

            // If this node has no value and only one child after removal,
            // merge them (path compression). Skip root to keep canonical form
            // (root always has empty fragment) and preserve structural sharing.
            if node.data.is_none() && new_children.len() == 1
                && !Arc::ptr_eq(node, &self.root)
            {
                let (only_byte, only_child) = new_children.iter().next().unwrap();

                let mut merged_fragment = node.key_fragment.clone();
                merged_fragment.push(*only_byte);
                merged_fragment.extend_from_slice(&only_child.key_fragment);

                let mut merged_node = TrieNode::new(merged_fragment);
                merged_node.children = only_child.children.clone();
                merged_node.data = only_child.data.clone();

                return Arc::new(merged_node);
            }

            // Create a new node with the updated children
            let mut new_node = TrieNode::new(node.key_fragment.clone());
            new_node.data = node.data.clone();
            new_node.children = new_children;

            return Arc::new(new_node);
        }

        // No matching child, key doesn't exist in the trie
        Arc::clone(node)
    }
}

// Implementation for types with a Default constructor
impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K> + Default> Default for Trie<K, V, KC> {
    fn default() -> Self {
        Self::new_with_converter(KC::default())
    }
}

// Generic new method for any key converter type
impl<K: Clone + Hash + Eq, V, KC: KeyToBytes<K>> Trie<K, V, KC> {
    /// Creates a new, empty trie with the default key converter type.
    ///
    /// This generic method creates a new trie with the default key converter.
    /// For string-like keys, you can use `new_str_key()` and for byte-slice keys,
    /// you can use `new_bytes_key()`.
    pub fn new() -> Self
    where
        KC: Default,
    {
        Self::new_with_converter(KC::default())
    }
}

// Implement PartialEq to enable efficient comparison of tries
impl<K: Clone + Hash + Eq, V: Hash + Eq + Clone, KC: KeyToBytes<K>> PartialEq for Trie<K, V, KC> {
    fn eq(&self, other: &Self) -> bool {
        // Fast path: same Arc means identical content
        if Arc::ptr_eq(&self.root, &other.root) {
            return true;
        }

        // Different sizes means definitely not equal
        if self.size != other.size {
            return false;
        }

        // Fast reject: different hashes means definitely not equal
        if self.root.hash() != other.root.hash() {
            return false;
        }

        // Hashes match — confirm with deep structural comparison
        // to avoid false positives from hash collisions
        self.root.deep_eq(&other.root)
    }
}

// Implement Eq for types where it makes sense
impl<K: Clone + Hash + Eq, V: Hash + Eq + Clone, KC: KeyToBytes<K>> Eq for Trie<K, V, KC> {}

#[cfg(feature = "serde")]
mod serde_impl {
    use super::*;
    use serde::de::{Deserialize, Deserializer, SeqAccess, Visitor};
    use serde::ser::{Serialize, SerializeSeq, Serializer};

    impl<K, V, KC> Serialize for Trie<K, V, KC>
    where
        K: Clone + Hash + Eq + Serialize,
        V: Clone + Serialize,
        KC: KeyToBytes<K>,
    {
        fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
            let mut seq = serializer.serialize_seq(Some(self.len()))?;
            for (k, v) in self.iter() {
                seq.serialize_element(&(k, v))?;
            }
            seq.end()
        }
    }

    impl<'de, K, V, KC> Deserialize<'de> for Trie<K, V, KC>
    where
        K: Clone + Hash + Eq + Deserialize<'de>,
        V: Clone + Deserialize<'de>,
        KC: KeyToBytes<K> + Default,
    {
        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
            struct TrieVisitor<K, V, KC>(std::marker::PhantomData<(K, V, KC)>);

            impl<'de, K, V, KC> Visitor<'de> for TrieVisitor<K, V, KC>
            where
                K: Clone + Hash + Eq + Deserialize<'de>,
                V: Clone + Deserialize<'de>,
                KC: KeyToBytes<K> + Default,
            {
                type Value = Trie<K, V, KC>;

                fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                    formatter.write_str("a sequence of key-value pairs")
                }

                fn visit_seq<A: SeqAccess<'de>>(
                    self,
                    mut seq: A,
                ) -> Result<Self::Value, A::Error> {
                    let mut trie = Trie::new_with_converter(KC::default());
                    while let Some((k, v)) = seq.next_element::<(K, V)>()? {
                        trie = trie.insert(k, v);
                    }
                    Ok(trie)
                }
            }

            deserializer.deserialize_seq(TrieVisitor(std::marker::PhantomData))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::key_converter::{BytesKeyConverter, StrKeyConverter};
    use std::sync::Arc;
    use std::collections::HashSet;

    #[test]
    fn test_new_trie_str_keys() {
        let trie = Trie::<String, i32, StrKeyConverter<String>>::new_str_key();
        assert!(trie.is_empty());
        assert_eq!(trie.len(), 0);
    }

    #[test]
    fn test_new_trie_bytes_keys() {
        let trie = Trie::<Vec<u8>, i32, BytesKeyConverter<Vec<u8>>>::new_bytes_key();
        assert!(trie.is_empty());
        assert_eq!(trie.len(), 0);
    }

    #[test]
    fn test_get_nonexistent_str() {
        let trie = Trie::<String, i32, StrKeyConverter<String>>::new_str_key();
        // To use get(&K), we need &String
        assert_eq!(trie.get(&"hello".to_string()), None);
    }

    #[test]
    fn test_get_nonexistent_bytes() {
        let trie = Trie::<Vec<u8>, i32, BytesKeyConverter<Vec<u8>>>::new_bytes_key();
        assert_eq!(trie.get(&b"hello".to_vec()), None);
    }

    #[test]
    fn test_insert_and_get_str() {
        let trie = Trie::<String, i32, StrKeyConverter<String>>::new_str_key();
        let trie = trie.insert("hello".to_string(), 42);

        assert_eq!(trie.len(), 1);
        assert_eq!(trie.get(&"hello".to_string()), Some(&42));
        assert_eq!(trie.get(&"world".to_string()), None);
    }

    #[test]
    fn test_insert_and_get_bytes() {
        let trie = Trie::<Vec<u8>, i32, BytesKeyConverter<Vec<u8>>>::new_bytes_key();
        let trie = trie.insert(b"hello".to_vec(), 42);

        assert_eq!(trie.len(), 1);
        assert_eq!(trie.get(&b"hello".to_vec()), Some(&42));
        assert_eq!(trie.get(&b"world".to_vec()), None);
    }

    #[test]
    fn test_insert_replace_str() {
        let trie = Trie::<String, i32, StrKeyConverter<String>>::new_str_key();
        let trie1 = trie.insert("hello".to_string(), 42);
        let trie2 = trie1.insert("hello".to_string(), 100);

        assert_eq!(trie1.len(), 1);
        assert_eq!(trie2.len(), 1);
        assert_eq!(trie1.get(&"hello".to_string()), Some(&42));
        assert_eq!(trie2.get(&"hello".to_string()), Some(&100));
    }

    #[test]
    fn test_insert_multiple_str() {
        let trie = Trie::<String, i32, StrKeyConverter<String>>::new_str_key();
        let trie1 = trie.insert("hello".to_string(), 42);
        let trie2 = trie1.insert("world".to_string(), 100);
        let trie3 = trie2.insert("hello world".to_string(), 200);

        assert_eq!(trie1.len(), 1);
        assert_eq!(trie2.len(), 2);
        assert_eq!(trie3.len(), 3);

        assert_eq!(trie3.get(&"hello".to_string()), Some(&42));
        assert_eq!(trie3.get(&"world".to_string()), Some(&100));
        assert_eq!(trie3.get(&"hello world".to_string()), Some(&200));
    }

    #[test]
    fn test_structural_sharing() {
        // Assuming StrKeyConverter for String keys
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        let trie1 = trie.insert("hello".to_string(), 42);
        let trie2 = trie1.insert("help".to_string(), 100);

        // The root nodes should be different
        assert!(!Arc::ptr_eq(&trie1.root, &trie2.root));

        // But they should share structure for the "hel" prefix
        // We need to navigate to the child nodes to check this
        let first_child1 = &trie1.root.children.get(&b'h').unwrap();
        let first_child2 = &trie2.root.children.get(&b'h').unwrap();

        // The 'h' nodes should be different since we modified this path
        assert!(!Arc::ptr_eq(first_child1, first_child2));

        // Now insert a completely different prefix
        let trie3 = trie2.insert("world".to_string(), 200);

        // The "hel" subtree should be shared between trie2 and trie3
        let prefix_node2 = &trie2.root.children.get(&b'h').unwrap();
        let prefix_node3 = &trie3.root.children.get(&b'h').unwrap();

        // Verify same pointer (same memory location) - structural sharing!
        assert!(Arc::ptr_eq(prefix_node2, prefix_node3));
    }

    #[test]
    fn test_hash_caching_with_trie_operations() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        
        // Insert a key and check that the hash is cached
        let trie1 = trie.insert("hello".to_string(), 42);
        let root_hash1 = trie1.root.hash();
        
        // Verify hash is cached
        assert_eq!(trie1.root.cached_hash.get(), Some(&root_hash1));
        
        // Insert another key and check that it has a different hash
        let trie2 = trie1.insert("world".to_string(), 100);
        let root_hash2 = trie2.root.hash();
        
        // The hashes should be different
        assert_ne!(root_hash1, root_hash2);
        
        // But both should be cached
        assert_eq!(trie1.root.cached_hash.get(), Some(&root_hash1));
        assert_eq!(trie2.root.cached_hash.get(), Some(&root_hash2));
        
        // Remove a key
        let (trie3, _) = trie2.remove(&"hello".to_string());
        let root_hash3 = trie3.root.hash();
        
        // The hash should be different from trie2
        assert_ne!(root_hash2, root_hash3);
        
        // And should be cached
        assert_eq!(trie3.root.cached_hash.get(), Some(&root_hash3));
    }
    
    #[test]
    fn test_subtree_size_caching() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        
        // Insert keys
        let trie1 = trie.insert("hello".to_string(), 42);
        let trie2 = trie1.insert("world".to_string(), 100);
        let trie3 = trie2.insert("test".to_string(), 200);
        
        // Check subtree size calculation
        // Force subtree_size calculation to populate the cache
        let size = trie3.root.subtree_size();
        assert_eq!(size, 3);
        
        // Verify it's cached
        assert_eq!(trie3.root.cached_subtree_size.get(), Some(&3));
        
        // Remove a key
        let (trie4, _) = trie3.remove(&"world".to_string());
        
        // Subtree size should be updated
        assert_eq!(trie4.root.subtree_size(), 2);
        
        // And cached
        assert_eq!(trie4.root.cached_subtree_size.get(), Some(&2));
    }

    #[test]
    fn test_complex_insert_remove_operations() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        let mut expected_keys = HashSet::new();
        
        // Build a trie with multiple operations
        let mut current_trie = trie;
        
        // Insert a series of keys
        let keys = vec![
            "hello",
            "help",
            "hell",
            "helicopter",
            "helipad",
            "world",
            "work",
            "worker",
            "working",
            "workflow"
        ];
        
        for (i, key) in keys.iter().enumerate() {
            current_trie = current_trie.insert(key.to_string(), i as u32);
            expected_keys.insert(key.to_string());

            assert_eq!(current_trie.len(), expected_keys.len(),
                "Size mismatch after inserting {:?}", key);
            assert_eq!(current_trie.root.subtree_size(), expected_keys.len(),
                "Subtree size mismatch after inserting {:?}", key);

            for expected_key in &expected_keys {
                assert!(current_trie.contains_key(expected_key),
                    "Key {:?} missing after inserting {:?}", expected_key, key);
            }
        }

        let all_items: Vec<_> = current_trie.iter().collect();
        assert_eq!(all_items.len(), expected_keys.len());

        let remove_keys = vec!["hell", "work", "hello", "helicopter", "workflow"];
        for key in remove_keys.iter() {
            let (new_trie, removed) = current_trie.remove(&key.to_string());
            assert!(removed.is_some(), "Key {:?} should have been present", key);
            current_trie = new_trie;
            expected_keys.remove(&key.to_string());

            assert_eq!(current_trie.len(), expected_keys.len(),
                "Size mismatch after removing {:?}", key);
            assert_eq!(current_trie.root.subtree_size(), expected_keys.len(),
                "Subtree size mismatch after removing {:?}", key);

            for expected_key in &expected_keys {
                assert!(current_trie.contains_key(expected_key),
                    "Key {:?} missing after removing {:?}", expected_key, key);
            }
            assert!(!current_trie.contains_key(&key.to_string()));
        }
        
        // Final check all remaining keys
        let final_items: Vec<_> = current_trie.iter().collect();
        assert_eq!(final_items.len(), expected_keys.len());
    }
    
    #[test]
    fn test_node_splitting_with_caching() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        
        // First insert a long key
        let trie1 = trie.insert("application".to_string(), 1);
        
        // Get the hash of the root node
        let hash1 = trie1.root.hash();
        assert!(trie1.root.cached_hash.get().is_some());
        
        // Now insert a key that shares a prefix, forcing a split
        let trie2 = trie1.insert("apple".to_string(), 2);
        
        // The hash should be different
        let hash2 = trie2.root.hash();
        assert_ne!(hash1, hash2);
        
        // Cache should be populated
        assert_eq!(trie2.root.cached_hash.get(), Some(&hash2));
        
        // Check subtree size
        assert_eq!(trie2.root.subtree_size(), 2);
        assert_eq!(trie2.root.cached_subtree_size.get(), Some(&2));
        
        // Insert one more key with a different prefix
        let trie3 = trie2.insert("banana".to_string(), 3);
        
        // Check hash and subtree size
        let hash3 = trie3.root.hash();
        assert_ne!(hash2, hash3);
        assert_eq!(trie3.root.subtree_size(), 3);
        
        // Caches should be populated
        assert_eq!(trie3.root.cached_hash.get(), Some(&hash3));
        assert_eq!(trie3.root.cached_subtree_size.get(), Some(&3));
    }
    
    #[test]
    fn test_path_compression_with_removes() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        
        // Insert keys with a common prefix
        let trie = trie.insert("compute".to_string(), 1);
        let trie = trie.insert("computer".to_string(), 2);
        let trie = trie.insert("computing".to_string(), 3);
        let trie = trie.insert("computational".to_string(), 4);
        
        // Verify size
        assert_eq!(trie.len(), 4);
        
        // Get initial hash
        let initial_hash = trie.root.hash();
        
        // Remove a key that should cause path compression
        let (trie2, _) = trie.remove(&"computer".to_string());
        
        // Verify size
        assert_eq!(trie2.len(), 3);
        
        // Hash should have changed
        let new_hash = trie2.root.hash();
        assert_ne!(initial_hash, new_hash);
        
        // Caches should be updated for hash
        assert_eq!(trie2.root.cached_hash.get(), Some(&new_hash));
        
        // Force computation of subtree size to populate cache
        let subtree_size = trie2.root.subtree_size();
        assert_eq!(subtree_size, 3);
        assert_eq!(trie2.root.cached_subtree_size.get(), Some(&3));
        
        // All remaining keys should be accessible
        assert!(trie2.contains_key(&"compute".to_string()));
        assert!(trie2.contains_key(&"computing".to_string()));
        assert!(trie2.contains_key(&"computational".to_string()));
        assert!(!trie2.contains_key(&"computer".to_string()));
    }

    #[test]
    fn test_node_splitting() {
        // Assuming StrKeyConverter for String keys
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();

        // Insert a key
        let trie1 = trie.insert("alphabet".to_string(), 1);

        // Insert another with common prefix - should cause splitting
        let trie2 = trie1.insert("alpha".to_string(), 2);

        assert_eq!(trie2.get(&"alphabet".to_string()), Some(&1));
        assert_eq!(trie2.get(&"alpha".to_string()), Some(&2));
    }

    #[test]
    fn test_larger_key_first() {
        // Assuming StrKeyConverter for String keys
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();

        // First insert the longer key
        let trie1 = trie.insert("alphabet".to_string(), 1);

        // Then insert the shorter one
        let trie2 = trie1.insert("alpha".to_string(), 2);

        assert_eq!(trie2.get(&"alphabet".to_string()), Some(&1));
        assert_eq!(trie2.get(&"alpha".to_string()), Some(&2));
    }

    #[test]
    fn test_shorter_key_first() {
        // Assuming StrKeyConverter for String keys
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();

        // First insert the shorter key
        let trie1 = trie.insert("alpha".to_string(), 1);

        // Then insert the longer one
        let trie2 = trie1.insert("alphabet".to_string(), 2);

        assert_eq!(trie2.get(&"alpha".to_string()), Some(&1));
        assert_eq!(trie2.get(&"alphabet".to_string()), Some(&2));
    }

    #[test]
    fn test_remove_existing_str() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 42)
            .insert("world".to_string(), 100);

        assert_eq!(trie.len(), 2);

        // Remove an existing key
        let (trie2, removed) = trie.remove(&"hello".to_string());

        assert_eq!(removed, Some(42));
        assert_eq!(trie2.len(), 1);
        assert_eq!(trie2.get(&"hello".to_string()), None);
        assert_eq!(trie2.get(&"world".to_string()), Some(&100));

        // The original trie should be unchanged
        assert_eq!(trie.len(), 2);
        assert_eq!(trie.get(&"hello".to_string()), Some(&42));
    }

    #[test]
    fn test_remove_nonexistent_str() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 42);

        // Remove a non-existent key
        let (trie2, removed) = trie.remove(&"world".to_string());

        assert_eq!(removed, None);
        assert_eq!(trie2.len(), 1);
        assert_eq!(trie2.get(&"hello".to_string()), Some(&42));

        // Removing from an empty trie
        let empty = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        let (empty2, removed) = empty.remove(&"anything".to_string());

        assert_eq!(removed, None);
        assert_eq!(empty2.len(), 0);
    }

    #[test]
    fn test_remove_with_compression() {
        // Assuming StrKeyConverter
        // Create a trie with keys that will cause path compression when removed
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("abc".to_string(), 1)
            .insert("abcde".to_string(), 2);

        // Remove the middle key, which should cause compression
        let (trie2, removed) = trie.remove(&"abc".to_string());

        assert_eq!(removed, Some(1));
        assert_eq!(trie2.len(), 1);
        assert_eq!(trie2.get(&"abc".to_string()), None);
        assert_eq!(trie2.get(&"abcde".to_string()), Some(&2));
    }

    #[test]
    fn test_remove_structural_sharing() {
        // Assuming StrKeyConverter
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("world".to_string(), 3);

        // Remove a key from one branch
        let (trie2, _) = trie.remove(&"world".to_string());

        // The "hel" branch should be shared between the two tries
        let h_node1 = trie.root.children.get(&b'h').unwrap();
        let h_node2 = trie2.root.children.get(&b'h').unwrap();

        // The 'h' nodes should be the same (structural sharing)
        assert!(Arc::ptr_eq(h_node1, h_node2));

        // But removing from a branch should create new nodes along that path
        let (trie3, _) = trie.remove(&"hello".to_string());

        let h_node3 = trie3.root.children.get(&b'h').unwrap();

        // The 'h' nodes should be different now
        assert!(!Arc::ptr_eq(h_node1, h_node3));
    }

    #[test]
    fn test_remove_produces_canonical_structure() {
        // Regression test: removing a sibling should merge the remaining
        // single child back into its parent (path compression).
        // A trie built by insertion should be structurally identical to one
        // that went through insert+remove to reach the same content.
        let direct = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("abc".to_string(), 1);

        let via_remove = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("abc".to_string(), 1)
            .insert("abd".to_string(), 2)
            .remove(&"abd".to_string()).0;

        // Both contain exactly {"abc": 1}
        assert_eq!(direct.len(), via_remove.len());
        assert_eq!(direct.get(&"abc".to_string()), via_remove.get(&"abc".to_string()));

        // The structural hash should be identical if path compression works
        assert_eq!(direct.root.hash(), via_remove.root.hash(),
            "Trie structure should be canonical after remove (path compression failed)");
    }

    #[test]
    fn test_remove_merge_with_data_child() {
        // Specifically tests the case where the remaining child has data
        // (the old code refused to merge in this case)
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("compute".to_string(), 1)
            .insert("computer".to_string(), 2)
            .insert("computing".to_string(), 3);

        let (after_remove, removed) = trie.remove(&"computing".to_string());
        assert_eq!(removed, Some(3));
        assert_eq!(after_remove.len(), 2);

        // Build the same trie directly
        let direct = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("compute".to_string(), 1)
            .insert("computer".to_string(), 2);

        assert_eq!(direct.root.hash(), after_remove.root.hash(),
            "After removing 'computing', structure should match direct construction");
    }

    #[test]
    fn test_remove_merge_with_multi_child() {
        // Tests merging when the remaining child has multiple children of its own
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("aa".to_string(), 1)
            .insert("ab".to_string(), 2)
            .insert("b".to_string(), 3);

        let (after_remove, _) = trie.remove(&"b".to_string());

        let direct = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("aa".to_string(), 1)
            .insert("ab".to_string(), 2);

        assert_eq!(direct.root.hash(), after_remove.root.hash(),
            "After removing 'b', structure should match direct construction");
    }

    #[test]
    fn test_get_prefix_of_existing_key_returns_none() {
        // Regression: get() must check that the node's key_fragment was fully
        // matched before returning data.  After path compression, a node can
        // hold data with a non-empty fragment; looking up a shorter key that
        // dispatches to that node should NOT return the node's data.
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1);

        // "h" is a prefix, not a stored key
        assert_eq!(trie.get(&"h".to_string()), None);
        assert!(!trie.contains_key(&"h".to_string()));
    }

    #[test]
    fn test_get_after_remove_with_merge() {
        // Regression: removing data from a node that then merges with its
        // sole child must not make the removed key retrievable via get().
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hell".to_string(), 1)
            .insert("hello".to_string(), 2);

        let (after, removed) = trie.remove(&"hell".to_string());
        assert_eq!(removed, Some(1));
        assert_eq!(after.get(&"hell".to_string()), None);
        assert_eq!(after.get(&"hello".to_string()), Some(&2));
    }

    #[test]
    fn test_remove_nonexistent_prefix_key() {
        // Regression: remove_recursive must not strip data from a node with
        // a non-empty key_fragment when the lookup key is exhausted.
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1);

        let (after, removed) = trie.remove(&"h".to_string());
        assert_eq!(removed, None);
        assert_eq!(after.get(&"hello".to_string()), Some(&1));
        assert_eq!(after.len(), 1);
    }

    #[test]
    fn test_remove_data_merges_with_sole_child() {
        // Regression: when removing data from a node that has exactly one
        // child (and no other data), path compression should merge them.
        let direct = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert(":\0".to_string(), 0);

        let via_remove = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert(":".to_string(), 0)
            .insert(":\0".to_string(), 0)
            .remove(&":".to_string()).0;

        assert_eq!(direct, via_remove);
    }

    #[test]
    fn test_longest_prefix_match_basic() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("/".to_string(), 0)
            .insert("/api".to_string(), 1)
            .insert("/api/users".to_string(), 2)
            .insert("/static".to_string(), 3);

        // Exact matches
        assert_eq!(trie.longest_prefix_match(&"/".to_string()), Some((&"/".to_string(), &0)));
        assert_eq!(trie.longest_prefix_match(&"/api".to_string()), Some((&"/api".to_string(), &1)));
        assert_eq!(trie.longest_prefix_match(&"/api/users".to_string()), Some((&"/api/users".to_string(), &2)));

        // Partial matches — returns longest stored prefix
        assert_eq!(trie.longest_prefix_match(&"/api/posts".to_string()), Some((&"/api".to_string(), &1)));
        assert_eq!(trie.longest_prefix_match(&"/api/users/123".to_string()), Some((&"/api/users".to_string(), &2)));
        assert_eq!(trie.longest_prefix_match(&"/static/css/main.css".to_string()), Some((&"/static".to_string(), &3)));
        assert_eq!(trie.longest_prefix_match(&"/other".to_string()), Some((&"/".to_string(), &0)));

        // No match
        assert_eq!(trie.longest_prefix_match(&"nope".to_string()), None);
    }

    #[test]
    fn test_longest_prefix_match_empty_key() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("".to_string(), 42)
            .insert("hello".to_string(), 1);

        assert_eq!(trie.longest_prefix_match(&"".to_string()), Some((&"".to_string(), &42)));
        assert_eq!(trie.longest_prefix_match(&"hello".to_string()), Some((&"hello".to_string(), &1)));
        assert_eq!(trie.longest_prefix_match(&"help".to_string()), Some((&"".to_string(), &42)));
    }

    #[test]
    fn test_longest_prefix_match_no_match() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("abc".to_string(), 1)
            .insert("def".to_string(), 2);

        // "xyz" shares no prefix with any stored key
        assert_eq!(trie.longest_prefix_match(&"xyz".to_string()), None);
        // "ab" is a prefix of "abc" but "ab" itself is not stored
        assert_eq!(trie.longest_prefix_match(&"ab".to_string()), None);
    }

    #[test]
    fn test_longest_prefix_match_empty_trie() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        assert_eq!(trie.longest_prefix_match(&"anything".to_string()), None);
    }

    #[test]
    fn test_iter_prefix() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("world".to_string(), 3)
            .insert("helicopter".to_string(), 4);

        let hel: Vec<_> = trie.iter_prefix("hel".to_string()).collect();
        assert_eq!(hel.len(), 3);
        assert!(hel.contains(&("hello".to_string(), 1)));
        assert!(hel.contains(&("help".to_string(), 2)));
        assert!(hel.contains(&("helicopter".to_string(), 4)));

        let world: Vec<_> = trie.iter_prefix("world".to_string()).collect();
        assert_eq!(world.len(), 1);
        assert!(world.contains(&("world".to_string(), 3)));

        let empty: Vec<_> = trie.iter_prefix("xyz".to_string()).collect();
        assert!(empty.is_empty());
    }

    #[test]
    fn test_from_iterator() {
        let entries = vec![
            ("hello".to_string(), 1u32),
            ("help".to_string(), 2),
            ("world".to_string(), 3),
        ];
        let trie: Trie<String, u32, StrKeyConverter<String>> = entries.into_iter().collect();

        assert_eq!(trie.len(), 3);
        assert_eq!(trie.get(&"hello".to_string()), Some(&1));
        assert_eq!(trie.get(&"help".to_string()), Some(&2));
        assert_eq!(trie.get(&"world".to_string()), Some(&3));
    }

    #[test]
    fn test_from_iterator_empty() {
        let entries: Vec<(String, u32)> = vec![];
        let trie: Trie<String, u32, StrKeyConverter<String>> = entries.into_iter().collect();
        assert!(trie.is_empty());
    }

    #[test]
    fn test_from_iterator_duplicates_last_wins() {
        let entries = vec![
            ("key".to_string(), 1u32),
            ("key".to_string(), 2),
        ];
        let trie: Trie<String, u32, StrKeyConverter<String>> = entries.into_iter().collect();

        assert_eq!(trie.len(), 1);
        assert_eq!(trie.get(&"key".to_string()), Some(&2));
    }

    #[test]
    fn test_equality_uses_deep_comparison() {
        // Verify that PartialEq performs deep comparison, not just hash comparison.
        // Two tries with identical content must be equal.
        let trie1 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("world".to_string(), 3);

        let trie2 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("world".to_string(), 3);

        // Same content → equal
        assert_eq!(trie1, trie2);

        // Different value → not equal
        let trie3 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 999)
            .insert("help".to_string(), 2)
            .insert("world".to_string(), 3);

        assert_ne!(trie1, trie3);

        // Different keys → not equal
        let trie4 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("other".to_string(), 3);

        assert_ne!(trie1, trie4);

        // Empty tries are equal
        let empty1 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        let empty2 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key();
        assert_eq!(empty1, empty2);
    }

    #[test]
    fn test_deep_eq_on_nodes_directly() {
        use crate::node::TrieNode;

        // Two nodes with same structure should be deep_eq
        let mut node1 = TrieNode::<String, u32>::with_key_value(
            vec![1, 2], "key".to_string(), 42,
        );
        let child1 = Arc::new(TrieNode::<String, u32>::with_key_value(
            vec![3], "child".to_string(), 99,
        ));
        node1.children.insert(3, child1);

        let mut node2 = TrieNode::<String, u32>::with_key_value(
            vec![1, 2], "key".to_string(), 42,
        );
        let child2 = Arc::new(TrieNode::<String, u32>::with_key_value(
            vec![3], "child".to_string(), 99,
        ));
        node2.children.insert(3, child2);

        assert!(node1.deep_eq(&node2));

        // Different value → not deep_eq
        let mut node3 = TrieNode::<String, u32>::with_key_value(
            vec![1, 2], "key".to_string(), 100,
        );
        let child3 = Arc::new(TrieNode::<String, u32>::with_key_value(
            vec![3], "child".to_string(), 99,
        ));
        node3.children.insert(3, child3);

        assert!(!node1.deep_eq(&node3));
    }

    #[test]
    fn test_insert_prefix_of_existing_key() {
        // Regression: inserting "hell" after "hello" must not overwrite "hello".
        // The bug was that insert_recursive returned early when key_bytes was
        // exhausted without checking if the current node's fragment was also
        // consumed, silently replacing data at the wrong node.
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1)
            .insert("help".to_string(), 2)
            .insert("hell".to_string(), 3);

        assert_eq!(trie.len(), 3);
        assert_eq!(trie.get(&"hello".to_string()), Some(&1));
        assert_eq!(trie.get(&"help".to_string()), Some(&2));
        assert_eq!(trie.get(&"hell".to_string()), Some(&3));
    }

    #[test]
    fn test_insert_empty_key() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("".to_string(), 42)
            .insert("hello".to_string(), 1);

        assert_eq!(trie.len(), 2);
        assert_eq!(trie.get(&"".to_string()), Some(&42));
        assert_eq!(trie.get(&"hello".to_string()), Some(&1));

        // Replace empty key
        let trie2 = trie.insert("".to_string(), 99);
        assert_eq!(trie2.len(), 2);
        assert_eq!(trie2.get(&"".to_string()), Some(&99));
    }

    // --- Structural sharing: deep nesting and Arc refcount ---

    #[test]
    fn test_structural_sharing_deep_nesting() {
        // Build a trie with many branches so we can verify sharing at depth
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("aaa".to_string(), 1)
            .insert("aab".to_string(), 2)
            .insert("aba".to_string(), 3)
            .insert("bbb".to_string(), 4);

        // Insert into a completely separate branch
        let trie2 = trie.insert("ccc".to_string(), 5);

        // The entire 'a' subtree should be shared (same Arc)
        let a1 = trie.root.children.get(&b'a').unwrap();
        let a2 = trie2.root.children.get(&b'a').unwrap();
        assert!(Arc::ptr_eq(a1, a2));

        // The entire 'b' subtree should also be shared
        let b1 = trie.root.children.get(&b'b').unwrap();
        let b2 = trie2.root.children.get(&b'b').unwrap();
        assert!(Arc::ptr_eq(b1, b2));

        // Now insert into the 'a' branch — 'b' should still be shared
        let trie3 = trie2.insert("aac".to_string(), 6);
        let b2 = trie2.root.children.get(&b'b').unwrap();
        let b3 = trie3.root.children.get(&b'b').unwrap();
        assert!(Arc::ptr_eq(b2, b3));

        // But 'a' should NOT be shared (we modified that branch)
        let a2 = trie2.root.children.get(&b'a').unwrap();
        let a3 = trie3.root.children.get(&b'a').unwrap();
        assert!(!Arc::ptr_eq(a2, a3));
    }

    #[test]
    fn test_structural_sharing_preserves_old_versions() {
        // The whole point of immutability: old snapshots remain valid
        let v1 = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("hello".to_string(), 1);
        let v2 = v1.insert("world".to_string(), 2);
        let (v3, _) = v2.remove(&"hello".to_string());

        // v1 still has its original data
        assert_eq!(v1.len(), 1);
        assert_eq!(v1.get(&"hello".to_string()), Some(&1));
        assert_eq!(v1.get(&"world".to_string()), None);

        // v2 has both
        assert_eq!(v2.len(), 2);
        assert_eq!(v2.get(&"hello".to_string()), Some(&1));
        assert_eq!(v2.get(&"world".to_string()), Some(&2));

        // v3 only has "world"
        assert_eq!(v3.len(), 1);
        assert_eq!(v3.get(&"hello".to_string()), None);
        assert_eq!(v3.get(&"world".to_string()), Some(&2));
    }

    #[test]
    fn test_structural_sharing_arc_refcounts() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("alpha".to_string(), 1)
            .insert("beta".to_string(), 2);

        let a_node = trie.root.children.get(&b'a').unwrap();
        let initial_count = Arc::strong_count(a_node);

        // Each new version that shares this subtree increments the refcount
        let trie2 = trie.insert("gamma".to_string(), 3);
        let a_node_shared = trie2.root.children.get(&b'a').unwrap();

        assert!(Arc::ptr_eq(a_node, a_node_shared));
        // Both trie and trie2 hold a reference, plus we have our local refs
        assert!(Arc::strong_count(a_node) > initial_count);

        // Inserting into the 'a' branch creates a new node, doesn't bump the old one further
        let trie3 = trie2.insert("alpha2".to_string(), 4);
        let a_node_new = trie3.root.children.get(&b'a').unwrap();
        assert!(!Arc::ptr_eq(a_node, a_node_new));
    }

    // --- BTreeMap: deterministic lexicographic iteration ---

    #[test]
    fn test_iteration_is_lexicographic() {
        // Insert in deliberately non-alphabetical order
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("zebra".to_string(), 1)
            .insert("apple".to_string(), 2)
            .insert("mango".to_string(), 3)
            .insert("banana".to_string(), 4)
            .insert("apricot".to_string(), 5);

        let keys: Vec<String> = trie.iter().map(|(k, _)| k.clone()).collect();
        let mut sorted = keys.clone();
        sorted.sort();
        assert_eq!(keys, sorted, "iteration should be in lexicographic order");
    }

    #[test]
    fn test_iteration_order_with_shared_prefixes() {
        // Keys that share prefixes — tests that BTreeMap child ordering
        // produces correct DFS lexicographic order
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("ab".to_string(), 1)
            .insert("abc".to_string(), 2)
            .insert("abd".to_string(), 3)
            .insert("a".to_string(), 4)
            .insert("b".to_string(), 5);

        let keys: Vec<String> = trie.iter().map(|(k, _)| k.clone()).collect();
        assert_eq!(keys, vec!["a", "ab", "abc", "abd", "b"]);
    }

    #[test]
    fn test_iter_prefix_order_is_lexicographic() {
        let trie = Trie::<String, u32, StrKeyConverter<String>>::new_str_key()
            .insert("/z/1".to_string(), 1)
            .insert("/a/2".to_string(), 2)
            .insert("/a/1".to_string(), 3)
            .insert("/m/1".to_string(), 4);

        let keys: Vec<String> = trie.iter_prefix("/a".to_string())
            .map(|(k, _)| k.clone()).collect();
        let mut sorted = keys.clone();
        sorted.sort();
        assert_eq!(keys, sorted, "iter_prefix should be in lexicographic order");
    }
}