//! mst node decoding, shared by the full-repo walk here and the sparse scanner in
//! `sparse_mst`.

use std::borrow::Borrow;
use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::ops::Deref;
use std::sync::Arc;

use bytes::Bytes;
use cid::Cid as IpldCid;
use jacquard_common::types::crypto::{KeyCodec, PublicKey};
use jacquard_common::types::string::Did;
use jacquard_common::types::tid::Tid;
use miette::{IntoDiagnostic, Result, WrapErr};
use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use smol_str::SmolStr;

use crate::car::BlockMap;
pub(crate) use crate::car::compute_cid;

/// signed atproto repository commit object stored as the root block of a repo car.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct AtpCommit {
    pub did: Did,
    pub version: i64,
    pub data: IpldCid,
    pub rev: Tid,
    pub prev: Option<IpldCid>,
    #[serde(
        with = "jacquard_common::serde_bytes_helper",
        skip_serializing_if = "Bytes::is_empty"
    )]
    pub sig: Bytes,
}

impl AtpCommit {
    pub(crate) fn from_cbor(bytes: &[u8]) -> Result<Self> {
        serde_ipld_dagcbor::from_slice(bytes).into_diagnostic()
    }

    pub(crate) fn to_cbor(&self) -> Result<Vec<u8>> {
        serde_ipld_dagcbor::to_vec(self).into_diagnostic()
    }

    pub(crate) fn to_cid(&self) -> Result<IpldCid> {
        Ok(compute_cid(&self.to_cbor()?))
    }

    pub(crate) fn unsigned_bytes(&self) -> Result<Vec<u8>> {
        Self {
            sig: Bytes::new(),
            ..self.clone()
        }
        .to_cbor()
    }

    pub(crate) fn verify(&self, pubkey: &PublicKey<'_>) -> Result<()> {
        let unsigned = self.unsigned_bytes()?;
        match pubkey.codec {
            KeyCodec::Secp256k1 => {
                use k256::ecdsa::{Signature, VerifyingKey, signature::Verifier};
                let vk = VerifyingKey::from_sec1_bytes(pubkey.bytes.as_ref()).into_diagnostic()?;
                let sig = Signature::from_slice(self.sig.as_ref()).into_diagnostic()?;
                vk.verify(&unsigned, &sig).into_diagnostic()
            }
            KeyCodec::P256 => {
                use p256::ecdsa::{Signature, VerifyingKey, signature::Verifier};
                let vk = VerifyingKey::from_sec1_bytes(pubkey.bytes.as_ref()).into_diagnostic()?;
                let sig = Signature::from_slice(self.sig.as_ref()).into_diagnostic()?;
                vk.verify(&unsigned, &sig).into_diagnostic()
            }
            other => miette::bail!("unsupported signing key codec: {other:?}"),
        }
    }

    #[cfg(test)]
    pub(crate) fn sign_k256(
        did: Did,
        data: IpldCid,
        rev: Tid,
        prev: Option<IpldCid>,
        key: &k256::ecdsa::SigningKey,
    ) -> Self {
        use k256::ecdsa::signature::Signer;
        let mut commit = Self {
            did,
            version: 3,
            data,
            rev,
            prev,
            sig: Bytes::new(),
        };
        let unsigned = commit.unsigned_bytes().expect("unsigned commit encodes");
        let sig: k256::ecdsa::Signature = key.sign(&unsigned);
        commit.sig = Bytes::copy_from_slice(&sig.to_bytes());
        commit
    }
}

/// mst layer (depth from the leaves) of a key: `floor(leading_zero_bits(sha256(key)) / 2)`.
pub(crate) fn layer_for_key(key: &str) -> usize {
    let hash = Sha256::digest(key.as_bytes());
    let mut zeros = 0usize;
    for &byte in hash.iter() {
        if byte == 0 {
            zeros += 8;
        } else {
            zeros += byte.leading_zeros() as usize;
            break;
        }
    }
    zeros / 2
}

/// inline capacity for mst keys (`collection/rkey`). 100% of keys across the benchmark
/// corpus are <= 62 bytes (`size_of::<MstKey>() == 64`), so normal keys never touch the heap.
const INLINE_KEY_CAP: usize = 62;

#[derive(Clone)]
pub(crate) enum MstKey {
    Inline { len: u8, buf: [u8; INLINE_KEY_CAP] },
    Heap(Arc<str>),
}

impl MstKey {
    #[inline]
    pub(crate) fn new(s: impl AsRef<str>) -> Self {
        let s = s.as_ref();
        let bytes = s.as_bytes();
        if bytes.len() <= INLINE_KEY_CAP {
            let mut buf = [0u8; INLINE_KEY_CAP];
            buf[..bytes.len()].copy_from_slice(bytes);
            Self::Inline {
                len: bytes.len() as u8,
                buf,
            }
        } else {
            Self::Heap(Arc::from(s))
        }
    }

    #[inline]
    pub(crate) fn as_bytes(&self) -> &[u8] {
        match self {
            Self::Inline { len, buf } => &buf[..usize::from(*len)],
            Self::Heap(s) => s.as_bytes(),
        }
    }

    #[inline]
    pub(crate) fn as_str(&self) -> &str {
        match self {
            Self::Inline { len, buf } => {
                std::str::from_utf8(&buf[..usize::from(*len)]).expect("MstKey holds valid utf8")
            }
            Self::Heap(s) => s,
        }
    }
}

impl Deref for MstKey {
    type Target = str;

    #[inline]
    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl AsRef<str> for MstKey {
    #[inline]
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl AsRef<[u8]> for MstKey {
    #[inline]
    fn as_ref(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl Borrow<str> for MstKey {
    #[inline]
    fn borrow(&self) -> &str {
        self.as_str()
    }
}

impl PartialEq for MstKey {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl Eq for MstKey {}

impl PartialOrd for MstKey {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MstKey {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.as_bytes().cmp(other.as_bytes())
    }
}

impl Hash for MstKey {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_str().hash(state);
    }
}

impl fmt::Debug for MstKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.as_str(), f)
    }
}

impl fmt::Display for MstKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_str(), f)
    }
}

impl PartialEq<str> for MstKey {
    #[inline]
    fn eq(&self, other: &str) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl PartialEq<&str> for MstKey {
    #[inline]
    fn eq(&self, other: &&str) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl PartialEq<SmolStr> for MstKey {
    #[inline]
    fn eq(&self, other: &SmolStr) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

#[cfg(test)]
impl From<SmolStr> for MstKey {
    #[inline]
    fn from(s: SmolStr) -> Self {
        Self::new(s.as_str())
    }
}

/// where [`walk`] reads nodes from. a read hands back owned bytes and can fail, so a car
/// that isn't held in memory can serve the walk too.
pub(crate) trait BlockSource {
    fn read(&self, cid: &IpldCid) -> Result<Option<Bytes>>;
}

impl<M: BlockMap> BlockSource for M {
    fn read(&self, cid: &IpldCid) -> Result<Option<Bytes>> {
        Ok(self.get(cid).cloned())
    }
}

#[derive(Serialize, Deserialize)]
pub(crate) struct RawTreeEntry<'a> {
    #[serde(rename = "k", with = "serde_bytes", borrow)]
    pub(crate) key_suffix: &'a [u8],
    #[serde(rename = "p")]
    pub(crate) prefix_len: u8,
    #[serde(rename = "t")]
    pub(crate) tree: Option<IpldCid>,
    #[serde(rename = "v")]
    pub(crate) value: IpldCid,
}

#[derive(Serialize)]
pub(crate) struct RawNodeData<'a> {
    #[serde(rename = "e")]
    pub(crate) entries: Vec<RawTreeEntry<'a>>,
    #[serde(rename = "l")]
    pub(crate) left: Option<IpldCid>,
}

#[derive(Debug, Clone)]
pub(crate) enum FlatEntry {
    Tree { cid: IpldCid },
    Leaf { key: MstKey, cid: IpldCid },
}

/// a node's subtrees and leaves in key order, with full keys rebuilt from prefix compression.
pub(crate) fn decode_node(bytes: &[u8]) -> Result<Vec<FlatEntry>> {
    let mut entries = Vec::new();
    decode_node_into(bytes, &mut entries)?;
    Ok(entries)
}

/// appends a decoded node's entries onto `entries`, reusing its allocation when empty.
pub(crate) fn decode_node_into(bytes: &[u8], entries: &mut Vec<FlatEntry>) -> Result<()> {
    let mut de = serde_ipld_dagcbor::de::Deserializer::from_slice(bytes);
    NodeSeed(entries).deserialize(&mut de).into_diagnostic()?;
    de.end().into_diagnostic()?;
    Ok(())
}

struct NodeSeed<'a>(&'a mut Vec<FlatEntry>);

impl<'de> DeserializeSeed<'de> for NodeSeed<'_> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_map(self)
    }
}

impl<'de> Visitor<'de> for NodeSeed<'_> {
    type Value = ();

    fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("an mst node map")
    }

    fn visit_map<A>(self, mut map: A) -> std::result::Result<Self::Value, A::Error>
    where
        A: MapAccess<'de>,
    {
        let base = self.0.len();
        let mut left: Option<IpldCid> = None;
        let mut seen_entries = false;

        while let Some(field) = map.next_key::<&str>()? {
            match field {
                "l" => {
                    left = map.next_value()?;
                }
                "e" => {
                    if seen_entries {
                        return Err(serde::de::Error::duplicate_field("e"));
                    }
                    seen_entries = true;
                    map.next_value_seed(EntriesSeed(&mut *self.0))?;
                }
                _ => {
                    map.next_value::<IgnoredAny>()?;
                }
            }
        }

        if !seen_entries {
            return Err(serde::de::Error::missing_field("e"));
        }
        if let Some(cid) = left {
            self.0.insert(base, FlatEntry::Tree { cid });
        }
        Ok(())
    }
}

struct EntriesSeed<'a>(&'a mut Vec<FlatEntry>);

impl<'de> DeserializeSeed<'de> for EntriesSeed<'_> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_seq(self)
    }
}

impl<'de> Visitor<'de> for EntriesSeed<'_> {
    type Value = ();

    fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("an mst entry list")
    }

    fn visit_seq<A>(self, mut seq: A) -> std::result::Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let count = seq.size_hint().unwrap_or(0);
        if count > 0 {
            self.0.reserve(count * 2 + 1);
        }

        // each key is stored as a prefix length into the previous key plus a suffix
        let mut key = KeyBuilder::new();
        let mut idx = 0usize;
        while let Some(entry) = seq.next_element::<RawTreeEntry<'de>>()? {
            let leaf_key = key
                .advance(idx > 0, usize::from(entry.prefix_len), entry.key_suffix)
                .map_err(serde::de::Error::custom)?;
            self.0.push(FlatEntry::Leaf {
                key: leaf_key,
                cid: entry.value,
            });
            if let Some(cid) = entry.tree {
                self.0.push(FlatEntry::Tree { cid });
            }
            idx += 1;
        }
        Ok(())
    }
}

enum KeyBuilder {
    Inline {
        len: usize,
        buf: [u8; INLINE_KEY_CAP],
    },
    Heap(String),
}

impl KeyBuilder {
    #[inline]
    fn new() -> Self {
        Self::Inline {
            len: 0,
            buf: [0u8; INLINE_KEY_CAP],
        }
    }

    #[inline]
    fn advance(
        &mut self,
        check_order: bool,
        prefix_len: usize,
        suffix_bytes: &[u8],
    ) -> std::result::Result<MstKey, String> {
        let suffix = std::str::from_utf8(suffix_bytes)
            .map_err(|err| format!("invalid utf8 in mst key suffix: {err}"))?;
        let prev_bytes = match self {
            Self::Inline { len, buf } => &buf[..*len],
            Self::Heap(s) => s.as_bytes(),
        };
        if !is_utf8_char_boundary(prev_bytes, prefix_len) {
            return Err(format!("invalid mst key prefix length {prefix_len}"));
        }
        if check_order && suffix.as_bytes() <= &prev_bytes[prefix_len..] {
            let prev = std::str::from_utf8(prev_bytes).expect("KeyBuilder holds valid utf8");
            let cur = format!("{}{suffix}", &prev[..prefix_len]);
            return Err(format!("mst key {cur} does not sort after {prev}"));
        }

        let new_len = prefix_len + suffix.len();
        match self {
            Self::Inline { len, buf } if new_len <= INLINE_KEY_CAP => {
                buf[prefix_len..new_len].copy_from_slice(suffix.as_bytes());
                *len = new_len;
                Ok(MstKey::Inline {
                    len: new_len as u8,
                    buf: *buf,
                })
            }
            Self::Inline { buf, .. } => {
                let mut s = String::with_capacity(new_len.max(128));
                s.push_str(
                    std::str::from_utf8(&buf[..prefix_len]).expect("KeyBuilder holds valid utf8"),
                );
                s.push_str(suffix);
                let key = MstKey::Heap(Arc::from(s.as_str()));
                *self = Self::Heap(s);
                Ok(key)
            }
            Self::Heap(s) => {
                s.truncate(prefix_len);
                s.push_str(suffix);
                Ok(MstKey::new(s.as_str()))
            }
        }
    }
}

#[inline]
fn is_utf8_char_boundary(bytes: &[u8], index: usize) -> bool {
    if index == bytes.len() {
        return true;
    }
    match bytes.get(index) {
        None => false,
        Some(&b) => (b as i8) >= -0x40,
    }
}

/// every leaf of the mst rooted at `root`, in key order.
///
/// the tree comes from an untrusted car, so anything a valid mst cannot contain is an error:
/// a missing node, a node reached twice (which also ends cid cycles), or keys that are not
/// strictly increasing. the walk keeps its own stack, so tree depth cannot overflow the
/// thread's.
// backfill walks with `walk`, so only tests and the benches collect every leaf at once
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) fn leaves(blocks: &impl BlockSource, root: IpldCid) -> Result<Vec<(MstKey, IpldCid)>> {
    walk(blocks, root).collect()
}

/// [`leaves`] one at a time, holding only the path being walked rather than every leaf. it
/// ends after the first error.
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) fn walk<B: BlockSource + ?Sized>(blocks: &B, root: IpldCid) -> Walk<'_, B> {
    Walk {
        blocks,
        visited: HashSet::with_hasher(ahash::RandomState::new()),
        stack: vec![FlatEntry::Tree { cid: root }],
        last: None,
    }
}

#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) struct Walk<'a, B: ?Sized> {
    blocks: &'a B,
    visited: HashSet<IpldCid, ahash::RandomState>,
    stack: Vec<FlatEntry>,
    last: Option<MstKey>,
}

#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
impl<B: BlockSource + ?Sized> Walk<'_, B> {
    fn step(&mut self) -> Result<Option<(MstKey, IpldCid)>> {
        while let Some(entry) = self.stack.pop() {
            match entry {
                FlatEntry::Leaf { key, cid } => {
                    if let Some(previous) = self.last.as_ref().filter(|previous| **previous >= key)
                    {
                        miette::bail!("mst key {key} does not sort after {previous}");
                    }
                    self.last = Some(key.clone());
                    return Ok(Some((key, cid)));
                }
                FlatEntry::Tree { cid } => {
                    if !self.visited.insert(cid) {
                        miette::bail!("mst node {cid} is reachable more than once");
                    }
                    let bytes = self
                        .blocks
                        .read(&cid)?
                        .ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
                    let start = self.stack.len();
                    decode_node_into(&bytes, &mut self.stack)
                        .wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
                    self.stack[start..].reverse();
                }
            }
        }
        Ok(None)
    }
}

impl<B: BlockSource + ?Sized> Iterator for Walk<'_, B> {
    type Item = Result<(MstKey, IpldCid)>;

    fn next(&mut self) -> Option<Self::Item> {
        let step = self.step();
        if step.is_err() {
            self.stack.clear();
        }
        step.transpose()
    }
}

/// shared prefix in bytes, ending on a char boundary and capped at `u8::MAX`.
fn common_prefix_len(a: &str, b: &str) -> usize {
    let mut len = 0usize;
    for (ca, cb) in a.chars().zip(b.chars()) {
        if ca != cb || len + ca.len_utf8() > usize::from(u8::MAX) {
            break;
        }
        len += ca.len_utf8();
    }
    len
}

fn encode_node(
    left: Option<IpldCid>,
    entries: &[(&str, IpldCid, Option<IpldCid>)],
) -> Result<(IpldCid, Bytes)> {
    let mut last = "";
    let raw_entries = entries
        .iter()
        .map(|&(key, value, tree)| {
            let prefix_len = common_prefix_len(last, key);
            last = key;
            RawTreeEntry {
                key_suffix: &key.as_bytes()[prefix_len..],
                prefix_len: prefix_len as u8,
                tree,
                value,
            }
        })
        .collect();
    let node = RawNodeData {
        entries: raw_entries,
        left,
    };
    let bytes = Bytes::from(serde_ipld_dagcbor::to_vec(&node).into_diagnostic()?);
    Ok((compute_cid(&bytes), bytes))
}

/// sha-256 has 256 bits (128 2-bit chunks), so no valid mst can be deeper than this.
const MAX_MST_DEPTH: usize = 128;

#[derive(Clone)]
enum FlatNode {
    Tree { cid: IpldCid, layer: usize },
    Leaf { key: MstKey, cid: IpldCid, layer: usize },
}

impl FlatNode {
    fn layer(&self) -> usize {
        match self {
            Self::Tree { layer, .. } | Self::Leaf { layer, .. } => *layer,
        }
    }
}

/// builds an mst from `(key, record_cid)` pairs, returning `(root_cid, nodes_in_preorder)`.
#[cfg_attr(not(any(feature = "indexer", test)), allow(dead_code))]
pub(crate) fn build_tree<K: AsRef<str>>(
    leaves: impl IntoIterator<Item = (K, IpldCid)>,
) -> Result<(IpldCid, Vec<(IpldCid, Bytes)>)> {
    let mut items: Vec<FlatNode> = Vec::new();
    let mut sorted = true;
    for (key, cid) in leaves {
        let key_str = key.as_ref();
        validate_mst_key(key_str)?;
        let mst_key = MstKey::new(key_str);
        if let Some(FlatNode::Leaf { key: prev, .. }) = items.last()
            && *prev >= mst_key
        {
            sorted = false;
        }
        items.push(FlatNode::Leaf {
            key: mst_key,
            cid,
            layer: layer_for_key(key_str),
        });
    }
    if !sorted {
        items.sort_by(|a, b| match (a, b) {
            (FlatNode::Leaf { key: ka, .. }, FlatNode::Leaf { key: kb, .. }) => ka.cmp(kb),
            _ => unreachable!(),
        });
        items.dedup_by(|a, b| match (a, b) {
            (
                FlatNode::Leaf { key: ka, cid: ca, .. },
                FlatNode::Leaf { key: kb, cid: cb, .. },
            ) if ka == kb => {
                *cb = *ca;
                true
            }
            _ => false,
        });
    }
    if items.is_empty() {
        let (cid, bytes) = encode_node(None, &[])?;
        return Ok((cid, vec![(cid, bytes)]));
    }
    let top_layer = items.iter().map(FlatNode::layer).max().unwrap_or(0);
    let mut nodes = Vec::new();
    let root = build_layer(&items, top_layer, None, Some(&mut nodes))?;
    Ok((root, nodes))
}

fn build_layer(
    slice: &[FlatNode],
    layer: usize,
    blocks: Option<&dyn BlockSource>,
    mut out: Option<&mut Vec<(IpldCid, Bytes)>>,
) -> Result<IpldCid> {
    if let [FlatNode::Tree { cid, layer: l }] = slice
        && *l == layer
    {
        return Ok(*cid);
    }
    let mut expanded: Vec<FlatNode>;
    let mut slice = slice;
    while slice
        .iter()
        .any(|n| matches!(n, FlatNode::Tree { layer: l, .. } if *l == layer))
    {
        let blocks = blocks.expect("tree node requires block source");
        let mut next = Vec::with_capacity(slice.len() + 4);
        for node in slice {
            if let FlatNode::Tree { cid, layer: l } = node
                && *l == layer
            {
                expand_tree(blocks, *cid, *l, &mut next)?;
            } else {
                next.push(node.clone());
            }
        }
        expanded = next;
        slice = &expanded;
    }

    let slot = out.as_deref_mut().map(|v| {
        let s = v.len();
        v.push((IpldCid::default(), Bytes::new()));
        s
    });

    let mut start = 0usize;
    let mut left: Option<IpldCid> = None;
    let mut entries: Vec<(&str, IpldCid, Option<IpldCid>)> = Vec::new();

    for (i, node) in slice.iter().enumerate() {
        if let FlatNode::Leaf { key, cid, layer: l } = node
            && *l == layer
        {
            if start < i {
                let child = build_layer(&slice[start..i], layer - 1, blocks, out.as_deref_mut())?;
                if let Some(last) = entries.last_mut() {
                    last.2 = Some(child);
                } else {
                    left = Some(child);
                }
            }
            entries.push((key.as_str(), *cid, None));
            start = i + 1;
        }
    }
    if start < slice.len() {
        let child = build_layer(&slice[start..], layer - 1, blocks, out.as_deref_mut())?;
        if let Some(last) = entries.last_mut() {
            last.2 = Some(child);
        } else {
            left = Some(child);
        }
    }

    let (node_cid, node_bytes) = encode_node(left, &entries)?;
    if let (Some(out), Some(slot)) = (out, slot) {
        out[slot] = (node_cid, node_bytes);
    }
    Ok(node_cid)
}

fn expand_tree(
    blocks: &dyn BlockSource,
    cid: IpldCid,
    layer: usize,
    out: &mut Vec<FlatNode>,
) -> Result<()> {
    let bytes = blocks
        .read(&cid)?
        .ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
    let flat =
        decode_node(&bytes).wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
    let child_layer = layer.saturating_sub(1);
    for entry in flat {
        match entry {
            FlatEntry::Leaf { key, cid } => out.push(FlatNode::Leaf {
                layer: layer_for_key(&key),
                key,
                cid,
            }),
            FlatEntry::Tree { cid } => out.push(FlatNode::Tree {
                cid,
                layer: child_layer,
            }),
        }
    }
    Ok(())
}

fn find_layer(blocks: &impl BlockSource, cid: IpldCid, depth: usize) -> Result<usize> {
    if depth > MAX_MST_DEPTH {
        miette::bail!("mst depth exceeds maximum ({MAX_MST_DEPTH})");
    }
    let bytes = blocks
        .read(&cid)?
        .ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
    let flat =
        decode_node(&bytes).wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
    for entry in &flat {
        if let FlatEntry::Leaf { key, .. } = entry {
            return Ok(layer_for_key(key));
        }
    }
    for entry in flat {
        if let FlatEntry::Tree { cid: child } = entry {
            return Ok(find_layer(blocks, child, depth + 1)? + 1);
        }
    }
    Ok(0)
}

fn validate_mst_key(key: &str) -> Result<()> {
    if key.is_empty() || key.len() > 830 {
        miette::bail!("invalid mst key length: {}", key.len());
    }
    let Some((coll, rkey)) = key.split_once('/') else {
        miette::bail!("mst key missing '/': {key}");
    };
    if coll.is_empty() || rkey.is_empty() || rkey.contains('/') {
        miette::bail!("invalid mst key format: {key}");
    }
    if !key
        .bytes()
        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b':' | b'~' | b'-' | b'/'))
    {
        miette::bail!("invalid characters in mst key: {key}");
    }
    Ok(())
}

/// lazy in-memory mst for inverting firehose commit operations against a commit car.
pub(crate) struct MstNode {
    root: IpldCid,
    nodes: Option<Vec<FlatNode>>,
    dirty: bool,
}

impl MstNode {
    pub(crate) fn load(root: IpldCid) -> Self {
        Self {
            root,
            nodes: None,
            dirty: false,
        }
    }

    fn seek(
        &mut self,
        blocks: &impl BlockSource,
        key: &str,
        split: bool,
    ) -> Result<std::result::Result<usize, usize>> {
        let key_layer = layer_for_key(key);
        if self.nodes.is_none() {
            let layer = find_layer(blocks, self.root, 0)?;
            self.nodes = Some(vec![FlatNode::Tree {
                cid: self.root,
                layer,
            }]);
        }
        let nodes = self.nodes.as_mut().unwrap();
        for _ in 0..=MAX_MST_DEPTH {
            let i = nodes
                .iter()
                .position(|n| matches!(n, FlatNode::Leaf { key: k, .. } if k.as_str() >= key))
                .unwrap_or(nodes.len());
            if let Some(FlatNode::Leaf { key: k, .. }) = nodes.get(i)
                && k.as_str() == key
            {
                return Ok(Ok(i));
            }
            if i > 0
                && let FlatNode::Tree { cid, layer } = nodes[i - 1]
                && (split || layer >= key_layer)
            {
                let mut children = Vec::new();
                expand_tree(blocks, cid, layer, &mut children)?;
                nodes.splice((i - 1)..i, children);
            } else {
                return Ok(Err(i));
            }
        }
        miette::bail!("mst depth exceeds maximum ({MAX_MST_DEPTH})")
    }

    pub(crate) fn get(&mut self, blocks: &impl BlockSource, key: &str) -> Result<Option<IpldCid>> {
        validate_mst_key(key)?;
        match self.seek(blocks, key, false)? {
            Ok(i) => {
                let FlatNode::Leaf { cid, .. } = self.nodes.as_ref().unwrap()[i] else {
                    unreachable!()
                };
                Ok(Some(cid))
            }
            Err(_) => Ok(None),
        }
    }

    pub(crate) fn add(&mut self, blocks: &impl BlockSource, key: &str, cid: IpldCid) -> Result<()> {
        validate_mst_key(key)?;
        let layer = layer_for_key(key);
        let leaf = FlatNode::Leaf {
            key: MstKey::new(key),
            cid,
            layer,
        };
        match self.seek(blocks, key, true)? {
            Ok(i) => self.nodes.as_mut().unwrap()[i] = leaf,
            Err(i) => self.nodes.as_mut().unwrap().insert(i, leaf),
        }
        self.dirty = true;
        Ok(())
    }

    pub(crate) fn delete(&mut self, blocks: &impl BlockSource, key: &str) -> Result<IpldCid> {
        validate_mst_key(key)?;
        match self.seek(blocks, key, false)? {
            Ok(i) => {
                let FlatNode::Leaf { cid, .. } = self.nodes.as_mut().unwrap().remove(i) else {
                    unreachable!()
                };
                self.dirty = true;
                Ok(cid)
            }
            Err(_) => miette::bail!("mst key {key} not found"),
        }
    }

    pub(crate) fn root_cid(&mut self, blocks: &impl BlockSource) -> Result<IpldCid> {
        if !self.dirty {
            return Ok(self.root);
        }
        let nodes = self.nodes.as_mut().unwrap();
        let top_layer = loop {
            let Some(max_layer) = nodes.iter().map(FlatNode::layer).max() else {
                let (cid, _) = encode_node(None, &[])?;
                self.root = cid;
                self.dirty = false;
                return Ok(cid);
            };
            if let Some(i) = nodes
                .iter()
                .position(|n| matches!(n, FlatNode::Tree { layer, .. } if *layer == max_layer))
            {
                let FlatNode::Tree { cid, layer } = nodes[i] else {
                    unreachable!()
                };
                let mut children = Vec::new();
                expand_tree(blocks, cid, layer, &mut children)?;
                nodes.splice(i..=i, children);
                continue;
            }
            break max_layer;
        };
        let cid = build_layer(nodes, top_layer, Some(blocks), None)?;
        self.root = cid;
        self.dirty = false;
        Ok(cid)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::car::CarBlocks;
    use cid::multihash::Multihash;
    use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256};

    pub(crate) fn cid(byte: u8) -> IpldCid {
        let mh = Multihash::<64>::wrap(SHA2_256, &[byte; 32]).unwrap();
        IpldCid::new_v1(DAG_CBOR, mh)
    }

    /// builds a node from full keys, prefix-compressing them like the wire format.
    pub(crate) fn node<'a>(
        entries: Vec<(&'a str, IpldCid, Option<IpldCid>)>,
        left: Option<IpldCid>,
    ) -> RawNodeData<'a> {
        let mut last = "";
        let entries = entries
            .into_iter()
            .map(|(key, value, tree)| {
                let prefix_len = common_prefix_len(last, key);
                last = key;
                RawTreeEntry {
                    key_suffix: &key.as_bytes()[prefix_len..],
                    prefix_len: prefix_len as u8,
                    tree,
                    value,
                }
            })
            .collect();
        RawNodeData { entries, left }
    }

    fn blocks(nodes: Vec<(IpldCid, RawNodeData<'_>)>) -> CarBlocks {
        nodes
            .into_iter()
            .map(|(cid, node)| (cid, serde_ipld_dagcbor::to_vec(&node).unwrap().into()))
            .collect()
    }

    fn keys(leaves: &[(MstKey, IpldCid)]) -> Vec<&str> {
        leaves.iter().map(|(key, _)| key.as_str()).collect()
    }

    #[test]
    fn leaves_match_jacquard_walk() {
        let mut input: Vec<(MstKey, IpldCid)> = (0..2000)
            .map(|i| {
                let collection = [
                    "app.bsky.feed.like",
                    "app.bsky.feed.post",
                    "sh.tangled.repo",
                ][i % 3];
                (
                    MstKey::new(format!("{collection}/{i:013}")),
                    cid((i % 251) as u8),
                )
            })
            .collect();
        let (root, nodes) = build_tree(input.iter().map(|(k, c)| (k.as_str(), *c))).unwrap();
        assert!(nodes.len() > 1, "test tree should span several nodes");

        let built_map: std::collections::BTreeMap<IpldCid, Bytes> =
            nodes.iter().cloned().collect();
        let ours = leaves(&nodes.into_iter().collect::<CarBlocks>(), root).unwrap();
        input.sort_by(|a, b| a.0.cmp(&b.0));
        assert_eq!(ours, input);

        // Deleting and re-adding keys through MstNode reproduces the exact root at each step.
        let mut mst = MstNode::load(root);
        let mut remaining = input.clone();
        for i in (0..input.len()).step_by(97).rev() {
            let (key, expected_cid) = remaining.remove(i);
            assert_eq!(mst.get(&built_map, &key).unwrap(), Some(expected_cid));
            assert_eq!(mst.delete(&built_map, &key).unwrap(), expected_cid);
            assert_eq!(mst.get(&built_map, &key).unwrap(), None);
            let (expected_root, _) =
                build_tree(remaining.iter().map(|(k, c)| (k.as_str(), *c))).unwrap();
            assert_eq!(mst.root_cid(&built_map).unwrap(), expected_root);
        }
        for (key, expected_cid) in &input {
            if mst.get(&built_map, key).unwrap().is_none() {
                mst.add(&built_map, key, *expected_cid).unwrap();
            }
        }
        assert_eq!(mst.root_cid(&built_map).unwrap(), root);
    }

    #[test]
    fn walks_an_empty_tree() {
        let root = cid(1);
        assert!(
            leaves(&blocks(vec![(root, node(vec![], None))]), root)
                .unwrap()
                .is_empty()
        );
    }

    #[test]
    fn rejects_a_missing_node() {
        let (root, missing) = (cid(1), cid(2));
        let tree = blocks(vec![(
            root,
            node(vec![("a/1", cid(9), Some(missing))], None),
        )]);

        let err = leaves(&tree, root).unwrap_err();
        assert!(err.to_string().contains("missing from car"), "{err}");
    }

    #[test]
    fn walk_hands_out_leaves_up_to_a_missing_node_and_nothing_after() {
        let (root, missing) = (cid(1), cid(2));
        let tree = blocks(vec![(
            root,
            node(
                vec![("a/1", cid(9), Some(missing)), ("b/1", cid(9), None)],
                None,
            ),
        )]);

        let mut walked = walk(&tree, root);
        assert_eq!(
            walked.next().unwrap().unwrap(),
            (MstKey::new("a/1"), cid(9))
        );
        let err = walked.next().unwrap().unwrap_err();
        assert!(err.to_string().contains("missing from car"), "{err}");
        // b/1 is still on the stack, but nothing after an error counts
        assert!(walked.next().is_none());
    }

    #[test]
    fn rejects_a_subtree_reachable_twice() {
        let (root, shared) = (cid(1), cid(2));
        let tree = blocks(vec![
            (
                root,
                node(
                    vec![("a/1", cid(9), Some(shared)), ("c/1", cid(9), Some(shared))],
                    None,
                ),
            ),
            (shared, node(vec![("b/1", cid(9), None)], None)),
        ]);

        let err = leaves(&tree, root).unwrap_err();
        assert!(
            err.to_string().contains("reachable more than once"),
            "{err}"
        );
    }

    #[test]
    fn stops_on_a_node_that_is_its_own_left_subtree() {
        // a car that is not cid-verified can claim any cid for a block
        let root = cid(1);
        let tree = blocks(vec![(root, node(vec![("a/1", cid(9), None)], Some(root)))]);

        let err = leaves(&tree, root).unwrap_err();
        assert!(
            err.to_string().contains("reachable more than once"),
            "{err}"
        );
    }

    #[test]
    fn rejects_keys_out_of_order() {
        let root = cid(1);
        let tree = blocks(vec![(
            root,
            node(vec![("b/1", cid(9), None), ("a/1", cid(9), None)], None),
        )]);

        let err = leaves(&tree, root).unwrap_err();
        assert!(
            format!("{err:?}").contains("does not sort after"),
            "{err:?}"
        );
    }

    #[test]
    fn rejects_a_duplicate_key_across_subtrees() {
        let (root, right) = (cid(1), cid(2));
        let tree = blocks(vec![
            (root, node(vec![("a/1", cid(9), Some(right))], None)),
            (right, node(vec![("a/1", cid(8), None)], None)),
        ]);

        let err = leaves(&tree, root).unwrap_err();
        assert!(err.to_string().contains("does not sort after"), "{err}");
    }

    #[test]
    fn rejects_a_prefix_that_splits_a_char() {
        let root = cid(1);
        let mut bad = node(vec![("é/1", cid(9), None), ("é/2", cid(9), None)], None);
        bad.entries[1].prefix_len = 1;

        let err = leaves(&blocks(vec![(root, bad)]), root).unwrap_err();
        assert!(
            format!("{err:?}").contains("invalid mst key prefix length 1"),
            "{err:?}"
        );
    }

    #[test]
    fn rebuilds_prefix_compressed_keys() {
        let long_prefix = "io.atcr.hold.extremely.long.collection.name.exceeding.inline/";
        let long_a = format!("{long_prefix}3k2b1f5moe22a");
        let long_b = format!("{long_prefix}3k2b1f5moe22b");
        assert!(long_a.len() > INLINE_KEY_CAP);
        let root = cid(1);
        let tree = blocks(vec![(
            root,
            node(
                vec![
                    ("app.bsky.feed.like/1", cid(9), None),
                    ("app.bsky.feed.post/1", cid(9), None),
                    ("app.bsky.feed.post/2", cid(9), None),
                    (long_a.as_str(), cid(9), None),
                    (long_b.as_str(), cid(9), None),
                    ("io.atcr.hold.short/1", cid(9), None),
                    // both start with byte 0xc3, so they share no whole-char prefix
                    ("è/1", cid(9), None),
                    ("é/1", cid(9), None),
                ],
                None,
            ),
        )]);

        assert_eq!(
            keys(&leaves(&tree, root).unwrap()),
            [
                "app.bsky.feed.like/1",
                "app.bsky.feed.post/1",
                "app.bsky.feed.post/2",
                long_a.as_str(),
                long_b.as_str(),
                "io.atcr.hold.short/1",
                "è/1",
                "é/1"
            ]
        );
    }

    /// a chain of `depth` nodes that only point left, ending in one leaf, for checking that a
    /// walk does not recurse per level. returns the top cid and every encoded node.
    pub(crate) fn left_chain(depth: usize, key: &str) -> (IpldCid, Vec<(IpldCid, bytes::Bytes)>) {
        let mut child_bytes =
            serde_ipld_dagcbor::to_vec(&node(vec![(key, cid(9), None)], None)).unwrap();
        let mut child = compute_cid(&child_bytes);
        let mut nodes = vec![(child, child_bytes.into())];
        for _ in 0..depth {
            child_bytes = serde_ipld_dagcbor::to_vec(&node(vec![], Some(child))).unwrap();
            child = compute_cid(&child_bytes);
            nodes.push((child, child_bytes.into()));
        }
        (child, nodes)
    }

    #[test]
    fn walks_a_tree_deeper_than_the_thread_stack() {
        let (top, nodes) = left_chain(100_000, "a/1");
        assert_eq!(
            keys(&leaves(&nodes.into_iter().collect::<CarBlocks>(), top).unwrap()),
            ["a/1"]
        );
    }

    #[test]
    fn commit_cbor_and_signature_verification_roundtrip() {
        let did = Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").unwrap();
        let rev = Tid::new("3jzfcijpj2z2a").unwrap();
        let key = k256::ecdsa::SigningKey::from_bytes(&[7u8; 32].into()).unwrap();
        let pubkey = PublicKey {
            codec: KeyCodec::Secp256k1,
            bytes: std::borrow::Cow::Owned(
                key.verifying_key()
                    .to_encoded_point(true)
                    .as_bytes()
                    .to_vec(),
            ),
        };
        let signed = AtpCommit::sign_k256(did, cid(1), rev, Some(cid(2)), &key);
        let cbor = signed.to_cbor().unwrap();
        let decoded = AtpCommit::from_cbor(&cbor).unwrap();
        assert_eq!(decoded, signed);
        assert_eq!(decoded.to_cid().unwrap(), compute_cid(&cbor));
        decoded.verify(&pubkey).unwrap();

        let mut tampered = decoded.clone();
        tampered.version = 2;
        assert!(tampered.verify(&pubkey).is_err());
    }
}