agentstategraph 1.2.8

//! State tree operations — get, set, delete values by path.
//!
//! The state tree is a Merkle DAG of Objects. Modifications create new
//! objects along the modified path (structural sharing), leaving all
//! other subtrees untouched.

use std::collections::BTreeMap;

use agentstategraph_core::{Atom, Node, Object, ObjectId, PathComponent, StatePath};
use agentstategraph_storage::{ObjectStore, StorageError};

/// Absolute ceiling on how deep `tree_list_paths` / `tree_search_values`
/// will ever recurse, regardless of the caller-supplied `max_depth`.
///
/// Set to twice `agentstategraph_core::MAX_PATH_DEPTH` (64 → 128) so we
/// can still enumerate legacy data that was written by older ASG versions
/// without the current path-depth cap. A caller passing `usize::MAX` is
/// clamped to this value, preventing unbounded traversal / OOM.
/// (security threat model v2, F2)
pub const LIST_PATHS_HARD_CEILING: usize = 128;

/// Absolute cap on the number of results returned by `tree_list_paths`
/// and `tree_search_values`. Once this many entries have been collected,
/// traversal short-circuits and the extras are dropped. Callers receive
/// the truncated vec without an explicit marker — keep the API simple.
/// (security threat model v2, F2)
pub const LIST_PATHS_MAX_RESULTS: usize = 10_000;

/// Read a value from the state tree at the given path.
/// Returns the Object at that path, or an error if the path doesn't exist.
pub fn tree_get(
    store: &dyn ObjectStore,
    root: &ObjectId,
    path: &StatePath,
) -> Result<Object, TreeError> {
    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;

    if path.is_root() {
        return Ok(root_obj);
    }

    let mut current = root_obj;
    for (i, component) in path.components().iter().enumerate() {
        current = navigate(&current, component, store, path, i)?;
    }

    Ok(current)
}

/// Set a value in the state tree at the given path.
/// Returns the new root ObjectId after the modification.
/// Creates intermediate maps as needed for new paths.
///
/// This operation uses structural sharing — only objects along the
/// modified path are new; all other subtrees retain their ObjectIds.
pub fn tree_set(
    store: &dyn ObjectStore,
    root: &ObjectId,
    path: &StatePath,
    value: &Object,
) -> Result<ObjectId, TreeError> {
    // Store the new value and get its ID
    let value_id = store.put_object(value)?;

    if path.is_root() {
        return Ok(value_id);
    }

    // Load the root
    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;

    // Recursively set the value, rebuilding the path from leaf to root
    let new_root = set_recursive(store, &root_obj, path.components(), 0, value_id)?;
    let new_root_id = store.put_object(&new_root)?;
    Ok(new_root_id)
}

/// Delete a value from the state tree at the given path.
/// Returns the new root ObjectId after the deletion.
pub fn tree_delete(
    store: &dyn ObjectStore,
    root: &ObjectId,
    path: &StatePath,
) -> Result<ObjectId, TreeError> {
    if path.is_root() {
        // Deleting root → return empty map
        let empty = Object::empty_map();
        let id = store.put_object(&empty)?;
        return Ok(id);
    }

    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;

    let new_root = delete_recursive(store, &root_obj, path.components(), 0)?;
    let new_root_id = store.put_object(&new_root)?;
    Ok(new_root_id)
}

/// Convert an Object to a serde_json::Value for user-facing output.
/// Recursively resolves all ObjectId references.
pub fn tree_to_json(store: &dyn ObjectStore, obj: &Object) -> Result<serde_json::Value, TreeError> {
    match obj {
        Object::Atom(atom) => Ok(atom_to_json(atom)),
        Object::Node(node) => match node {
            Node::Map(entries) => {
                let mut map = serde_json::Map::new();
                for (key, child_id) in entries {
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    map.insert(key.clone(), tree_to_json(store, &child)?);
                }
                Ok(serde_json::Value::Object(map))
            }
            Node::List(items) => {
                let mut arr = Vec::new();
                for child_id in items {
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    arr.push(tree_to_json(store, &child)?);
                }
                Ok(serde_json::Value::Array(arr))
            }
            Node::Set(items) => {
                let mut arr = Vec::new();
                for child_id in items {
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    arr.push(tree_to_json(store, &child)?);
                }
                Ok(serde_json::Value::Array(arr))
            }
        },
    }
}

/// Like [`tree_to_json`] but stops at `max_depth` levels below `obj`: a Node at
/// the depth limit is rendered as a compact `{ "_truncated": true, ... }`
/// placeholder (its child keys, or length) instead of being recursed into, so
/// deep subtrees are never loaded or materialized. Used for cheap shallow reads
/// of a large ref (plan perf-slow-endpoints t-007). Depth is counted from
/// `obj`: `max_depth = 0` summarizes a Node immediately; an Atom is always
/// returned as its value.
pub fn tree_to_json_capped(
    store: &dyn ObjectStore,
    obj: &Object,
    max_depth: usize,
) -> Result<serde_json::Value, TreeError> {
    to_json_capped(store, obj, max_depth, 0)
}

fn to_json_capped(
    store: &dyn ObjectStore,
    obj: &Object,
    max_depth: usize,
    depth: usize,
) -> Result<serde_json::Value, TreeError> {
    match obj {
        Object::Atom(atom) => Ok(atom_to_json(atom)),
        Object::Node(node) => {
            if depth >= max_depth {
                // Boundary: summarize from the node itself — its keys/length are
                // in hand, so we never load (or serialize) any child object.
                let summary = match node {
                    Node::Map(entries) => serde_json::json!({
                        "_truncated": true,
                        "_keys": entries.keys().cloned().collect::<Vec<_>>(),
                    }),
                    Node::List(items) | Node::Set(items) => serde_json::json!({
                        "_truncated": true,
                        "_len": items.len(),
                    }),
                };
                return Ok(summary);
            }
            match node {
                Node::Map(entries) => {
                    let mut map = serde_json::Map::new();
                    for (key, child_id) in entries {
                        let child = store
                            .get_object(child_id)?
                            .ok_or(TreeError::ObjectNotFound(*child_id))?;
                        map.insert(
                            key.clone(),
                            to_json_capped(store, &child, max_depth, depth + 1)?,
                        );
                    }
                    Ok(serde_json::Value::Object(map))
                }
                Node::List(items) | Node::Set(items) => {
                    let mut arr = Vec::new();
                    for child_id in items {
                        let child = store
                            .get_object(child_id)?
                            .ok_or(TreeError::ObjectNotFound(*child_id))?;
                        arr.push(to_json_capped(store, &child, max_depth, depth + 1)?);
                    }
                    Ok(serde_json::Value::Array(arr))
                }
            }
        }
    }
}

/// List all paths in the state tree under a given prefix.
/// Returns leaf paths (paths that point to atoms/values, not intermediate maps).
pub fn tree_list_paths(
    store: &dyn ObjectStore,
    root: &ObjectId,
    prefix: &str,
    max_depth: usize,
) -> Result<Vec<String>, TreeError> {
    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;

    // If prefix is non-empty, navigate to the subtree first
    let (start_obj, base_path) = if prefix.is_empty() || prefix == "/" {
        (root_obj, String::new())
    } else {
        let path = StatePath::parse(prefix).map_err(|e| TreeError::PathNotFound(e.to_string()))?;
        let obj = tree_get(store, root, &path)?;
        (obj, prefix.to_string())
    };

    // Clamp the caller-supplied depth to the hard ceiling so `usize::MAX`
    // cannot produce unbounded recursion. See LIST_PATHS_HARD_CEILING.
    let depth = max_depth.min(LIST_PATHS_HARD_CEILING);
    let mut paths = Vec::new();
    collect_paths(store, &start_obj, &base_path, depth, 0, &mut paths)?;
    // Truncate to the absolute result cap — a malicious/wide tree cannot
    // make us materialize more than LIST_PATHS_MAX_RESULTS entries.
    if paths.len() > LIST_PATHS_MAX_RESULTS {
        paths.truncate(LIST_PATHS_MAX_RESULTS);
    }
    Ok(paths)
}

fn collect_paths(
    store: &dyn ObjectStore,
    obj: &Object,
    current_path: &str,
    max_depth: usize,
    depth: usize,
    paths: &mut Vec<String>,
) -> Result<(), TreeError> {
    if depth > max_depth {
        return Ok(());
    }
    if paths.len() >= LIST_PATHS_MAX_RESULTS {
        return Ok(());
    }

    match obj {
        Object::Atom(_) => {
            let path = if current_path.is_empty() {
                "/".to_string()
            } else {
                current_path.to_string()
            };
            paths.push(path);
        }
        Object::Node(node) => match node {
            Node::Map(entries) => {
                if entries.is_empty() {
                    let path = if current_path.is_empty() {
                        "/".to_string()
                    } else {
                        current_path.to_string()
                    };
                    paths.push(path);
                } else {
                    for (key, child_id) in entries {
                        let child_path = format!("{}/{}", current_path, key);
                        let child = store
                            .get_object(child_id)?
                            .ok_or(TreeError::ObjectNotFound(*child_id))?;
                        collect_paths(store, &child, &child_path, max_depth, depth + 1, paths)?;
                    }
                }
            }
            Node::List(items) => {
                for (i, child_id) in items.iter().enumerate() {
                    let child_path = format!("{}/{}", current_path, i);
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    collect_paths(store, &child, &child_path, max_depth, depth + 1, paths)?;
                }
            }
            Node::Set(items) => {
                for (i, child_id) in items.iter().enumerate() {
                    let child_path = format!("{}/{}", current_path, i);
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    collect_paths(store, &child, &child_path, max_depth, depth + 1, paths)?;
                }
            }
        },
    }
    Ok(())
}

/// Search state values for a query string. Returns matching paths + values.
pub fn tree_search_values(
    store: &dyn ObjectStore,
    root: &ObjectId,
    query: &str,
    max_results: usize,
) -> Result<Vec<(String, String)>, TreeError> {
    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;

    let query_lower = query.to_lowercase();
    // Clamp the caller's max_results to LIST_PATHS_MAX_RESULTS so a huge
    // query can't force us to materialize an unbounded result set.
    let limit = max_results.min(LIST_PATHS_MAX_RESULTS);
    let mut results = Vec::new();
    search_recursive(store, &root_obj, "", &query_lower, limit, 0, &mut results)?;
    Ok(results)
}

fn search_recursive(
    store: &dyn ObjectStore,
    obj: &Object,
    current_path: &str,
    query: &str,
    max_results: usize,
    depth: usize,
    results: &mut Vec<(String, String)>,
) -> Result<(), TreeError> {
    if results.len() >= max_results {
        return Ok(());
    }
    // Hard ceiling on traversal depth — mirrors tree_list_paths so a
    // pathologically deep or cyclic-seeming graph can't infinite-loop.
    if depth > LIST_PATHS_HARD_CEILING {
        return Ok(());
    }

    match obj {
        Object::Atom(atom) => {
            let value_str = match atom {
                Atom::String(s) => s.clone(),
                Atom::Int(i) => i.to_string(),
                Atom::UInt(u) => u.to_string(),
                Atom::Float(f) => f.to_string(),
                Atom::Bool(b) => b.to_string(),
                _ => return Ok(()),
            };
            if value_str.to_lowercase().contains(query) {
                let path = if current_path.is_empty() {
                    "/".to_string()
                } else {
                    current_path.to_string()
                };
                results.push((path, value_str));
            }
        }
        Object::Node(node) => match node {
            Node::Map(entries) => {
                for (key, child_id) in entries {
                    if results.len() >= max_results {
                        break;
                    }
                    // Also match on key names
                    if key.to_lowercase().contains(query) {
                        let path = format!("{}/{}", current_path, key);
                        results.push((path.clone(), format!("[key match: {}]", key)));
                    }
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    let child_path = format!("{}/{}", current_path, key);
                    search_recursive(
                        store,
                        &child,
                        &child_path,
                        query,
                        max_results,
                        depth + 1,
                        results,
                    )?;
                }
            }
            Node::List(items) | Node::Set(items) => {
                for (i, child_id) in items.iter().enumerate() {
                    if results.len() >= max_results {
                        break;
                    }
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    let child_path = format!("{}/{}", current_path, i);
                    search_recursive(
                        store,
                        &child,
                        &child_path,
                        query,
                        max_results,
                        depth + 1,
                        results,
                    )?;
                }
            }
        },
    }
    Ok(())
}

/// Collect EVERY `(path, value)` string leaf of the tree at `root`, with no
/// query filter and no result cap — the one-time full walk that seeds the
/// leaf-value index (plan perf-slow-endpoints t-006). Unlike
/// [`tree_search_values`] it does not emit key-name matches; the index is over
/// values only. Depth is bounded by [`LIST_PATHS_HARD_CEILING`], mirroring the
/// other tree walks.
pub fn tree_collect_leaves(
    store: &dyn ObjectStore,
    root: &ObjectId,
) -> Result<Vec<(String, String)>, TreeError> {
    let root_obj = store
        .get_object(root)?
        .ok_or(TreeError::ObjectNotFound(*root))?;
    let mut out = Vec::new();
    collect_leaves_recursive(store, &root_obj, "", 0, &mut out)?;
    Ok(out)
}

fn collect_leaves_recursive(
    store: &dyn ObjectStore,
    obj: &Object,
    current_path: &str,
    depth: usize,
    out: &mut Vec<(String, String)>,
) -> Result<(), TreeError> {
    if depth > LIST_PATHS_HARD_CEILING {
        return Ok(());
    }
    match obj {
        Object::Atom(atom) => {
            let value_str = match atom {
                Atom::String(s) => s.clone(),
                Atom::Int(i) => i.to_string(),
                Atom::UInt(u) => u.to_string(),
                Atom::Float(f) => f.to_string(),
                Atom::Bool(b) => b.to_string(),
                _ => return Ok(()),
            };
            let path = if current_path.is_empty() {
                "/".to_string()
            } else {
                current_path.to_string()
            };
            out.push((path, value_str));
        }
        Object::Node(node) => match node {
            Node::Map(entries) => {
                for (key, child_id) in entries {
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    let child_path = format!("{}/{}", current_path, key);
                    collect_leaves_recursive(store, &child, &child_path, depth + 1, out)?;
                }
            }
            Node::List(items) | Node::Set(items) => {
                for (i, child_id) in items.iter().enumerate() {
                    let child = store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?;
                    let child_path = format!("{}/{}", current_path, i);
                    collect_leaves_recursive(store, &child, &child_path, depth + 1, out)?;
                }
            }
        },
    }
    Ok(())
}

/// Stringify an atom the way the value index/search does. `None` for atom
/// kinds that are not indexed (Null, Bytes).
fn atom_value_string(atom: &Atom) -> Option<String> {
    match atom {
        Atom::String(s) => Some(s.clone()),
        Atom::Int(i) => Some(i.to_string()),
        Atom::UInt(u) => Some(u.to_string()),
        Atom::Float(f) => Some(f.to_string()),
        Atom::Bool(b) => Some(b.to_string()),
        _ => None,
    }
}

/// Collect just the leaf PATHS under `obj` (for removals during a diff).
fn collect_leaf_paths(
    store: &dyn ObjectStore,
    obj: &Object,
    current_path: &str,
    depth: usize,
    out: &mut Vec<String>,
) -> Result<(), TreeError> {
    if depth > LIST_PATHS_HARD_CEILING {
        return Ok(());
    }
    match obj {
        Object::Atom(atom) => {
            if atom_value_string(atom).is_some() {
                out.push(if current_path.is_empty() {
                    "/".to_string()
                } else {
                    current_path.to_string()
                });
            }
        }
        Object::Node(node) => {
            let children: Vec<(String, ObjectId)> = match node {
                Node::Map(entries) => entries
                    .iter()
                    .map(|(k, id)| (format!("{}/{}", current_path, k), *id))
                    .collect(),
                Node::List(items) | Node::Set(items) => items
                    .iter()
                    .enumerate()
                    .map(|(i, id)| (format!("{}/{}", current_path, i), *id))
                    .collect(),
            };
            for (child_path, child_id) in children {
                let child = store
                    .get_object(&child_id)?
                    .ok_or(TreeError::ObjectNotFound(child_id))?;
                collect_leaf_paths(store, &child, &child_path, depth + 1, out)?;
            }
        }
    }
    Ok(())
}

/// Diff the string leaves of two trees, returning `(removed_paths, added
/// (path, value))` for the transition `old_root -> new_root`. Content-addressed:
/// a subtree whose ObjectId is unchanged is skipped whole, so the cost is
/// proportional to what a commit actually changed — not the tree size. This is
/// what makes incremental index maintenance on the write path cheap. `old_root
/// = None` treats every leaf of `new_root` as added. See plan t-006.
pub fn tree_diff_leaves(
    store: &dyn ObjectStore,
    old_root: Option<&ObjectId>,
    new_root: &ObjectId,
) -> Result<(Vec<String>, Vec<(String, String)>), TreeError> {
    let mut removed = Vec::new();
    let mut added = Vec::new();
    diff_recursive(
        store,
        old_root.copied(),
        Some(*new_root),
        "",
        0,
        &mut removed,
        &mut added,
    )?;
    Ok((removed, added))
}

#[allow(clippy::too_many_arguments)]
fn diff_recursive(
    store: &dyn ObjectStore,
    old_id: Option<ObjectId>,
    new_id: Option<ObjectId>,
    path: &str,
    depth: usize,
    removed: &mut Vec<String>,
    added: &mut Vec<(String, String)>,
) -> Result<(), TreeError> {
    // Content-addressing: identical ids (or both absent) => identical subtree.
    if old_id == new_id {
        return Ok(());
    }
    if depth > LIST_PATHS_HARD_CEILING {
        return Ok(());
    }
    let old_obj = match old_id {
        Some(id) => store.get_object(&id)?,
        None => None,
    };
    let new_obj = match new_id {
        Some(id) => store.get_object(&id)?,
        None => None,
    };
    match (old_obj, new_obj) {
        (None, None) => {}
        (Some(o), None) => collect_leaf_paths(store, &o, path, depth, removed)?,
        (None, Some(n)) => collect_leaves_recursive(store, &n, path, depth, added)?,
        (Some(Object::Atom(_)), Some(Object::Atom(new_atom))) => {
            // Same path, changed value: drop the old row, add the new one.
            let leaf = if path.is_empty() { "/" } else { path }.to_string();
            removed.push(leaf.clone());
            if let Some(v) = atom_value_string(&new_atom) {
                added.push((leaf, v));
            }
        }
        (Some(o @ Object::Node(_)), Some(Object::Atom(new_atom))) => {
            collect_leaf_paths(store, &o, path, depth, removed)?;
            let leaf = if path.is_empty() { "/" } else { path }.to_string();
            if let Some(v) = atom_value_string(&new_atom) {
                added.push((leaf, v));
            }
        }
        (Some(Object::Atom(_)), Some(n @ Object::Node(_))) => {
            let leaf = if path.is_empty() { "/" } else { path }.to_string();
            removed.push(leaf);
            collect_leaves_recursive(store, &n, path, depth, added)?;
        }
        (Some(Object::Node(on)), Some(Object::Node(nn))) => match (on, nn) {
            (Node::Map(oe), Node::Map(ne)) => {
                let mut keys: std::collections::BTreeSet<&String> = oe.keys().collect();
                keys.extend(ne.keys());
                for k in keys {
                    let child_path = format!("{}/{}", path, k);
                    diff_recursive(
                        store,
                        oe.get(k).copied(),
                        ne.get(k).copied(),
                        &child_path,
                        depth + 1,
                        removed,
                        added,
                    )?;
                }
            }
            (Node::List(oi) | Node::Set(oi), Node::List(ni) | Node::Set(ni)) => {
                let max = oi.len().max(ni.len());
                for i in 0..max {
                    let child_path = format!("{}/{}", path, i);
                    diff_recursive(
                        store,
                        oi.get(i).copied(),
                        ni.get(i).copied(),
                        &child_path,
                        depth + 1,
                        removed,
                        added,
                    )?;
                }
            }
            // Container kind changed (map <-> sequence): treat as full replace.
            (old_node, new_node) => {
                collect_leaf_paths(store, &Object::Node(old_node), path, depth, removed)?;
                collect_leaves_recursive(store, &Object::Node(new_node), path, depth, added)?;
            }
        },
    }
    Ok(())
}

/// Convert a serde_json::Value to Objects and store them.
/// Returns the root ObjectId.
pub fn json_to_tree(
    store: &dyn ObjectStore,
    value: &serde_json::Value,
) -> Result<ObjectId, TreeError> {
    let obj = json_to_object(store, value)?;
    let id = store.put_object(&obj)?;
    Ok(id)
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

fn navigate(
    current: &Object,
    component: &PathComponent,
    store: &dyn ObjectStore,
    path: &StatePath,
    _depth: usize,
) -> Result<Object, TreeError> {
    match (current, component) {
        (Object::Node(Node::Map(entries)), PathComponent::Key(key)) => {
            let child_id = entries
                .get(key)
                .ok_or_else(|| TreeError::PathNotFound(path.to_string()))?;
            store
                .get_object(child_id)?
                .ok_or(TreeError::ObjectNotFound(*child_id))
        }
        (Object::Node(Node::List(items)), PathComponent::Index(idx)) => {
            let child_id = items.get(*idx).ok_or(TreeError::IndexOutOfBounds {
                index: *idx,
                length: items.len(),
            })?;
            store
                .get_object(child_id)?
                .ok_or(TreeError::ObjectNotFound(*child_id))
        }
        (Object::Node(Node::Map(_)), PathComponent::Index(_)) => Err(TreeError::TypeMismatch {
            path: path.to_string(),
            expected: "list".to_string(),
            found: "map".to_string(),
        }),
        (Object::Node(Node::List(_)), PathComponent::Key(_)) => Err(TreeError::TypeMismatch {
            path: path.to_string(),
            expected: "map".to_string(),
            found: "list".to_string(),
        }),
        (Object::Atom(_), _) => Err(TreeError::CannotNavigateAtom(path.to_string())),
        _ => Err(TreeError::PathNotFound(path.to_string())),
    }
}

fn set_recursive(
    store: &dyn ObjectStore,
    current: &Object,
    components: &[PathComponent],
    depth: usize,
    value_id: ObjectId,
) -> Result<Object, TreeError> {
    if depth >= components.len() {
        // We've consumed all path components — return the object at value_id
        return store
            .get_object(&value_id)?
            .ok_or(TreeError::ObjectNotFound(value_id));
    }

    let component = &components[depth];
    let is_last = depth == components.len() - 1;

    match (current, component) {
        (Object::Node(Node::Map(entries)), PathComponent::Key(key)) => {
            let child_value_id = if is_last {
                value_id
            } else {
                // Navigate deeper
                let child = if let Some(child_id) = entries.get(key) {
                    store
                        .get_object(child_id)?
                        .ok_or(TreeError::ObjectNotFound(*child_id))?
                } else {
                    // Path doesn't exist yet — create intermediate maps
                    Object::empty_map()
                };
                let new_child = set_recursive(store, &child, components, depth + 1, value_id)?;
                store.put_object(&new_child)?
            };

            let mut new_entries = entries.clone();
            new_entries.insert(key.clone(), child_value_id);
            Ok(Object::map(new_entries))
        }
        (Object::Node(Node::List(items)), PathComponent::Index(idx)) => {
            if *idx >= items.len() {
                return Err(TreeError::IndexOutOfBounds {
                    index: *idx,
                    length: items.len(),
                });
            }

            let child_value_id = if is_last {
                value_id
            } else {
                let child = store
                    .get_object(&items[*idx])?
                    .ok_or_else(|| TreeError::ObjectNotFound(items[*idx]))?;
                let new_child = set_recursive(store, &child, components, depth + 1, value_id)?;
                store.put_object(&new_child)?
            };

            let mut new_items = items.clone();
            new_items[*idx] = child_value_id;
            Ok(Object::list(new_items))
        }
        (Object::Atom(Atom::Null), PathComponent::Key(key)) => {
            // A null is a placeholder: start the subtree there. Any other
            // value is data, and writing below it must not replace it — that
            // turned `/a = 30` plus a write to `/a/b` into `{"b": ...}` with
            // the 30 gone, and `Ok`.
            let mut entries = BTreeMap::new();
            if is_last {
                entries.insert(key.clone(), value_id);
            } else {
                let child = Object::empty_map();
                let new_child = set_recursive(store, &child, components, depth + 1, value_id)?;
                let child_id = store.put_object(&new_child)?;
                entries.insert(key.clone(), child_id);
            }
            Ok(Object::map(entries))
        }
        (Object::Atom(atom), _) => Err(TreeError::TypeMismatch {
            path: format!("at depth {}", depth),
            expected: "map or list".to_string(),
            found: format!("a {} value", atom_kind(atom)),
        }),
        _ => Err(TreeError::TypeMismatch {
            path: format!("at depth {}", depth),
            expected: "map or list".to_string(),
            found: "incompatible type".to_string(),
        }),
    }
}

fn atom_kind(atom: &Atom) -> &'static str {
    match atom {
        Atom::Null => "null",
        Atom::Bool(_) => "boolean",
        Atom::Int(_) | Atom::UInt(_) | Atom::Float(_) => "number",
        Atom::String(_) => "string",
        Atom::Bytes(_) => "bytes",
    }
}

fn delete_recursive(
    store: &dyn ObjectStore,
    current: &Object,
    components: &[PathComponent],
    depth: usize,
) -> Result<Object, TreeError> {
    let component = &components[depth];
    let is_last = depth == components.len() - 1;

    match (current, component) {
        (Object::Node(Node::Map(entries)), PathComponent::Key(key)) => {
            if !entries.contains_key(key) {
                return Err(TreeError::PathNotFound(format!("key '{}' not found", key)));
            }

            if is_last {
                let mut new_entries = entries.clone();
                new_entries.remove(key);
                Ok(Object::map(new_entries))
            } else {
                let child_id = entries.get(key).unwrap();
                let child = store
                    .get_object(child_id)?
                    .ok_or(TreeError::ObjectNotFound(*child_id))?;
                let new_child = delete_recursive(store, &child, components, depth + 1)?;
                let new_child_id = store.put_object(&new_child)?;
                let mut new_entries = entries.clone();
                new_entries.insert(key.clone(), new_child_id);
                Ok(Object::map(new_entries))
            }
        }
        (Object::Node(Node::List(items)), PathComponent::Index(idx)) => {
            if *idx >= items.len() {
                return Err(TreeError::IndexOutOfBounds {
                    index: *idx,
                    length: items.len(),
                });
            }

            if is_last {
                let mut new_items = items.clone();
                new_items.remove(*idx);
                Ok(Object::list(new_items))
            } else {
                let child = store
                    .get_object(&items[*idx])?
                    .ok_or_else(|| TreeError::ObjectNotFound(items[*idx]))?;
                let new_child = delete_recursive(store, &child, components, depth + 1)?;
                let new_child_id = store.put_object(&new_child)?;
                let mut new_items = items.clone();
                new_items[*idx] = new_child_id;
                Ok(Object::list(new_items))
            }
        }
        _ => Err(TreeError::PathNotFound(format!(
            "cannot navigate at depth {}",
            depth
        ))),
    }
}

fn atom_to_json(atom: &Atom) -> serde_json::Value {
    match atom {
        Atom::Null => serde_json::Value::Null,
        Atom::Bool(b) => serde_json::Value::Bool(*b),
        Atom::Int(i) => serde_json::json!(*i),
        Atom::UInt(u) => serde_json::json!(*u),
        Atom::Float(f) => serde_json::json!(*f),
        Atom::String(s) => serde_json::Value::String(s.clone()),
        Atom::Bytes(b) => {
            use serde_json::Value;
            // Encode bytes as base64 string
            Value::String(format!("base64:{}", base64_encode(b)))
        }
    }
}

fn json_to_object(store: &dyn ObjectStore, value: &serde_json::Value) -> Result<Object, TreeError> {
    match value {
        serde_json::Value::Null => Ok(Object::null()),
        serde_json::Value::Bool(b) => Ok(Object::bool(*b)),
        serde_json::Value::Number(n) => Ok(Object::from_json_number(n)),
        serde_json::Value::String(s) => Ok(Object::string(s.clone())),
        serde_json::Value::Array(arr) => {
            let mut ids = Vec::new();
            for item in arr {
                let child = json_to_object(store, item)?;
                let id = store.put_object(&child)?;
                ids.push(id);
            }
            Ok(Object::list(ids))
        }
        serde_json::Value::Object(map) => {
            let mut entries = BTreeMap::new();
            for (key, val) in map {
                let child = json_to_object(store, val)?;
                let id = store.put_object(&child)?;
                entries.insert(key.clone(), id);
            }
            Ok(Object::map(entries))
        }
    }
}

/// Simple base64 encoding without external dependency.
fn base64_encode(data: &[u8]) -> String {
    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut result = String::new();
    for chunk in data.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
        let triple = (b0 << 16) | (b1 << 8) | b2;
        result.push(CHARS[((triple >> 18) & 0x3F) as usize] as char);
        result.push(CHARS[((triple >> 12) & 0x3F) as usize] as char);
        if chunk.len() > 1 {
            result.push(CHARS[((triple >> 6) & 0x3F) as usize] as char);
        } else {
            result.push('=');
        }
        if chunk.len() > 2 {
            result.push(CHARS[(triple & 0x3F) as usize] as char);
        } else {
            result.push('=');
        }
    }
    result
}

// ---------------------------------------------------------------------------
// Errors
// ---------------------------------------------------------------------------

#[derive(Debug, thiserror::Error)]
pub enum TreeError {
    #[error("object not found: {0}")]
    ObjectNotFound(ObjectId),

    #[error("path not found: {0}")]
    PathNotFound(String),

    #[error("index {index} out of bounds (length {length})")]
    IndexOutOfBounds { index: usize, length: usize },

    #[error("type mismatch at {path}: expected {expected}, found {found}")]
    TypeMismatch {
        path: String,
        expected: String,
        found: String,
    },

    #[error("cannot navigate into atom at: {0}")]
    CannotNavigateAtom(String),

    #[error("invalid JSON: {0}")]
    InvalidJson(String),

    #[error("storage error: {0}")]
    Storage(#[from] StorageError),
}

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

    fn setup() -> (SqliteStorage, ObjectId) {
        let store = SqliteStorage::in_memory().expect("in-memory sqlite");

        // Build a test state tree:
        // {
        //   "name": "test-cluster",
        //   "nodes": [
        //     { "hostname": "jetson-01", "status": "healthy" },
        //     { "hostname": "jetson-02", "status": "unhealthy" }
        //   ],
        //   "config": { "network": { "subnet": "10.0.0.0/24" } }
        // }
        let root_id = json_to_tree(
            &store,
            &serde_json::json!({
                "name": "test-cluster",
                "nodes": [
                    { "hostname": "jetson-01", "status": "healthy" },
                    { "hostname": "jetson-02", "status": "unhealthy" }
                ],
                "config": {
                    "network": { "subnet": "10.0.0.0/24" }
                }
            }),
        )
        .unwrap();

        (store, root_id)
    }

    #[test]
    fn test_get_root() {
        let (store, root_id) = setup();
        let root = tree_get(&store, &root_id, &StatePath::root()).unwrap();
        assert!(matches!(root, Object::Node(Node::Map(_))));
    }

    #[test]
    fn test_get_simple_key() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/name").unwrap();
        let obj = tree_get(&store, &root_id, &path).unwrap();
        assert_eq!(obj, Object::string("test-cluster"));
    }

    #[test]
    fn test_get_nested_path() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/config/network/subnet").unwrap();
        let obj = tree_get(&store, &root_id, &path).unwrap();
        assert_eq!(obj, Object::string("10.0.0.0/24"));
    }

    #[test]
    fn test_get_array_index() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/nodes/0/hostname").unwrap();
        let obj = tree_get(&store, &root_id, &path).unwrap();
        assert_eq!(obj, Object::string("jetson-01"));
    }

    #[test]
    fn test_get_nonexistent_path() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/nonexistent").unwrap();
        assert!(tree_get(&store, &root_id, &path).is_err());
    }

    #[test]
    fn test_set_existing_value() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/nodes/1/status").unwrap();

        let new_root = tree_set(&store, &root_id, &path, &Object::string("healthy")).unwrap();

        // Verify the change
        let obj = tree_get(&store, &new_root, &path).unwrap();
        assert_eq!(obj, Object::string("healthy"));

        // Verify old root is unchanged (immutability)
        let old_obj = tree_get(&store, &root_id, &path).unwrap();
        assert_eq!(old_obj, Object::string("unhealthy"));
    }

    #[test]
    fn test_set_new_key() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/version").unwrap();

        let new_root = tree_set(&store, &root_id, &path, &Object::string("1.0")).unwrap();

        let obj = tree_get(&store, &new_root, &path).unwrap();
        assert_eq!(obj, Object::string("1.0"));

        // Original keys still accessible
        let name = tree_get(&store, &new_root, &StatePath::parse("/name").unwrap()).unwrap();
        assert_eq!(name, Object::string("test-cluster"));
    }

    #[test]
    fn test_set_creates_intermediate_maps() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/new/deeply/nested/value").unwrap();

        let new_root = tree_set(&store, &root_id, &path, &Object::int(42)).unwrap();

        let obj = tree_get(&store, &new_root, &path).unwrap();
        assert_eq!(obj, Object::int(42));
    }

    #[test]
    fn test_structural_sharing() {
        let (store, root_id) = setup();

        // Modify only /name — /config subtree should be shared
        let path = StatePath::parse("/name").unwrap();
        let new_root = tree_set(&store, &root_id, &path, &Object::string("new-name")).unwrap();

        // Get the config ObjectId from both roots — should be the same
        let old_root_obj = store.get_object(&root_id).unwrap().unwrap();
        let new_root_obj = store.get_object(&new_root).unwrap().unwrap();

        if let (Object::Node(Node::Map(old_entries)), Object::Node(Node::Map(new_entries))) =
            (&old_root_obj, &new_root_obj)
        {
            assert_eq!(
                old_entries.get("config"),
                new_entries.get("config"),
                "unmodified subtrees should share the same ObjectId (structural sharing)"
            );
            assert_ne!(
                old_entries.get("name"),
                new_entries.get("name"),
                "modified values should have different ObjectIds"
            );
        } else {
            panic!("roots should be maps");
        }
    }

    #[test]
    fn test_delete_key() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/name").unwrap();

        let new_root = tree_delete(&store, &root_id, &path).unwrap();

        assert!(tree_get(&store, &new_root, &path).is_err());

        // Other keys still accessible
        let subnet = tree_get(
            &store,
            &new_root,
            &StatePath::parse("/config/network/subnet").unwrap(),
        )
        .unwrap();
        assert_eq!(subnet, Object::string("10.0.0.0/24"));
    }

    #[test]
    fn test_delete_array_element() {
        let (store, root_id) = setup();
        let path = StatePath::parse("/nodes/0").unwrap();

        let new_root = tree_delete(&store, &root_id, &path).unwrap();

        // After deleting index 0, old index 1 is now at index 0
        let hostname = tree_get(
            &store,
            &new_root,
            &StatePath::parse("/nodes/0/hostname").unwrap(),
        )
        .unwrap();
        assert_eq!(hostname, Object::string("jetson-02"));
    }

    #[test]
    fn test_list_paths_truncates_at_max_results() {
        let store = SqliteStorage::in_memory().expect("in-memory sqlite");
        // Build a wide tree: a single map with > LIST_PATHS_MAX_RESULTS keys.
        let mut obj = serde_json::Map::new();
        let wide = LIST_PATHS_MAX_RESULTS + 500;
        for i in 0..wide {
            obj.insert(format!("k{}", i), serde_json::Value::from(i as i64));
        }
        let root_id = json_to_tree(&store, &serde_json::Value::Object(obj)).unwrap();

        let paths = tree_list_paths(&store, &root_id, "/", 10).unwrap();
        assert_eq!(
            paths.len(),
            LIST_PATHS_MAX_RESULTS,
            "list_paths must cap results at LIST_PATHS_MAX_RESULTS"
        );
    }

    #[test]
    fn test_list_paths_clamps_usize_max_depth() {
        // A usize::MAX max_depth used to drive unbounded recursion; now
        // it must clamp to LIST_PATHS_HARD_CEILING and return promptly.
        let (store, root_id) = setup();
        let paths = tree_list_paths(&store, &root_id, "/", usize::MAX).unwrap();
        // Just assert we got here without panicking / OOMing, and that we
        // returned a plausible, bounded result set.
        assert!(!paths.is_empty());
        assert!(paths.len() <= LIST_PATHS_MAX_RESULTS);
    }

    #[test]
    fn test_to_json_capped_truncates_nodes_at_depth() {
        let (store, root_id) = setup();
        let root = tree_get(&store, &root_id, &StatePath::root()).unwrap();
        let capped = tree_to_json_capped(&store, &root, 1).unwrap();

        // Atom child is returned as-is.
        assert_eq!(capped["name"], "test-cluster");
        // A List child at the cap becomes a length placeholder, not its items.
        assert_eq!(capped["nodes"]["_truncated"], true);
        assert_eq!(capped["nodes"]["_len"], 2);
        assert!(capped["nodes"].get("0").is_none());
        // A Map child at the cap becomes a keys placeholder.
        assert_eq!(capped["config"]["_truncated"], true);
        assert_eq!(capped["config"]["_keys"][0], "network");
        assert!(capped["config"]["network"].is_null());
    }

    #[test]
    fn test_to_json_capped_deep_enough_matches_full() {
        let (store, root_id) = setup();
        let root = tree_get(&store, &root_id, &StatePath::root()).unwrap();
        // A cap deeper than the tree yields the same result as the full walk.
        let full = tree_to_json(&store, &root).unwrap();
        let capped = tree_to_json_capped(&store, &root, 64).unwrap();
        assert_eq!(full, capped);
    }

    #[test]
    fn test_json_roundtrip() {
        let (store, root_id) = setup();
        let root = tree_get(&store, &root_id, &StatePath::root()).unwrap();
        let json = tree_to_json(&store, &root).unwrap();

        // Re-import
        let new_root_id = json_to_tree(&store, &json).unwrap();

        // Same content should produce same ObjectId
        assert_eq!(
            root_id, new_root_id,
            "JSON roundtrip should produce same ObjectId"
        );
    }
}