//! Binary Merkle tree primitives.
//!
//! Leaves are `blake3("leaf-v1" || cell_key || len(value) || value)`. Internal
//! nodes are `blake3("node-v1" || left || right)`. Odd-numbered layers lift
//! the unpaired hash by pairing it with itself. The empty tree commits to
//! `blake3("empty-tree-v1")`. Domain-separation prefixes guard against
//! collisions across schemes.

use ontoref_types::{canonical, AttrId, EntityId, Value};
use serde::{Deserialize, Serialize};

/// 32-byte Merkle key for a `(entity, attr)` cell.
pub type CellKey = [u8; 32];

/// 32-byte Merkle root.
pub type StateRoot = [u8; 32];

/// One step in a Merkle inclusion proof.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PathStep {
    /// Whether the proven leaf sits on the right at this level (so its
    /// sibling is on the left).
    pub on_right: bool,
    /// Sibling hash at this level.
    pub sibling: [u8; 32],
}

/// Inclusion proof for a single cell.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Witness {
    /// Cell key (= `cell_key(entity, attr)`).
    pub key: CellKey,
    /// Canonical encoding of the cell's value.
    pub value: Vec<u8>,
    /// Path of sibling hashes from leaf to root.
    pub path: Vec<PathStep>,
}

/// Domain-separated Merkle key for the `(entity, attr)` cell.
pub fn cell_key(entity: &EntityId, attr: &AttrId) -> CellKey {
    let mut hasher = blake3::Hasher::new();
    hasher.update(b"cell-v1");
    hasher.update(&[0x01]);
    push_lp(&mut hasher, entity.as_str().as_bytes());
    push_lp(&mut hasher, attr.as_str().as_bytes());
    *hasher.finalize().as_bytes()
}

/// Canonical-encoded bytes for a value, suitable as Merkle leaf payload.
pub fn encode_value(value: &Value) -> Vec<u8> {
    canonical::encode(value).expect("canonical encoding never fails for owned types")
}

pub(crate) fn leaf_hash(key: &CellKey, value: &[u8]) -> [u8; 32] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(b"leaf-v1");
    hasher.update(key);
    hasher.update(&(value.len() as u64).to_be_bytes());
    hasher.update(value);
    *hasher.finalize().as_bytes()
}

pub(crate) fn node_hash(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(b"node-v1");
    hasher.update(left);
    hasher.update(right);
    *hasher.finalize().as_bytes()
}

/// State root of an empty commit layer. Exposed so callers can compare
/// against "fresh substrate" without re-deriving the magic-string hash.
pub fn empty_root() -> StateRoot {
    *blake3::hash(b"empty-tree-v1").as_bytes()
}

/// Compute the Merkle root over a slice of (key, value-bytes) pairs.
///
/// The caller is responsible for providing entries sorted by `key`; this
/// function does not sort, so the same set of entries must always be
/// presented in the same order.
pub(crate) fn merkle_root(entries: &[(CellKey, &[u8])]) -> StateRoot {
    if entries.is_empty() {
        return empty_root();
    }
    let mut layer: Vec<[u8; 32]> = entries.iter().map(|(k, v)| leaf_hash(k, v)).collect();
    while layer.len() > 1 {
        let mut next = Vec::with_capacity(layer.len().div_ceil(2));
        for chunk in layer.chunks(2) {
            let hash = match chunk {
                [left, right] => node_hash(left, right),
                [only] => node_hash(only, only),
                _ => unreachable!("chunks(2) yields slices of length 1 or 2"),
            };
            next.push(hash);
        }
        layer = next;
    }
    layer[0]
}

/// Build an inclusion proof for `target_key` against the tree formed by
/// `entries`. Returns `None` if the key is absent.
///
/// `entries` must be sorted by key (same constraint as `merkle_root`).
pub(crate) fn build_witness(entries: &[(CellKey, &[u8])], target_key: &CellKey) -> Option<Witness> {
    let index = entries.iter().position(|(k, _)| k == target_key)?;
    let value = entries[index].1.to_vec();

    let mut layer: Vec<[u8; 32]> = entries.iter().map(|(k, v)| leaf_hash(k, v)).collect();
    let mut idx = index;
    let mut path: Vec<PathStep> = Vec::new();

    while layer.len() > 1 {
        let on_right = idx % 2 == 1;
        let sibling_idx = if on_right { idx - 1 } else { idx + 1 };
        let sibling = if sibling_idx < layer.len() {
            layer[sibling_idx]
        } else {
            // Odd man out: paired with itself.
            layer[idx]
        };
        path.push(PathStep { on_right, sibling });

        let mut next = Vec::with_capacity(layer.len().div_ceil(2));
        for chunk in layer.chunks(2) {
            let hash = match chunk {
                [left, right] => node_hash(left, right),
                [only] => node_hash(only, only),
                _ => unreachable!("chunks(2) yields slices of length 1 or 2"),
            };
            next.push(hash);
        }
        idx /= 2;
        layer = next;
    }

    Some(Witness {
        key: *target_key,
        value,
        path,
    })
}

/// Verify that `witness` is a valid inclusion proof for `root`.
pub fn verify_witness(witness: &Witness, root: &StateRoot) -> bool {
    let mut current = leaf_hash(&witness.key, &witness.value);
    for step in &witness.path {
        current = if step.on_right {
            node_hash(&step.sibling, &current)
        } else {
            node_hash(&current, &step.sibling)
        };
    }
    current == *root
}

fn push_lp(hasher: &mut blake3::Hasher, bytes: &[u8]) {
    hasher.update(&(bytes.len() as u32).to_be_bytes());
    hasher.update(bytes);
}
