//! The chunk-tree form `wai.prior.tree` (extensions/prior-carriage.md §19).
//!
//! A tree head (`WTR1`) names a one-file set by its length, its SHA-256 and its BLAKE3 hash,
//! and lists the BLAKE3 chaining value of each of its pieces: runs of 2^g of BLAKE3's 1024-byte
//! chunks, each a subtree of BLAKE3's tree over the whole set. The pin covers the head, so a
//! sink verifies each piece as it arrives, from any peer, and refuses a changed piece at that
//! piece; the pieces' values give the root, which is the set's BLAKE3 hash.

use super::sealed;
use crate::codecs::CAP_PRIOR_SEALED;
use crate::pin::sha256;
use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, HasherExt, Mode};

/// The magic a tree head begins with.
pub const MAGIC: [u8; 4] = *b"WTR1";
/// The largest piece exponent: a piece of 64 MiB.
pub const MAX_G: u8 = 16;
/// The most pieces a head lists.
pub const MAX_PIECES: u64 = 65536;

pub const TREE_HEAD_MALFORMED: &str = "tree_head_malformed";
pub const TREE_INNER_FORM: &str = "tree_inner_form";
pub const TREE_ROOT_MISMATCH: &str = "tree_root_mismatch";
pub const TREE_PIECE_MISMATCH: &str = "tree_piece_mismatch";
pub const TREE_DIGEST_MISMATCH: &str = "tree_digest_mismatch";

/// The inner forms a tree admits: a sealed set's, and a sealed set.
pub fn inner_form_admitted(form: &str) -> bool {
    sealed::inner_form_admitted(form) || form == CAP_PRIOR_SEALED
}

/// A tree head, read.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Head {
    pub g: u8,
    pub inner_form: String,
    pub len: u64,
    /// The SHA-256 of the bytes: the set's digest in its inner form.
    pub inner: [u8; 32],
    /// The BLAKE3 hash of the bytes.
    pub root: [u8; 32],
    pub values: Vec<[u8; 32]>,
}

impl Head {
    /// The length of a piece.
    pub fn piece_len(&self) -> u64 {
        1024u64 << self.g
    }
    /// Piece `i`'s byte range, `[start, end)`.
    pub fn range(&self, i: usize) -> (u64, u64) {
        let p = self.piece_len();
        (i as u64 * p, ((i as u64 + 1) * p).min(self.len))
    }
}

/// The chaining value of `bytes` as the subtree of BLAKE3's tree that begins at byte `start`.
fn subtree_cv(bytes: &[u8], start: u64) -> ChainingValue {
    let mut h = blake3::Hasher::new();
    h.set_input_offset(start);
    h.update(bytes);
    h.finalize_non_root()
}

/// The chaining value of pieces `v[..]` merged as BLAKE3 merges subtrees: the left subtree the
/// largest power of two of pieces below their count.
fn merged(v: &[[u8; 32]]) -> ChainingValue {
    if v.len() == 1 {
        return v[0];
    }
    let l = 1usize << (usize::BITS - 1 - (v.len() - 1).leading_zeros());
    merge_subtrees_non_root(&merged(&v[..l]), &merged(&v[l..]), Mode::Hash)
}

/// The root the values give: BLAKE3's root node over them, at least two.
pub fn root_of(v: &[[u8; 32]]) -> [u8; 32] {
    let l = 1usize << (usize::BITS - 1 - (v.len() - 1).leading_zeros());
    *merge_subtrees_root(&merged(&v[..l]), &merged(&v[l..]), Mode::Hash).as_bytes()
}

fn le(b: &[u8]) -> u64 {
    b.iter().rev().fold(0u64, |a, x| a << 8 | *x as u64)
}

/// §19's head checks, in order: the layout (`tree_head_malformed`), the inner form
/// (`tree_inner_form`), and the root the values give (`tree_root_mismatch`).
pub fn read_head(b: &[u8]) -> Result<Head, &'static str> {
    let bad = TREE_HEAD_MALFORMED;
    if b.len() < 6 || b[..4] != MAGIC {
        return Err(bad);
    }
    let (g, k) = (b[4], b[5] as usize);
    if g > MAX_G || k > 64 || b.len() < 82 + k {
        return Err(bad);
    }
    let form = &b[6..6 + k];
    if !form.iter().all(|c| (0x21..=0x7e).contains(c)) {
        return Err(bad);
    }
    let inner_form = String::from_utf8(form.to_vec()).map_err(|_| bad)?;
    let len = le(&b[6 + k..14 + k]);
    let inner: [u8; 32] = b[14 + k..46 + k].try_into().map_err(|_| bad)?;
    let root: [u8; 32] = b[46 + k..78 + k].try_into().map_err(|_| bad)?;
    let n = le(&b[78 + k..82 + k]);
    let p = 1024u64 << g;
    if len == 0 || n != len.div_ceil(p) || !(2..=MAX_PIECES).contains(&n) || b.len() as u64 != 82 + k as u64 + 32 * n {
        return Err(bad);
    }
    if !inner_form_admitted(&inner_form) {
        return Err(TREE_INNER_FORM);
    }
    let values: Vec<[u8; 32]> = b[82 + k..].chunks(32).map(|c| c.try_into().expect("32 bytes")).collect();
    if root_of(&values) != root {
        return Err(TREE_ROOT_MISMATCH);
    }
    Ok(Head { g, inner_form, len, inner, root, values })
}

/// The head of `bytes`, a set in `inner_form`, with pieces of 1024·2^`g` bytes.
pub fn write_head(bytes: &[u8], inner_form: &str, g: u8) -> Vec<u8> {
    let p = 1024usize << g;
    let values: Vec<[u8; 32]> = bytes.chunks(p).enumerate().map(|(i, c)| subtree_cv(c, (i * p) as u64)).collect();
    let mut out = MAGIC.to_vec();
    out.push(g);
    out.push(inner_form.len() as u8);
    out.extend_from_slice(inner_form.as_bytes());
    out.extend_from_slice(&(bytes.len() as u64).to_le_bytes());
    out.extend_from_slice(&sha256(bytes));
    out.extend_from_slice(blake3::hash(bytes).as_bytes());
    out.extend_from_slice(&(values.len() as u32).to_le_bytes());
    for v in &values {
        out.extend_from_slice(v);
    }
    out
}

/// Whether `piece` is piece `i` of the set `head` names.
pub fn piece_verifies(head: &Head, i: usize, piece: &[u8]) -> bool {
    let (s, e) = head.range(i);
    piece.len() as u64 == e - s && subtree_cv(piece, s) == head.values[i]
}

/// A source's answer to a request for the bytes `[start, end)`: its content coding, if any,
/// and the bytes it sends.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Range {
    pub content_encoding: Option<String>,
    pub body: Vec<u8>,
}

/// A source: its name, and what it answers for a range (`None`: no answer).
pub type Source<'a> = (&'a str, &'a dyn Fn(u64, u64) -> Option<Range>);

/// Each piece asked, from which source, and what came of it; then the set's bytes or the
/// reason it is not held.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TreeFetch {
    pub attempts: Vec<(String, usize, &'static str)>,
    pub result: Result<Vec<u8>, &'static str>,
}

/// Fetches every piece the head names, piece `i` first from source `i mod m` and then from the
/// sources after it, verifying each as it arrives (§19); then requires the whole set's SHA-256
/// to be the head's `inner`.
pub fn fetch_pieces(head: &Head, sources: &[Source<'_>]) -> TreeFetch {
    let mut attempts = Vec::new();
    let mut bytes = Vec::with_capacity(head.len as usize);
    for i in 0..head.values.len() {
        let (s, e) = head.range(i);
        let mut got = None;
        for j in 0..sources.len() {
            let (name, serve) = sources[(i + j) % sources.len()];
            let outcome = match serve(s, e) {
                None => "set_location_unavailable",
                Some(r) if r.content_encoding.is_some() => "content_coding_present",
                Some(r) if r.body.len() as u64 != e - s => "content_range_mismatch",
                Some(r) if !piece_verifies(head, i, &r.body) => TREE_PIECE_MISMATCH,
                Some(r) => {
                    got = Some(r.body);
                    "held"
                }
            };
            attempts.push((name.to_string(), i, outcome));
            if got.is_some() {
                break;
            }
        }
        match got {
            Some(b) => bytes.extend_from_slice(&b),
            None => return TreeFetch { attempts, result: Err("prior_not_held") },
        }
    }
    let result = if sha256(&bytes) == head.inner { Ok(bytes) } else { Err(TREE_DIGEST_MISMATCH) };
    TreeFetch { attempts, result }
}

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

    fn data(n: usize) -> Vec<u8> {
        (0..n as u32).map(|i| (i.wrapping_mul(2654435761) >> 13) as u8).collect()
    }

    #[test]
    fn the_root_is_the_blake3_hash_for_every_length_and_piece_size() {
        for g in 0..3u8 {
            for n in [1025usize, 2048, 2049, 3 * 1024 + 7, 5 * 1024, 9 * 1024 + 1, 33 * 1024 - 1] {
                if n.div_ceil(1024 << g) < 2 {
                    continue;
                }
                let d = data(n);
                let h = read_head(&write_head(&d, "", g)).expect("reads");
                assert_eq!(h.root, *blake3::hash(&d).as_bytes(), "{g} {n}");
                for i in 0..h.values.len() {
                    let (s, e) = h.range(i);
                    assert!(piece_verifies(&h, i, &d[s as usize..e as usize]));
                }
            }
        }
    }

    #[test]
    fn a_flipped_piece_is_refused_at_that_piece() {
        let d = data(5000);
        let head = read_head(&write_head(&d, "", 0)).expect("reads");
        let mut bad = d.clone();
        bad[3 * 1024 + 5] ^= 1;
        let good = |s: u64, e: u64| Some(Range { content_encoding: None, body: d[s as usize..e as usize].to_vec() });
        let flipped = |s: u64, e: u64| Some(Range { content_encoding: None, body: bad[s as usize..e as usize].to_vec() });
        let f = fetch_pieces(&head, &[("a", &good), ("b", &flipped)]);
        assert_eq!(f.result.as_deref(), Ok(&d[..]));
        let refused: Vec<_> = f.attempts.iter().filter(|a| a.2 != "held").collect();
        assert_eq!(refused, [&("b".to_string(), 3, TREE_PIECE_MISMATCH)]);
    }

    #[test]
    fn heads_that_do_not_read() {
        let d = data(5000);
        let mut h = write_head(&d, "", 0);
        assert!(read_head(&write_head(&data(900), "", 0)).is_err());
        for inner in ["wai.prior.nnc", "wai.prior.set", "wai.prior.tree", "x.unregistered"] {
            assert_eq!(read_head(&write_head(&d, inner, 0)), Err(TREE_INNER_FORM), "{inner}");
        }
        for inner in ["", "wai.prior.safetensors", "wai.prior.gguf", "wai.prior.sealed"] {
            assert!(read_head(&write_head(&d, inner, 0)).is_ok(), "{inner}");
        }
        let last = h.len() - 1;
        h[last] ^= 1;
        assert_eq!(read_head(&h), Err(TREE_ROOT_MISMATCH));
        assert_eq!(read_head(&h[..last]), Err(TREE_HEAD_MALFORMED));
    }
}
