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

use std::collections::HashSet;

use cid::Cid as IpldCid;
use jacquard_repo::mst::NodeData;
use miette::{IntoDiagnostic, Result, WrapErr};
use smol_str::SmolStr;

use crate::car::CarBlocks;

#[derive(Debug, Clone)]
pub(crate) enum FlatEntry {
    Tree { cid: IpldCid },
    Leaf { key: SmolStr, 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 node: NodeData = serde_ipld_dagcbor::from_slice(bytes).into_diagnostic()?;
    decode_node_entries(&node)
}

fn decode_node_entries(node: &NodeData) -> Result<Vec<FlatEntry>> {
    let mut entries = Vec::with_capacity(node.entries.len() * 2 + 1);
    if let Some(cid) = node.left {
        entries.push(FlatEntry::Tree { cid });
    }

    // each key is stored as a prefix length into the previous key plus a suffix
    let mut key = String::new();
    for entry in &node.entries {
        let suffix = std::str::from_utf8(&entry.key_suffix)
            .into_diagnostic()
            .wrap_err("invalid utf8 in mst key suffix")?;
        let prefix_len = usize::from(entry.prefix_len);
        if !key.is_char_boundary(prefix_len) {
            return Err(miette::miette!(
                "invalid mst key prefix length {prefix_len}"
            ));
        }
        key.truncate(prefix_len);
        key.push_str(suffix);

        entries.push(FlatEntry::Leaf {
            key: SmolStr::new(&key),
            cid: entry.value,
        });
        if let Some(cid) = entry.tree {
            entries.push(FlatEntry::Tree { cid });
        }
    }

    Ok(entries)
}

/// 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.
pub(crate) fn leaves(blocks: &CarBlocks, root: IpldCid) -> Result<Vec<(SmolStr, IpldCid)>> {
    let mut found: Vec<(SmolStr, IpldCid)> = Vec::new();
    let mut visited = HashSet::with_hasher(ahash::RandomState::new());
    let mut stack = vec![FlatEntry::Tree { cid: root }];

    while let Some(entry) = stack.pop() {
        match entry {
            FlatEntry::Leaf { key, cid } => {
                let out_of_order = found.last().filter(|(previous, _)| *previous >= key);
                if let Some((previous, _)) = out_of_order {
                    miette::bail!("mst key {key} does not sort after {previous}");
                }
                found.push((key, cid));
            }
            FlatEntry::Tree { cid } => {
                if !visited.insert(cid) {
                    miette::bail!("mst node {cid} is reachable more than once");
                }
                let bytes = blocks
                    .get(&cid)
                    .ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?;
                let entries = decode_node(bytes)
                    .wrap_err_with(|| format!("failed to decode mst node {cid}"))?;
                stack.extend(entries.into_iter().rev());
            }
        }
    }

    Ok(found)
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use cid::multihash::Multihash;
    use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256};
    use jacquard_repo::mst::TreeEntry;
    use jacquard_repo::mst::util::compute_cid;

    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(
        entries: Vec<(&str, IpldCid, Option<IpldCid>)>,
        left: Option<IpldCid>,
    ) -> NodeData {
        let mut last = "";
        let entries = entries
            .into_iter()
            .map(|(key, value, tree)| {
                let prefix_len = common_prefix_len(last, key);
                last = key;
                TreeEntry {
                    key_suffix: key.as_bytes()[prefix_len..].to_vec().into(),
                    prefix_len: prefix_len as u8,
                    tree,
                    value,
                }
            })
            .collect();
        NodeData { left, entries }
    }

    /// shared prefix in bytes, ending on a char boundary.
    fn common_prefix_len(a: &str, b: &str) -> usize {
        a.chars()
            .zip(b.chars())
            .take_while(|(a, b)| a == b)
            .map(|(a, _)| a.len_utf8())
            .sum()
    }

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

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

    #[tokio::test]
    async fn leaves_match_jacquard_walk() {
        use jacquard_repo::{MemoryBlockStore, Mst};

        let mut mst = Mst::new(std::sync::Arc::new(MemoryBlockStore::new()));
        for i in 0..2000 {
            let collection = [
                "app.bsky.feed.like",
                "app.bsky.feed.post",
                "sh.tangled.repo",
            ][i % 3];
            let key = format!("{collection}/{i:013}");
            mst = mst.add(&key, cid((i % 251) as u8)).await.unwrap();
        }
        let (root, nodes) = mst.collect_blocks().await.unwrap();
        assert!(nodes.len() > 1, "test tree should span several nodes");

        let ours = leaves(&nodes.into_iter().collect(), root).unwrap();
        assert_eq!(ours, mst.leaves().await.unwrap());
    }

    #[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 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!(err.to_string().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 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),
                    // 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",
                "è/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).unwrap();
        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).unwrap();
            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(), top).unwrap()),
            ["a/1"]
        );
    }
}