//! The `WHP1` head proof (http-delivery §6): the BLAKE3 chaining values that prove an
//! object's bytes [0, covered_len) against its content hash without the rest.
//!
//! The content hash is the BLAKE3 root over `X = "moq-jwp:object\x01" ‖ envelope`
//! (jwp-receipts §2.1). The proof carries the chaining value of every subtree of `X`
//! that lies wholly at or after `c = 15 + covered_len`, in the order §6.3's procedure
//! consumes them; a verifier computes every subtree before `c` from the bytes it holds
//! and merges up to the root.

use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, HasherExt, Mode};

use super::head::{covered_len, parse_head, HeadParse};
use super::{RangeRefusal, DOMAIN_LEN};
use crate::jwp_hash::DOMAIN_OBJECT;

/// The file's magic, "WHP1".
pub const MAGIC: &[u8; 4] = b"WHP1";
// Kept out of the crate's public constants, which the C header generator exports.
const MAX_CV: usize = HeadProof::MAX_CV;
const FIXED_LEN: usize = HeadProof::FIXED_LEN;

const CHUNK: u64 = 1024;

/// A parsed head proof.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HeadProof {
    pub cid: [u8; 32],
    pub envelope_len: u64,
    pub covered_len: u64,
    pub cvs: Vec<[u8; 32]>,
}

/// The length of the left subtree of a BLAKE3 subtree of `n` bytes (`n` > 1024):
/// `1024 · 2^⌊log₂((n−1)/1024)⌋`, the largest power-of-two number of chunks that
/// leaves at least one byte to the right.
fn left_len(n: u64) -> u64 {
    let q = (n - 1) / CHUNK;
    CHUNK << (63 - q.leading_zeros())
}

/// The root of a tree from its two children: BLAKE3's parent node with the ROOT flag.
fn merge_root(l: &ChainingValue, r: &ChainingValue) -> [u8; 32] {
    *merge_subtrees_root(l, r, Mode::Hash).as_bytes()
}

/// The non-root chaining value of the subtree `x[..len]` whose first chunk is at
/// byte offset `start` of the whole input.
fn subtree_cv(x: &[u8], start: u64) -> ChainingValue {
    let mut h = blake3::Hasher::new();
    h.set_input_offset(start);
    h.update(x);
    h.finalize_non_root()
}

/// §6.3's `cv(start, len)`: the chaining value of the subtree [start, start + len) of
/// `X`, computed from `x_prefix` (`X`'s first `c` bytes) where the subtree lies wholly
/// before `c`, taken from `it` where it lies wholly at or after `c`, and merged from
/// its two children otherwise.
fn node_cv(
    x_prefix: &[u8],
    start: u64,
    len: u64,
    c: u64,
    it: &mut dyn Iterator<Item = [u8; 32]>,
) -> Result<ChainingValue, RangeRefusal> {
    if start >= c {
        it.next().ok_or(RangeRefusal::HeadProofMismatch)
    } else if start + len <= c {
        Ok(subtree_cv(&x_prefix[start as usize..(start + len) as usize], start))
    } else {
        let l = left_len(len);
        let left = node_cv(x_prefix, start, l, c, it)?;
        let right = node_cv(x_prefix, start + l, len - l, c, it)?;
        Ok(merge_subtrees_non_root(&left, &right, Mode::Hash))
    }
}

/// The producer's walk: the chaining values of the subtrees at or after `c`, in the
/// order [`node_cv`] consumes them.
fn collect_cvs(x: &[u8], start: u64, len: u64, c: u64, out: &mut Vec<[u8; 32]>) {
    if start >= c {
        out.push(subtree_cv(&x[start as usize..(start + len) as usize], start));
    } else if start + len > c {
        let l = left_len(len);
        collect_cvs(x, start, l, c, out);
        collect_cvs(x, start + l, len - l, c, out);
    }
}

impl HeadProof {
    /// The most chaining values a proof may carry.
    pub const MAX_CV: usize = 64;
    /// The fixed part of the file: magic, cid, `envelope_len`, `covered_len`, `n_cv`.
    pub const FIXED_LEN: usize = 53;

    /// The head proof of `envelope`, whose head must parse (§5.1 included).
    pub fn build(envelope: &[u8]) -> Result<HeadProof, RangeRefusal> {
        let envelope_len = envelope.len() as u64;
        let head = match parse_head(envelope, envelope_len)? {
            HeadParse::Head(h) => h,
            HeadParse::Need(n) => {
                return Err(RangeRefusal::HeadMalformed(format!("the head needs {n} bytes; the object has {envelope_len}")))
            }
        };
        let cid = super::content_hash(envelope);
        let covered = covered_len(head.head_len, envelope_len);
        let mut cvs = Vec::new();
        if covered < envelope_len {
            let mut x = Vec::with_capacity(DOMAIN_OBJECT.len() + envelope.len());
            x.extend_from_slice(DOMAIN_OBJECT);
            x.extend_from_slice(envelope);
            let (l_total, c) = (x.len() as u64, DOMAIN_LEN + covered);
            let l = left_len(l_total);
            collect_cvs(&x, 0, l, c, &mut cvs);
            collect_cvs(&x, l, l_total - l, c, &mut cvs);
        }
        Ok(HeadProof { cid, envelope_len, covered_len: covered, cvs })
    }

    /// The `WHP1` file.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(FIXED_LEN + 32 * self.cvs.len());
        out.extend_from_slice(MAGIC);
        out.extend_from_slice(&self.cid);
        out.extend_from_slice(&self.envelope_len.to_le_bytes());
        out.extend_from_slice(&self.covered_len.to_le_bytes());
        out.push(self.cvs.len() as u8);
        for cv in &self.cvs {
            out.extend_from_slice(cv);
        }
        out
    }

    /// Read a `WHP1` file: its magic, its length, `n_cv` at most [`MAX_CV`], and a
    /// `covered_len` that is the object's length or ends `X` on a chunk boundary
    /// before the object's end (§6.3 check 1, the part the file alone decides).
    pub fn from_bytes(b: &[u8]) -> Result<HeadProof, RangeRefusal> {
        if b.len() < FIXED_LEN || &b[..4] != MAGIC {
            return Err(RangeRefusal::HeadProofMalformed("not a WHP1 file"));
        }
        let n = b[52] as usize;
        if n > MAX_CV {
            return Err(RangeRefusal::HeadProofMalformed("more than 64 chaining values"));
        }
        if b.len() != FIXED_LEN + 32 * n {
            return Err(RangeRefusal::HeadProofMalformed("the file is not 53 + 32·n_cv bytes"));
        }
        let cid: [u8; 32] = b[4..36].try_into().expect("32 bytes");
        let envelope_len = u64::from_le_bytes(b[36..44].try_into().expect("8 bytes"));
        let covered = u64::from_le_bytes(b[44..52].try_into().expect("8 bytes"));
        // `X` is the domain and the envelope: its length must be a u64 too.
        if envelope_len > u64::MAX - DOMAIN_LEN || covered > envelope_len {
            return Err(RangeRefusal::HeadProofMalformed("envelope_len or covered_len past what a u64 offset into X can name"));
        }
        let aligned = covered < envelope_len && (DOMAIN_LEN + covered) % CHUNK == 0;
        if !(covered == envelope_len || aligned) {
            return Err(RangeRefusal::HeadProofMalformed("covered_len neither ends a chunk nor is the object's length"));
        }
        if covered == envelope_len && n != 0 {
            return Err(RangeRefusal::HeadProofMalformed("a proof that covers the whole object carries chaining values"));
        }
        let cvs = b[FIXED_LEN..].chunks_exact(32).map(|c| c.try_into().expect("32 bytes")).collect();
        Ok(HeadProof { cid, envelope_len, covered_len: covered, cvs })
    }

    /// §6.3: check this proof against the object's `cid` and the bytes `prefix` (at
    /// least [0, covered_len) of the object), in order: the proof names this object
    /// (`head_proof_malformed`); its chaining values are consumed exactly and the
    /// root equals the cid (`head_proof_mismatch`); and its `envelope_len` is the
    /// length the verified head gives and its `covered_len` §2's formula for that head
    /// (`head_proof_malformed`). No byte of `prefix` is parsed before its hash is
    /// checked.
    pub fn verify(&self, cid: &[u8; 32], prefix: &[u8]) -> Result<(), RangeRefusal> {
        if &self.cid != cid {
            return Err(RangeRefusal::HeadProofMalformed("the proof names another object"));
        }
        if (prefix.len() as u64) < self.covered_len {
            return Err(RangeRefusal::HeadProofMalformed("fewer bytes are held than the proof covers"));
        }
        let covered = &prefix[..self.covered_len as usize];
        if self.covered_len == self.envelope_len {
            if !self.cvs.is_empty() {
                return Err(RangeRefusal::HeadProofMalformed("a proof that covers the whole object carries chaining values"));
            }
            if &super::content_hash(covered) != cid {
                return Err(RangeRefusal::HeadProofMismatch);
            }
        } else {
            let mut x = Vec::with_capacity(DOMAIN_OBJECT.len() + covered.len());
            x.extend_from_slice(DOMAIN_OBJECT);
            x.extend_from_slice(covered);
            let past = RangeRefusal::HeadProofMalformed("envelope_len or covered_len past what a u64 offset into X can name");
            let l_total = DOMAIN_LEN.checked_add(self.envelope_len).ok_or(past.clone())?;
            let c = DOMAIN_LEN.checked_add(self.covered_len).ok_or(past)?;
            if c % CHUNK != 0 || c >= l_total {
                return Err(RangeRefusal::HeadProofMalformed("covered_len neither ends a chunk nor is the object's length"));
            }
            let mut it = self.cvs.iter().copied();
            let l = left_len(l_total);
            let left = node_cv(&x, 0, l, c, &mut it)?;
            let right = node_cv(&x, l, l_total - l, c, &mut it)?;
            if it.next().is_some() || &merge_root(&left, &right) != cid {
                return Err(RangeRefusal::HeadProofMismatch);
            }
        }
        // The bytes are the object's: now read the head they hold, whose fields give the
        // object's length (§5.1).
        let head = match parse_head(covered, self.envelope_len) {
            Ok(HeadParse::Head(h)) => h,
            Ok(HeadParse::Need(_)) => return Err(RangeRefusal::HeadProofMalformed("covered_len ends before the head does")),
            Err(RangeRefusal::ObjectLengthMismatch { .. }) => {
                return Err(RangeRefusal::HeadProofMalformed("envelope_len is not the length the head gives"))
            }
            Err(e) => return Err(e),
        };
        if covered_len(head.head_len, self.envelope_len) != self.covered_len {
            return Err(RangeRefusal::HeadProofMalformed("covered_len is not §2's formula for this head"));
        }
        Ok(())
    }
}

#[cfg(test)]
pub(super) fn left_len_for_test(n: u64) -> u64 {
    left_len(n)
}
