//! Unified state tree facade and state root.
//!
//! [`StateTree`] wraps a single Jellyfish Merkle Tree backed by redb. Three logical
//! namespaces (identity, confidential, secrets) share the JMT via a 1-byte prefix at
//! `leaf_key[0]` — there is one version counter, one state root, and one proof cache.

use crate::{
    cache::ProofCache,
    error::{Result, StateTreeError},
    key::NamespaceKey,
    storage::{self, RedbNodeStore, TransactionReader},
};
use borsh::{BorshDeserialize, BorshSerialize};
use jmt::{JellyfishMerkleTree, KeyHash, RootHash};
use std::{path::Path, sync::Arc};

/// The 32-byte state root hash of the unified JMT.
///
/// Cross-boundary consumers (BLS signing, on-chain state commits) should use
/// [`AsRef<[u8]>`] to obtain the canonical byte order.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)]
pub struct StateRoot(pub [u8; 32]);

/// JMT's empty-tree root: `*b"SPARSE_MERKLE_PLACEHOLDER_HASH__"`.
const EMPTY_ROOT_HASH: [u8; 32] = *b"SPARSE_MERKLE_PLACEHOLDER_HASH__";

impl StateRoot {
    /// JMT placeholder root for an empty tree. Use this for non-inclusion
    /// proofs against trees with no committed version, instead of re-reading
    /// `tree.root()` on a separate transaction.
    pub const fn placeholder() -> Self {
        Self(EMPTY_ROOT_HASH)
    }
}

impl AsRef<[u8]> for StateRoot {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl From<RootHash> for StateRoot {
    fn from(rh: RootHash) -> Self {
        Self(rh.0)
    }
}

impl From<StateRoot> for RootHash {
    fn from(sr: StateRoot) -> Self {
        RootHash(sr.0)
    }
}

/// Configuration for [`StateTree`] resource limits.
#[derive(Debug, Clone)]
pub struct StateTreeConfig {
    /// Maximum allowed value size in bytes (default: 8 MiB).
    pub max_value_bytes: usize,
    /// Maximum entries per `commit_batch` call (default: 10 000).
    pub max_batch_entries: usize,
    /// Maximum entries in the proof cache (default: 10 000).
    pub proof_cache_capacity: u64,
}

impl Default for StateTreeConfig {
    fn default() -> Self {
        Self {
            max_value_bytes: 8 * 1024 * 1024,
            max_batch_entries: 10_000,
            proof_cache_capacity: 10_000,
        }
    }
}

/// Single unified Jellyfish Merkle Tree backed by redb.
///
/// Three logical namespaces (identity, confidential, secrets) share one JMT via
/// a namespace-byte prefix at `leaf_key[0]`. One version counter, one state root.
pub struct StateTree {
    store: Arc<RedbNodeStore>,
    cache: ProofCache,
    config: StateTreeConfig,
}

impl std::fmt::Debug for StateTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StateTree").finish_non_exhaustive()
    }
}

impl StateTree {
    /// Open or create the state tree backed by the redb file at `path`.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        Self::open_with_config(path, StateTreeConfig::default())
    }

    /// Open with explicit resource-limit configuration.
    pub fn open_with_config(path: impl AsRef<Path>, config: StateTreeConfig) -> Result<Self> {
        let store = Arc::new(RedbNodeStore::open(path)?);
        let cache = ProofCache::new(config.proof_cache_capacity);
        Ok(Self { store, cache, config })
    }

    /// Insert or update a leaf value. Returns the new unified [`StateRoot`].
    #[tracing::instrument(skip(self, value), fields(namespace = ?key.namespace(), value_bytes = value.len()))]
    pub fn put(&self, key: &NamespaceKey, value: Vec<u8>) -> Result<StateRoot> {
        let started = std::time::Instant::now();
        let value_bytes = value.len() as u64;

        if value.len() > self.config.max_value_bytes {
            crate::metrics::commit_failure(crate::metrics::REASON_VALUE_TOO_LARGE);
            return Err(StateTreeError::ValueTooLarge {
                size: value.len(),
                cap: self.config.max_value_bytes,
            });
        }
        let key_hash = key.key_hash();
        let write_txn = self.store.begin_write()?;
        let read_txn = self.store.begin_read()?;
        let reader = TransactionReader::new(read_txn)?;

        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let version = storage::next_version(&write_txn)?;
        let (root_hash, batch) = tree
            .put_value_set(vec![(key_hash, Some(value.clone()))], version)
            .map_err(|source| {
                crate::metrics::commit_failure(crate::metrics::REASON_JMT);
                StateTreeError::Jmt { version, source }
            })?;

        drop(reader);
        storage::write_nodes(&write_txn, &batch)?;
        storage::write_value(&write_txn, version, key_hash, Some(value))?;
        write_txn.commit()?;

        self.cache.invalidate();
        crate::metrics::commit_success(crate::metrics::OP_PUT, started.elapsed().as_secs_f64(), 1, value_bytes);
        Ok(root_hash.into())
    }

    /// Get the raw leaf value for a key. Returns `None` if absent or tombstoned.
    #[tracing::instrument(level = "debug", skip(self), fields(namespace = ?key.namespace()))]
    pub fn get(&self, key: &NamespaceKey) -> Result<Option<Vec<u8>>> {
        let key_hash = key.key_hash();
        self.store.read(|txn| {
            let version = match storage::current_version(txn)? {
                Some(v) => v,
                None => return Ok(None),
            };
            storage::read_value(txn, version, key_hash)
        })
    }

    /// Delete (tombstone) a leaf. Returns the new unified [`StateRoot`].
    #[tracing::instrument(skip(self), fields(namespace = ?key.namespace()))]
    pub fn delete(&self, key: &NamespaceKey) -> Result<StateRoot> {
        let started = std::time::Instant::now();
        let key_hash = key.key_hash();
        let write_txn = self.store.begin_write()?;
        let read_txn = self.store.begin_read()?;
        let reader = TransactionReader::new(read_txn)?;

        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let version = storage::next_version(&write_txn)?;
        let (root_hash, batch) = tree.put_value_set(vec![(key_hash, None)], version).map_err(|source| {
            crate::metrics::commit_failure(crate::metrics::REASON_JMT);
            StateTreeError::Jmt { version, source }
        })?;

        drop(reader);
        storage::write_nodes(&write_txn, &batch)?;
        storage::write_value(&write_txn, version, key_hash, None)?;
        write_txn.commit()?;

        self.cache.invalidate();
        crate::metrics::commit_success(crate::metrics::OP_DELETE, started.elapsed().as_secs_f64(), 1, 0);
        Ok(root_hash.into())
    }

    /// The current unified state root. Returns the JMT placeholder hash if no
    /// writes have occurred.
    #[tracing::instrument(level = "debug", skip(self))]
    pub fn root(&self) -> Result<StateRoot> {
        let read_txn = self.store.begin_read()?;
        let version = match storage::current_version(&read_txn)? {
            Some(v) => v,
            None => return Ok(StateRoot(EMPTY_ROOT_HASH)),
        };

        let reader = TransactionReader::new(read_txn)?;
        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let root = tree
            .get_root_hash(version)
            .map_err(|source| StateTreeError::Jmt { version, source })?;
        Ok(root.into())
    }

    /// The current committed JMT version, or `None` if the tree has never been
    /// written to.
    ///
    /// Increments per commit (`put`, `delete`, or non-empty `commit_batch`).
    /// An empty `commit_batch` is a no-op and does not bump the version.
    /// Used by callers that need to assert local replica progress against the
    /// on-chain `StateCommitRegistry.currentSequenceNo` without paying for a
    /// full proof or root materialization.
    #[tracing::instrument(level = "debug", skip(self))]
    pub fn version(&self) -> Result<Option<u64>> {
        self.store.read(storage::current_version)
    }

    /// Atomically insert, update, and/or delete multiple leaves across any namespaces.
    ///
    /// Returns the new unified [`StateRoot`]. No-ops on an empty batch, returning the
    /// current root without bumping the version.
    #[tracing::instrument(skip(self, entries), fields(batch_size = entries.len()))]
    pub fn commit_batch(&self, entries: Vec<(NamespaceKey, Option<Vec<u8>>)>) -> Result<StateRoot> {
        let started = std::time::Instant::now();

        if entries.is_empty() {
            return self.root();
        }

        let batch_size = entries.len();

        if batch_size > self.config.max_batch_entries {
            crate::metrics::commit_failure(crate::metrics::REASON_BATCH_TOO_LARGE);
            return Err(StateTreeError::BatchTooLarge {
                count: batch_size,
                cap: self.config.max_batch_entries,
            });
        }

        let mut total_value_bytes: u64 = 0;
        for (_, value) in &entries {
            if let Some(v) = value {
                if v.len() > self.config.max_value_bytes {
                    crate::metrics::commit_failure(crate::metrics::REASON_VALUE_TOO_LARGE);
                    return Err(StateTreeError::ValueTooLarge {
                        size: v.len(),
                        cap: self.config.max_value_bytes,
                    });
                }
                total_value_bytes += v.len() as u64;
            }
        }

        let write_txn = self.store.begin_write()?;
        let read_txn = self.store.begin_read()?;
        let reader = TransactionReader::new(read_txn)?;

        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let version = storage::next_version(&write_txn)?;

        let (key_hashes, values): (Vec<KeyHash>, Vec<Option<Vec<u8>>>) =
            entries.into_iter().map(|(k, v)| (k.key_hash(), v)).unzip();

        // Clone is unavoidable: JMT consumes ownership and we still need values for write_value.
        let jmt_set = key_hashes.iter().zip(values.iter()).map(|(kh, v)| (*kh, v.clone()));
        let (root_hash, batch) = tree.put_value_set(jmt_set, version).map_err(|source| {
            crate::metrics::commit_failure(crate::metrics::REASON_JMT);
            StateTreeError::Jmt { version, source }
        })?;

        drop(reader);
        storage::write_nodes(&write_txn, &batch)?;
        for (key_hash, value) in key_hashes.into_iter().zip(values) {
            storage::write_value(&write_txn, version, key_hash, value)?;
        }
        write_txn.commit()?;

        self.cache.invalidate();

        crate::metrics::commit_success(
            crate::metrics::OP_BATCH,
            started.elapsed().as_secs_f64(),
            batch_size,
            total_value_bytes,
        );

        Ok(root_hash.into())
    }

    /// Access the proof cache.
    pub fn proof_cache(&self) -> &ProofCache {
        &self.cache
    }

    /// Generate a Merkle inclusion or exclusion proof for a key.
    ///
    /// "Key not present" is encoded two ways and callers must handle both:
    /// - `Ok(None)` — the tree has never been written to (no committed JMT
    ///   version exists). There is no `LeafProof` to verify; callers that need
    ///   to assert non-membership against `EMPTY_ROOT_HASH` must short-circuit
    ///   on this branch.
    /// - `Ok(Some(LeafProof { value: None, .. }))` — the tree is populated and
    ///   the key is genuinely absent. The returned proof is a non-membership
    ///   proof that can be verified via [`LeafProof::verify`].
    ///
    /// For populated trees with an existing key, returns
    /// `Ok(Some(LeafProof { value: Some(_), .. }))`.
    #[tracing::instrument(level = "debug", skip(self), fields(namespace = ?key.namespace()))]
    pub fn proof(&self, key: &NamespaceKey) -> Result<Option<LeafProof>> {
        let kh = key.key_hash();

        if let Some(cached) = self.cache.get(&kh.0) {
            return Ok(Some(cached));
        }

        let cache_version = self.cache.version();

        let read_txn = self.store.begin_read()?;
        let version = match storage::current_version(&read_txn)? {
            Some(v) => v,
            None => return Ok(None),
        };

        let reader = TransactionReader::new(read_txn)?;
        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let jmt_err = |source| StateTreeError::Jmt { version, source };

        let (value, proof) = tree.get_with_proof(kh, version).map_err(jmt_err)?;
        let root = tree.get_root_hash(version).map_err(jmt_err)?;

        let leaf = LeafProof {
            value,
            proof,
            root: root.into(),
            version,
        };
        if !self.cache.try_insert(kh.0, leaf.clone(), cache_version) {
            tracing::debug!(
                key_hash = ?kh.0,
                expected_version = cache_version,
                current_version = self.cache.version(),
                "proof cache insert skipped: version advanced during proof generation"
            );
            crate::metrics::cache_outcome(crate::metrics::OUTCOME_MISS_VERSION_CHANGED);
        }
        Ok(Some(leaf))
    }
}

/// A Merkle proof for a single leaf in the unified JMT.
///
/// Covers both inclusion (`value` is `Some`) and exclusion (`value` is `None`).
/// Verification requires only the expected root hash — no tree handle needed.
///
/// Wire format uses borsh for deterministic, compact encoding that is independent
/// of jmt's serde representation.
#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)]
pub struct LeafProof {
    value: Option<Vec<u8>>,
    proof: jmt::proof::SparseMerkleProof<sha2::Sha256>,
    root: StateRoot,
    version: u64,
}

impl LeafProof {
    /// The leaf value, or `None` for exclusion (non-membership) proofs.
    pub fn value(&self) -> Option<&[u8]> {
        self.value.as_deref()
    }

    /// The JMT sparse merkle proof.
    pub fn proof(&self) -> &jmt::proof::SparseMerkleProof<sha2::Sha256> {
        &self.proof
    }

    /// The state root this proof was generated against.
    pub fn root(&self) -> &StateRoot {
        &self.root
    }

    /// The JMT version (monotonic sequence number).
    pub fn version(&self) -> u64 {
        self.version
    }

    /// Serialize to the canonical wire format (borsh).
    ///
    /// Returns `Result` rather than `Vec<u8>` because borsh's `to_vec` can fail
    /// on size-class overflow when `value: Option<Vec<u8>>` exceeds the encoder's
    /// internal limits. Per `rust.md`, `.expect()` on runtime data is banned in
    /// production paths even when the payload looks structurally well-formed.
    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        borsh::to_vec(self).map_err(|e| StateTreeError::Serialization(e.to_string()))
    }

    /// Deserialize from the canonical wire format (borsh).
    ///
    /// **Unbounded.** borsh decodes `Option<Vec<u8>>` and `SparseMerkleProof`
    /// using `u32` length prefixes — a malicious payload can declare ~4 GB and
    /// cause a pre-allocation before the actual byte count is reached. Use
    /// [`from_bytes_with_limit`](Self::from_bytes_with_limit) at every system
    /// boundary that ingests proofs from an untrusted source (P2P delta
    /// ingestion, bootnode replay, RPC inputs).
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        borsh::from_slice(bytes).map_err(|e| StateTreeError::Serialization(e.to_string()))
    }

    /// Deserialize from the canonical wire format (borsh) with an explicit
    /// upper bound on the encoded length.
    ///
    /// Returns `Err(StateTreeError::Serialization)` if `bytes.len() > max_bytes`
    /// before any allocation, so a hostile encoded length cannot trigger a
    /// pre-allocation DoS. Use this at every untrusted-input boundary; reserve
    /// [`from_bytes`](Self::from_bytes) for trusted in-process round-trips.
    pub fn from_bytes_with_limit(bytes: &[u8], max_bytes: usize) -> Result<Self> {
        if bytes.len() > max_bytes {
            return Err(StateTreeError::Serialization(format!(
                "proof exceeds size limit: {} bytes > {} max",
                bytes.len(),
                max_bytes
            )));
        }
        borsh::from_slice(bytes).map_err(|e| StateTreeError::Serialization(e.to_string()))
    }

    /// Verify this proof against an expected state root for the given key.
    pub fn verify(&self, key: &NamespaceKey, expected_root: &StateRoot) -> Result<()> {
        if self.root != *expected_root {
            return Err(StateTreeError::ProofInvalid {
                reason: format!(
                    "root mismatch: proof root {:?} != expected {:?}",
                    self.root.0, expected_root.0
                ),
            });
        }
        self.proof
            .verify(RootHash::from(*expected_root), key.key_hash(), self.value.as_deref())
            .map_err(|e| StateTreeError::ProofInvalid { reason: e.to_string() })
    }
}

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

    fn test_tree() -> StateTree {
        let dir = tempfile::tempdir().unwrap();
        StateTree::open(dir.path().join("test.redb")).unwrap()
    }

    #[test]
    fn empty_root_is_jmt_placeholder() {
        let tree = test_tree();
        assert_eq!(tree.root().unwrap(), StateRoot(EMPTY_ROOT_HASH));
    }

    #[test]
    fn version_is_none_for_empty_tree() {
        let tree = test_tree();
        assert_eq!(tree.version().unwrap(), None);
    }

    #[test]
    fn version_increments_per_commit() {
        let tree = test_tree();
        let k1 = NamespaceKey::identity([0xab; 32]);
        let k2 = NamespaceKey::identity([0xcd; 32]);

        tree.put(&k1, vec![1]).unwrap();
        let v1 = tree.version().unwrap().expect("version after first put");

        tree.put(&k2, vec![2]).unwrap();
        let v2 = tree.version().unwrap().expect("version after second put");

        assert!(v2 > v1, "version must monotonically increase: {v1} -> {v2}");

        tree.delete(&k1).unwrap();
        let v3 = tree.version().unwrap().expect("version after delete");
        assert!(v3 > v2, "delete must bump version: {v2} -> {v3}");
    }

    #[test]
    fn empty_commit_batch_does_not_bump_version() {
        let tree = test_tree();
        let k = NamespaceKey::identity([0xab; 32]);
        tree.put(&k, vec![1]).unwrap();
        let v_before = tree.version().unwrap().expect("version after put");

        tree.commit_batch(vec![]).unwrap();
        let v_after = tree.version().unwrap().expect("version after empty batch");

        assert_eq!(v_before, v_after, "empty commit_batch is a no-op");
    }

    #[test]
    fn put_changes_root() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root = tree.put(&key, vec![1, 2, 3]).unwrap();

        assert_ne!(root, StateRoot(EMPTY_ROOT_HASH));
        assert_eq!(tree.root().unwrap(), root);
    }

    #[test]
    fn get_returns_inserted_value() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        tree.put(&key, vec![1, 2, 3]).unwrap();

        let value = tree.get(&key).unwrap();
        assert_eq!(value, Some(vec![1, 2, 3]));
    }

    #[test]
    fn get_missing_returns_none() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        assert_eq!(tree.get(&key).unwrap(), None);
    }

    #[test]
    fn delete_changes_root() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root_after_put = tree.put(&key, vec![1, 2, 3]).unwrap();
        let root_after_delete = tree.delete(&key).unwrap();
        assert_ne!(root_after_put, root_after_delete);
    }

    #[test]
    fn delete_removes_value() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        tree.put(&key, vec![1, 2, 3]).unwrap();
        tree.delete(&key).unwrap();
        assert_eq!(tree.get(&key).unwrap(), None);
    }

    #[test]
    fn root_is_deterministic() {
        let tree1 = test_tree();
        let tree2 = test_tree();

        let key = NamespaceKey::identity([0xab; 32]);
        let value = vec![1, 2, 3];

        let root1 = tree1.put(&key, value.clone()).unwrap();
        let root2 = tree2.put(&key, value).unwrap();

        assert_eq!(root1, root2);
    }

    #[test]
    fn commit_batch_inserts_multiple() {
        let tree = test_tree();
        let k1 = NamespaceKey::identity([0x01; 32]);
        let k2 = NamespaceKey::identity([0x02; 32]);
        let k3 = NamespaceKey::identity([0x03; 32]);

        let root = tree
            .commit_batch(vec![
                (k1.clone(), Some(vec![10])),
                (k2.clone(), Some(vec![20])),
                (k3.clone(), Some(vec![30])),
            ])
            .unwrap();

        assert_ne!(root, StateRoot(EMPTY_ROOT_HASH));
        assert_eq!(tree.get(&k1).unwrap(), Some(vec![10]));
        assert_eq!(tree.get(&k2).unwrap(), Some(vec![20]));
        assert_eq!(tree.get(&k3).unwrap(), Some(vec![30]));
        assert_eq!(tree.root().unwrap(), root);
    }

    #[test]
    fn commit_batch_mixed_insert_and_delete() {
        let tree = test_tree();
        let k1 = NamespaceKey::identity([0x01; 32]);
        let k2 = NamespaceKey::identity([0x02; 32]);

        tree.put(&k1, vec![1]).unwrap();
        tree.put(&k2, vec![2]).unwrap();

        let root = tree
            .commit_batch(vec![(k1.clone(), None), (k2.clone(), Some(vec![22]))])
            .unwrap();

        assert_eq!(tree.get(&k1).unwrap(), None);
        assert_eq!(tree.get(&k2).unwrap(), Some(vec![22]));
        assert_eq!(tree.root().unwrap(), root);
    }

    #[test]
    fn commit_batch_empty_returns_current_root() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root_after_put = tree.put(&key, vec![1]).unwrap();

        let cache_version_before = tree.proof_cache().version();
        let root = tree.commit_batch(vec![]).unwrap();

        assert_eq!(root, root_after_put);
        assert_eq!(tree.proof_cache().version(), cache_version_before);
    }

    #[test]
    fn commit_batch_is_atomic_single_version() {
        let tree1 = test_tree();
        let tree2 = test_tree();

        let k1 = NamespaceKey::identity([0x01; 32]);
        let k2 = NamespaceKey::identity([0x02; 32]);

        let root_batch = tree1
            .commit_batch(vec![(k1.clone(), Some(vec![1])), (k2.clone(), Some(vec![2]))])
            .unwrap();

        tree2.put(&k1, vec![1]).unwrap();
        let root_sequential = tree2.put(&k2, vec![2]).unwrap();

        assert_eq!(root_batch, root_sequential);
        assert_eq!(tree1.get(&k1).unwrap(), tree2.get(&k1).unwrap());
        assert_eq!(tree1.get(&k2).unwrap(), tree2.get(&k2).unwrap());
    }

    #[test]
    fn commit_batch_invalidates_cache() {
        let tree = test_tree();
        let v_before = tree.proof_cache().version();
        tree.commit_batch(vec![(NamespaceKey::identity([0x01; 32]), Some(vec![1]))])
            .unwrap();
        assert_eq!(tree.proof_cache().version(), v_before + 1);
    }

    #[test]
    fn data_persists_across_reopen() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("persist.redb");

        let key = NamespaceKey::identity([0xaa; 32]);
        let root_before = {
            let tree = StateTree::open(&path).unwrap();
            tree.put(&key, vec![42]).unwrap()
        };

        let tree = StateTree::open(&path).unwrap();
        assert_eq!(tree.get(&key).unwrap(), Some(vec![42]));
        assert_eq!(tree.root().unwrap(), root_before);
    }

    #[test]
    fn put_same_key_overwrites() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);

        tree.put(&key, vec![1]).unwrap();
        let root = tree.put(&key, vec![2]).unwrap();

        assert_eq!(tree.get(&key).unwrap(), Some(vec![2]));
        assert_eq!(tree.root().unwrap(), root);
    }

    #[test]
    fn namespaces_share_unified_root() {
        let tree = test_tree();
        let raw = [0xab; 32];

        let root1 = tree.put(&NamespaceKey::identity(raw), vec![1]).unwrap();
        let root2 = tree.put(&NamespaceKey::confidential(raw), vec![2]).unwrap();
        let root3 = tree.put(&NamespaceKey::secrets(raw), vec![3]).unwrap();

        assert_eq!(tree.get(&NamespaceKey::identity(raw)).unwrap(), Some(vec![1]));
        assert_eq!(tree.get(&NamespaceKey::confidential(raw)).unwrap(), Some(vec![2]));
        assert_eq!(tree.get(&NamespaceKey::secrets(raw)).unwrap(), Some(vec![3]));

        assert_ne!(root1, root2);
        assert_ne!(root2, root3);
        assert_eq!(tree.root().unwrap(), root3);
    }

    #[test]
    fn cross_namespace_batch() {
        let tree = test_tree();
        let root = tree
            .commit_batch(vec![
                (NamespaceKey::identity([0x01; 32]), Some(vec![1])),
                (NamespaceKey::confidential([0x02; 32]), Some(vec![2])),
                (NamespaceKey::secrets([0x03; 32]), Some(vec![3])),
            ])
            .unwrap();

        assert_eq!(tree.get(&NamespaceKey::identity([0x01; 32])).unwrap(), Some(vec![1]));
        assert_eq!(
            tree.get(&NamespaceKey::confidential([0x02; 32])).unwrap(),
            Some(vec![2])
        );
        assert_eq!(tree.get(&NamespaceKey::secrets([0x03; 32])).unwrap(), Some(vec![3]));
        assert_eq!(tree.root().unwrap(), root);
    }

    #[test]
    fn cache_invalidation_on_write() {
        let tree = test_tree();

        let v = tree.proof_cache().version();

        tree.put(&NamespaceKey::identity([0x01; 32]), vec![1]).unwrap();
        assert_eq!(tree.proof_cache().version(), v + 1);

        tree.put(&NamespaceKey::confidential([0x01; 32]), vec![2]).unwrap();
        assert_eq!(tree.proof_cache().version(), v + 2);
    }

    // --- Proof tests ---

    #[test]
    fn proof_inclusion_verifies() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root = tree.put(&key, vec![1, 2, 3]).unwrap();

        let leaf = tree.proof(&key).unwrap().expect("non-empty tree");
        assert_eq!(leaf.value(), Some([1, 2, 3].as_slice()));
        assert_eq!(*leaf.root(), root);
        leaf.verify(&key, &root).unwrap();
    }

    #[test]
    fn proof_exclusion_verifies() {
        let tree = test_tree();
        let existing = NamespaceKey::identity([0x01; 32]);
        tree.put(&existing, vec![1]).unwrap();

        let absent = NamespaceKey::identity([0x02; 32]);
        let leaf = tree.proof(&absent).unwrap().expect("non-empty tree");
        assert_eq!(leaf.value(), None);

        let root = tree.root().unwrap();
        leaf.verify(&absent, &root).unwrap();
    }

    #[test]
    fn proof_against_stale_root_fails() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let old_root = tree.put(&key, vec![1]).unwrap();
        let leaf = tree.proof(&key).unwrap().unwrap();

        let new_root = tree.put(&key, vec![2]).unwrap();
        assert_ne!(old_root, new_root);

        let err = leaf.verify(&key, &new_root).unwrap_err();
        assert!(
            matches!(err, StateTreeError::ProofInvalid { .. }),
            "expected ProofInvalid, got: {err}"
        );
    }

    #[test]
    fn proof_deterministic() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        tree.put(&key, vec![1, 2, 3]).unwrap();

        let p1 = tree.proof(&key).unwrap().unwrap();
        let p2 = tree.proof(&key).unwrap().unwrap();

        let bytes1 = borsh::to_vec(p1.proof()).unwrap();
        let bytes2 = borsh::to_vec(p2.proof()).unwrap();
        assert_eq!(bytes1, bytes2);
        assert_eq!(p1.value(), p2.value());
        assert_eq!(p1.root(), p2.root());
        assert_eq!(p1.version(), p2.version());
    }

    #[test]
    fn proof_empty_tree_returns_none() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        assert!(tree.proof(&key).unwrap().is_none());
    }

    #[test]
    fn proof_verify_wrong_key_fails() {
        let tree = test_tree();
        let key_a = NamespaceKey::identity([0x01; 32]);
        let key_b = NamespaceKey::identity([0x02; 32]);
        let root = tree.put(&key_a, vec![1]).unwrap();

        let leaf = tree.proof(&key_a).unwrap().unwrap();
        let err = leaf.verify(&key_b, &root).unwrap_err();
        assert!(
            matches!(err, StateTreeError::ProofInvalid { .. }),
            "expected ProofInvalid, got: {err}"
        );
    }

    #[test]
    fn proof_cross_namespace_independent() {
        let tree = test_tree();
        let raw = [0xab; 32];
        tree.put(&NamespaceKey::identity(raw), vec![1]).unwrap();
        tree.put(&NamespaceKey::confidential(raw), vec![2]).unwrap();

        let id_key = NamespaceKey::identity(raw);
        let id_proof = tree.proof(&id_key).unwrap().unwrap();
        let root = tree.root().unwrap();

        id_proof.verify(&id_key, &root).unwrap();

        let conf_key = NamespaceKey::confidential(raw);
        let err = id_proof.verify(&conf_key, &root).unwrap_err();
        assert!(
            matches!(err, StateTreeError::ProofInvalid { .. }),
            "expected ProofInvalid, got: {err}"
        );
    }

    #[test]
    fn proof_cache_hit_returns_same_proof() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        tree.put(&key, vec![1, 2, 3]).unwrap();

        let p1 = tree.proof(&key).unwrap().unwrap();
        let p2 = tree.proof(&key).unwrap().unwrap();

        assert_eq!(p1.value(), p2.value());
        assert_eq!(p1.root(), p2.root());
        assert_eq!(p1.version(), p2.version());

        let bytes1 = borsh::to_vec(p1.proof()).unwrap();
        let bytes2 = borsh::to_vec(p2.proof()).unwrap();
        assert_eq!(bytes1, bytes2);
    }

    #[test]
    fn proof_cache_invalidated_on_write() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        tree.put(&key, vec![1]).unwrap();

        let p1 = tree.proof(&key).unwrap().unwrap();
        assert_eq!(p1.value(), Some([1].as_slice()));

        tree.put(&key, vec![2]).unwrap();

        let p2 = tree.proof(&key).unwrap().unwrap();
        assert_eq!(p2.value(), Some([2].as_slice()));
        assert_ne!(p1.root(), p2.root());
    }

    #[test]
    fn proof_tampered_value_fails_verification() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root = tree.put(&key, vec![1, 2, 3]).unwrap();

        let leaf = tree.proof(&key).unwrap().unwrap();
        assert_eq!(leaf.value(), Some([1, 2, 3].as_slice()));

        let tampered = LeafProof {
            value: Some(vec![99, 99, 99]),
            proof: leaf.proof.clone(),
            root: leaf.root,
            version: leaf.version,
        };
        let err = tampered.verify(&key, &root).unwrap_err();
        assert!(
            matches!(err, StateTreeError::ProofInvalid { .. }),
            "expected ProofInvalid, got: {err}"
        );
    }

    #[test]
    fn proof_single_leaf_tree() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0x01; 32]);
        let root = tree.put(&key, vec![42]).unwrap();

        let leaf = tree.proof(&key).unwrap().unwrap();
        assert_eq!(leaf.value(), Some([42].as_slice()));
        leaf.verify(&key, &root).unwrap();

        let absent = NamespaceKey::identity([0x02; 32]);
        let exclusion = tree.proof(&absent).unwrap().unwrap();
        assert_eq!(exclusion.value(), None);
        exclusion.verify(&absent, &root).unwrap();
    }

    #[test]
    fn proof_cache_respects_capacity() {
        let dir = tempfile::tempdir().unwrap();
        let config = StateTreeConfig {
            proof_cache_capacity: 4,
            ..Default::default()
        };
        let tree = StateTree::open_with_config(dir.path().join("cap.redb"), config).unwrap();

        for i in 0u8..20 {
            let key = NamespaceKey::identity([i; 32]);
            tree.put(&key, vec![i]).unwrap();
        }
        for i in 0u8..20 {
            let key = NamespaceKey::identity([i; 32]);
            tree.proof(&key).unwrap();
        }

        // moka may briefly exceed max_capacity before async eviction runs
        assert!(
            tree.proof_cache().entry_count() <= 8,
            "cache entry count {} should be bounded near max_capacity of 4",
            tree.proof_cache().entry_count()
        );
    }

    // --- Rightmost leaf metadata tests (NEWT-1066) ---

    #[test]
    fn rightmost_leaf_tracks_across_puts() {
        use jmt::storage::TreeReader;

        let dir = tempfile::tempdir().unwrap();
        let tree = StateTree::open(dir.path().join("rl.redb")).unwrap();

        let read_and_get_rightmost = |store: &RedbNodeStore| -> Option<KeyHash> {
            let txn = store.begin_read().unwrap();
            let reader = TransactionReader::new(txn).unwrap();
            reader.get_rightmost_leaf().unwrap().map(|(_, leaf)| leaf.key_hash())
        };

        assert!(read_and_get_rightmost(&tree.store).is_none());

        let k1 = NamespaceKey::identity([0x10; 32]);
        tree.put(&k1, vec![1]).unwrap();
        let rl1 = read_and_get_rightmost(&tree.store).unwrap();
        assert_eq!(rl1, k1.key_hash());

        let k2 = NamespaceKey::identity([0x20; 32]);
        tree.put(&k2, vec![2]).unwrap();
        let rl2 = read_and_get_rightmost(&tree.store).unwrap();

        let expected = std::cmp::max(k1.key_hash().0, k2.key_hash().0);
        assert_eq!(rl2.0, expected);
    }

    #[test]
    fn rightmost_leaf_tracks_across_batch() {
        use jmt::storage::TreeReader;

        let dir = tempfile::tempdir().unwrap();
        let tree = StateTree::open(dir.path().join("rl_batch.redb")).unwrap();

        let keys: Vec<NamespaceKey> = (0u8..10).map(|i| NamespaceKey::identity([i; 32])).collect();
        let entries: Vec<_> = keys.iter().map(|k| (k.clone(), Some(vec![42]))).collect();
        tree.commit_batch(entries).unwrap();

        let expected_rightmost = keys.iter().map(|k| k.key_hash()).max_by_key(|kh| kh.0).unwrap();

        let txn = tree.store.begin_read().unwrap();
        let reader = TransactionReader::new(txn).unwrap();
        let (_, leaf) = reader.get_rightmost_leaf().unwrap().expect("should have leaves");
        assert_eq!(leaf.key_hash(), expected_rightmost);
    }

    #[test]
    fn rightmost_leaf_tracks_across_namespaces() {
        use jmt::storage::TreeReader;

        let dir = tempfile::tempdir().unwrap();
        let tree = StateTree::open(dir.path().join("rl_ns.redb")).unwrap();

        tree.commit_batch(vec![
            (NamespaceKey::identity([0x01; 32]), Some(vec![1])),
            (NamespaceKey::confidential([0x02; 32]), Some(vec![2])),
            (NamespaceKey::secrets([0x03; 32]), Some(vec![3])),
        ])
        .unwrap();

        let all_hashes = [
            NamespaceKey::identity([0x01; 32]).key_hash(),
            NamespaceKey::confidential([0x02; 32]).key_hash(),
            NamespaceKey::secrets([0x03; 32]).key_hash(),
        ];
        let expected = all_hashes.iter().max_by_key(|kh| kh.0).unwrap();

        let txn = tree.store.begin_read().unwrap();
        let reader = TransactionReader::new(txn).unwrap();
        let (_, leaf) = reader.get_rightmost_leaf().unwrap().expect("should have leaves");
        assert_eq!(leaf.key_hash(), *expected);
    }

    /// Regression: deleting the rightmost leaf and inserting a smaller one must not
    /// cause `get_rightmost_leaf` to return a deleted node.
    #[test]
    fn rightmost_leaf_not_stale_after_delete_and_smaller_insert() {
        use jmt::storage::TreeReader;

        let dir = tempfile::tempdir().unwrap();
        let tree = StateTree::open(dir.path().join("rl_del.redb")).unwrap();

        let mut keys = [
            NamespaceKey::identity([0x01; 32]),
            NamespaceKey::identity([0xFF; 32]),
            NamespaceKey::identity([0x80; 32]),
        ];
        keys.sort_by_key(|k| k.key_hash().0);
        let k_smallest = &keys[0];
        let k_middle = &keys[1];
        let k_largest = &keys[2];

        tree.put(k_smallest, vec![1]).unwrap();
        tree.put(k_largest, vec![2]).unwrap();

        tree.delete(k_largest).unwrap();
        tree.put(k_middle, vec![3]).unwrap();

        let txn = tree.store.begin_read().unwrap();
        let reader = TransactionReader::new(txn).unwrap();
        let (_, leaf) = reader.get_rightmost_leaf().unwrap().expect("should have leaves");

        assert_ne!(
            leaf.key_hash(),
            k_largest.key_hash(),
            "rightmost metadata must not point at a deleted leaf"
        );

        let live_hashes = [k_smallest.key_hash(), k_middle.key_hash()];
        let expected_rightmost = live_hashes.iter().max_by_key(|kh| kh.0).unwrap();
        assert_eq!(
            leaf.key_hash(),
            *expected_rightmost,
            "rightmost metadata should reflect the actual rightmost live leaf"
        );
    }

    /// Deleting the only leaf in the tree must clear the rightmost metadata.
    #[test]
    fn rightmost_leaf_cleared_when_tree_emptied() {
        use jmt::storage::TreeReader;

        let dir = tempfile::tempdir().unwrap();
        let tree = StateTree::open(dir.path().join("rl_empty.redb")).unwrap();

        let key = NamespaceKey::identity([0x42; 32]);
        tree.put(&key, vec![1]).unwrap();
        tree.delete(&key).unwrap();

        let txn = tree.store.begin_read().unwrap();
        let reader = TransactionReader::new(txn).unwrap();
        assert!(
            reader.get_rightmost_leaf().unwrap().is_none(),
            "rightmost should be None after deleting the only leaf"
        );
    }

    /// Hardcoded test vector: inserting `b"known_vector_value"` at
    /// `NamespaceKey::identity(b"known_vector_key")` into a fresh tree must produce
    /// this exact root. Catches silent jmt-side hasher/node-ordering/serialization
    /// changes across upgrades that `root_is_deterministic` would miss.
    #[test]
    fn root_matches_known_vector() {
        let tree = test_tree();
        let key = NamespaceKey::identity(b"known_vector_key");
        let root = tree.put(&key, b"known_vector_value".to_vec()).unwrap();

        let expected = StateRoot([
            0xd2, 0x45, 0xc8, 0x47, 0x90, 0x8e, 0x05, 0x93, 0x46, 0x3f, 0xdc, 0x87, 0xb2, 0x89, 0xc9, 0xb4, 0xf3, 0x75,
            0x31, 0x32, 0xc8, 0x5f, 0xcf, 0x32, 0x54, 0xff, 0xa2, 0x2f, 0x3a, 0xe4, 0xf7, 0x9f,
        ]);
        assert_eq!(
            root, expected,
            "root hash changed — this breaks inter-operator determinism. \
             If the jmt crate was intentionally upgraded, update the expected vector."
        );
    }

    #[test]
    fn proof_borsh_roundtrip() {
        let tree = test_tree();
        let key = NamespaceKey::identity([0xab; 32]);
        let root = tree.put(&key, vec![1, 2, 3]).unwrap();

        let leaf = tree.proof(&key).unwrap().unwrap();
        let bytes = leaf.to_bytes().unwrap();
        let restored = LeafProof::from_bytes(&bytes).unwrap();

        assert_eq!(restored.value(), Some([1, 2, 3].as_slice()));
        assert_eq!(*restored.root(), root);
        restored.verify(&key, &root).unwrap();
    }
}
