//! Inclusion proof types and verification.

use serde::{Deserialize, Serialize};

use crate::{node_hash, KtError};

/// A single step in an inclusion proof path.
///
/// `hash` is the sibling hash; `sibling_is_left` is `true` when the sibling
/// is the left child (meaning the node being proved is the right child).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PathStep {
    pub hash: [u8; 32],
    pub sibling_is_left: bool,
}

/// A Merkle inclusion proof for a single leaf.
///
/// ## Wire format
///
/// Serialised with `bincode` and transported as the `inclusionProof :Data` field
/// in the `resolveUser` Cap'n Proto response. Clients call `verify_inclusion` to
/// authenticate the server's response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InclusionProof {
    /// 0-based index of this leaf in the log.
    pub leaf_index: u64,
    /// Number of leaves in the tree at the time the proof was generated.
    pub tree_size: u64,
    /// The 32-byte leaf hash (pre-computed from `leaf_hash(username, identity_key)`).
    pub leaf_hash: [u8; 32],
    /// Path steps from leaf level to root level (leaf-to-root order).
    pub path: Vec<PathStep>,
    /// Merkle root at the time the proof was generated.
    pub root: [u8; 32],
}

impl InclusionProof {
    /// Serialise to bytes (bincode).
    pub fn to_bytes(&self) -> Result<Vec<u8>, KtError> {
        bincode::serialize(self)
            .map_err(|e| KtError::Serialisation(e.to_string()))
    }

    /// Deserialise from bytes (bincode).
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KtError> {
        bincode::deserialize(bytes)
            .map_err(|e| KtError::Serialisation(e.to_string()))
    }
}

/// Verify that `(username, identity_key)` appears at `proof.leaf_index` in a
/// Merkle log with root `proof.root` and `proof.tree_size` leaves.
///
/// Returns `Ok(())` on success, `Err(KtError::RootMismatch)` on failure.
///
/// The caller should additionally check that `proof.root` matches a root they
/// obtained from a trusted source (e.g. a previously-pinned root or one returned
/// by a second server for cross-verification).
pub fn verify_inclusion(
    proof: &InclusionProof,
    username: &str,
    identity_key: &[u8],
) -> Result<(), KtError> {
    let expected_leaf = crate::leaf_hash(username, identity_key);
    if expected_leaf != proof.leaf_hash {
        return Err(KtError::RootMismatch);
    }

    let computed_root = recompute_root(proof.leaf_hash, &proof.path)?;

    if computed_root != proof.root {
        return Err(KtError::RootMismatch);
    }

    Ok(())
}

/// Recompute the Merkle root from a leaf hash + direction-annotated sibling path.
///
/// Each `PathStep` records the sibling hash and whether that sibling is on the
/// left (meaning the current node is on the right). This is leaf-to-root order.
fn recompute_root(leaf: [u8; 32], path: &[PathStep]) -> Result<[u8; 32], KtError> {
    let mut current = leaf;
    for step in path {
        current = if step.sibling_is_left {
            // Sibling is left, current is right.
            node_hash(&step.hash, &current)
        } else {
            // Sibling is right, current is left.
            node_hash(&current, &step.hash)
        };
    }
    Ok(current)
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::tree::MerkleLog;

    fn log_with(entries: &[(&str, &[u8])]) -> MerkleLog {
        let mut log = MerkleLog::new();
        for (u, k) in entries {
            log.append(u, k);
        }
        log
    }

    fn verify_all(log: &MerkleLog, entries: &[(&str, &[u8])]) {
        for (i, (u, k)) in entries.iter().enumerate() {
            let proof = log.inclusion_proof(i as u64).unwrap();
            verify_inclusion(&proof, u, k).unwrap_or_else(|e| {
                panic!("proof verification failed for leaf {i}: {e}");
            });
        }
    }

    #[test]
    fn single_leaf_verifies() {
        let log = log_with(&[("alice", b"KEY1")]);
        verify_all(&log, &[("alice", b"KEY1")]);
    }

    #[test]
    fn two_leaves_verify() {
        let log = log_with(&[("alice", b"K1"), ("bob", b"K2")]);
        verify_all(&log, &[("alice", b"K1"), ("bob", b"K2")]);
    }

    #[test]
    fn three_leaves_verify() {
        let log = log_with(&[("alice", b"K1"), ("bob", b"K2"), ("charlie", b"K3")]);
        verify_all(&log, &[("alice", b"K1"), ("bob", b"K2"), ("charlie", b"K3")]);
    }

    #[test]
    fn power_of_two_leaves_verify() {
        let entries: Vec<(String, Vec<u8>)> = (0u8..8)
            .map(|i| (format!("user{i}"), vec![i; 32]))
            .collect();
        let refs: Vec<(&str, &[u8])> = entries.iter().map(|(u, k)| (u.as_str(), k.as_slice())).collect();
        let log = log_with(&refs);
        verify_all(&log, &refs);
    }

    #[test]
    fn seven_leaves_all_verify() {
        let entries: Vec<(String, Vec<u8>)> = (0u8..7)
            .map(|i| (format!("u{i}"), vec![i; 32]))
            .collect();
        let refs: Vec<(&str, &[u8])> = entries.iter().map(|(u, k)| (u.as_str(), k.as_slice())).collect();
        let log = log_with(&refs);
        verify_all(&log, &refs);
    }

    #[test]
    fn wrong_identity_key_fails() {
        let log = log_with(&[("alice", b"REAL_KEY")]);
        let proof = log.inclusion_proof(0).unwrap();
        assert!(matches!(
            verify_inclusion(&proof, "alice", b"WRONG_KEY"),
            Err(KtError::RootMismatch)
        ));
    }

    #[test]
    fn tampered_sibling_fails() {
        let log = log_with(&[("alice", b"K1"), ("bob", b"K2"), ("charlie", b"K3")]);
        let mut proof = log.inclusion_proof(0).unwrap();
        if !proof.path.is_empty() {
            proof.path[0].hash[0] ^= 0xff;
        }
        assert!(matches!(
            verify_inclusion(&proof, "alice", b"K1"),
            Err(KtError::RootMismatch)
        ));
    }

    #[test]
    fn proof_serialise_roundtrip() {
        let log = log_with(&[("alice", b"K1"), ("bob", b"K2")]);
        let proof = log.inclusion_proof(0).unwrap();
        let bytes = proof.to_bytes().unwrap();
        let proof2 = InclusionProof::from_bytes(&bytes).unwrap();
        verify_inclusion(&proof2, "alice", b"K1").unwrap();
    }
}