//! The manifest repository's log of live-provenance §4.7: an RFC 9162 §2.1 Merkle tree over
//! SHA-256, and its receipts of inclusion as RFC 9942 §5.2.1 encodes them (`RFC9162_SHA256`,
//! a `COSE_Sign1_Tagged` over the tree's root, detached, with the inclusion proof in the
//! unprotected header).

use super::cose::{self, int_map, nint, Sign1};
use crate::live_watermark::cbor::{encode, read_one, Item};
use ed25519_dalek::SigningKey;
use sha2::{Digest, Sha256};

/// RFC 9942's `vds` header parameter.
pub const H_VDS: i64 = 395;
/// RFC 9942's `vdp` header parameter.
pub const H_VDP: i64 = 396;
/// `RFC9162_SHA256` (RFC 9942 §5.1).
pub const VDS_RFC9162_SHA256: u64 = 1;
/// The inclusion proof's label in a `vdp` map (RFC 9942 §5.2.1).
pub const INCLUSION: i64 = -1;

/// RFC 9162 §2.1.1: the leaf hash, SHA-256(0x00 ‖ d).
pub fn leaf(entry: &[u8]) -> [u8; 32] {
    let mut h = Sha256::new();
    h.update([0u8]);
    h.update(entry);
    h.finalize().into()
}

/// RFC 9162 §2.1.1: an interior node, SHA-256(0x01 ‖ left ‖ right).
pub fn node(l: &[u8; 32], r: &[u8; 32]) -> [u8; 32] {
    let mut h = Sha256::new();
    h.update([1u8]);
    h.update(l);
    h.update(r);
    h.finalize().into()
}

/// The largest power of two below `n` (n > 1).
fn split(n: usize) -> usize {
    let mut k = 1;
    while k * 2 < n {
        k *= 2;
    }
    k
}

/// RFC 9162 §2.1.1: the Merkle Tree Hash of `entries`.
pub fn mth(entries: &[Vec<u8>]) -> [u8; 32] {
    match entries.len() {
        0 => Sha256::digest([]).into(),
        1 => leaf(&entries[0]),
        n => {
            let k = split(n);
            node(&mth(&entries[..k]), &mth(&entries[k..]))
        }
    }
}

/// RFC 9162 §2.1.3.1: the inclusion proof of entry `m`.
pub fn path(m: usize, entries: &[Vec<u8>]) -> Vec<[u8; 32]> {
    let n = entries.len();
    if n <= 1 {
        return vec![];
    }
    let k = split(n);
    if m < k {
        let mut p = path(m, &entries[..k]);
        p.push(mth(&entries[k..]));
        p
    } else {
        let mut p = path(m - k, &entries[k..]);
        p.push(mth(&entries[..k]));
        p
    }
}

/// RFC 9162 §2.1.3.2: the root an inclusion proof of `entry` at `leaf_index` in a tree of
/// `tree_size` leads to, or none where the proof fails.
pub fn root_of(entry: &[u8], leaf_index: u64, tree_size: u64, inclusion_path: &[[u8; 32]]) -> Option<[u8; 32]> {
    if leaf_index >= tree_size {
        return None;
    }
    let (mut f, mut s) = (leaf_index, tree_size - 1);
    let mut r = leaf(entry);
    for p in inclusion_path {
        if s == 0 {
            return None;
        }
        if f & 1 == 1 || f == s {
            r = node(p, &r);
            if f & 1 == 0 {
                while f & 1 == 0 && f != 0 {
                    f >>= 1;
                    s >>= 1;
                }
            }
        } else {
            r = node(&r, p);
        }
        f >>= 1;
        s >>= 1;
    }
    (s == 0).then_some(r)
}

fn protected() -> Item {
    int_map(vec![(1, nint(cose::ALG_EDDSA)), (H_VDS, Item::Uint(VDS_RFC9162_SHA256))])
}

/// A receipt of inclusion (RFC 9942 §5.2.1) of entry `index` of `entries`, signed by `key`.
pub fn receipt(entries: &[Vec<u8>], index: usize, key: &SigningKey) -> Vec<u8> {
    let content = encode(&Item::Array(vec![
        Item::Uint(entries.len() as u64),
        Item::Uint(index as u64),
        Item::Array(path(index, entries).iter().map(|h| Item::Bytes(h.to_vec())).collect()),
    ]));
    let unprotected = int_map(vec![(H_VDP, int_map(vec![(INCLUSION, Item::Array(vec![Item::Bytes(content)]))]))]);
    cose::sign1(&protected(), unprotected, None, &mth(entries), key)
}

/// A receipt of inclusion, read and verified: its protected header is `{1: -8, 395: 1}`, its
/// unprotected header one inclusion proof and nothing else, its payload detached; the proof
/// of `entry` leads to a root (RFC 9942 §5.2.1 step 1) under which `key`'s signature verifies
/// (step 2). The tree's size and the entry's index.
pub fn verify_receipt(receipt: &[u8], entry: &[u8], key: &[u8; 32]) -> Option<(u64, u64)> {
    let s = Sign1::read(receipt)?;
    if s.protected != protected() || s.payload.is_some() {
        return None;
    }
    let Item::Map(u) = &s.unprotected else { return None };
    let [(label, vdp)] = u.as_slice() else { return None };
    if label.int() != Some(H_VDP) {
        return None;
    }
    let Item::Map(v) = vdp else { return None };
    let [(il, Item::Array(proofs))] = v.as_slice() else { return None };
    let [Item::Bytes(content)] = proofs.as_slice() else { return None };
    if il.int() != Some(INCLUSION) {
        return None;
    }
    let Item::Array(c) = read_one(content).ok()? else { return None };
    let [Item::Uint(size), Item::Uint(index), Item::Array(p)] = c.as_slice() else { return None };
    if p.is_empty() {
        return None;
    }
    let p: Vec<[u8; 32]> =
        p.iter().map(|h| if let Item::Bytes(b) = h { b.as_slice().try_into().ok() } else { None }).collect::<Option<_>>()?;
    let root = root_of(entry, *index, *size, &p)?;
    s.verify(key, &root).then_some((*size, *index))
}

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

    #[test]
    fn every_inclusion_proof_leads_to_the_tree_hash() {
        for n in 1..=9usize {
            let entries: Vec<Vec<u8>> = (0..n).map(|i| vec![i as u8; 3]).collect();
            let root = mth(&entries);
            for m in 0..n {
                assert_eq!(root_of(&entries[m], m as u64, n as u64, &path(m, &entries)), Some(root), "{m} of {n}");
                if n > 1 {
                    assert_ne!(root_of(&entries[m], ((m + 1) % n) as u64, n as u64, &path(m, &entries)), Some(root));
                }
            }
        }
    }

    #[test]
    fn a_receipt_verifies_for_its_entry_and_key_only() {
        let k = SigningKey::from_bytes(&[4; 32]);
        let entries: Vec<Vec<u8>> = (0..6u8).map(|i| vec![i; 32]).collect();
        let r = receipt(&entries, 4, &k);
        let pk = k.verifying_key().to_bytes();
        assert_eq!(verify_receipt(&r, &entries[4], &pk), Some((6, 4)));
        assert_eq!(verify_receipt(&r, &entries[3], &pk), None);
        assert_eq!(verify_receipt(&r, &entries[4], &[9; 32]), None);
    }
}
