//! Writing version-2 B-trees: a header (`BTHD`) and its nodes, leaves
//! (`BTLF`) and, for more records than one leaf holds, internal nodes
//! (`BTIN`) to any depth.
//!
//! Node capacities come from [`crate::btree_v2::node_info`], the arithmetic
//! libhdf5 uses (`H5B2__hdr_init`) and the reader decodes pointers with, so
//! the pointer widths the writer encodes are the ones every reader expects.

use crate::addr::saturating_usize;
#[cfg(not(feature = "std"))]
use alloc::{format, vec, vec::Vec};

use crate::btree_v2::{NodeInfo, bytes_for_max_records, node_info};
use crate::checksum::jenkins_lookup3;
use crate::error::FormatError;

/// How a B-tree is laid out: its record type and node geometry, as the
/// header records them.
#[derive(Debug, Clone, Copy)]
pub(crate) struct BTreeV2Params {
    /// Record type (5: link names, 6: link creation order, 8: attribute
    /// names, 9: attribute creation order, 10/11: chunks).
    pub(crate) tree_type: u8,
    /// Bytes per node.
    pub(crate) node_size: u32,
    /// Bytes per record.
    pub(crate) record_size: u16,
    /// Split and merge percentages. The writer fills nodes itself; these
    /// only tell libhdf5 when to split and merge as it modifies the tree.
    pub(crate) split_percent: u8,
    pub(crate) merge_percent: u8,
}

/// Size of a B-tree v2 header.
pub(crate) fn header_size(offset_size: u8, length_size: u8) -> usize {
    4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + offset_size as usize + 2 + length_size as usize + 4
}

/// Deepest tree the writer builds. Even at the smallest fan-out libhdf5's
/// arithmetic allows, a few levels hold more records than any file could.
const MAX_WRITE_DEPTH: u16 = 32;

/// Write a B-tree v2 holding `records` (`record_size` bytes each,
/// concatenated, already in the tree's key order) at `addr`: the header,
/// then its nodes, each `node_size` bytes. No records gives a header with
/// an undefined root.
///
/// The tree is as shallow as the node size allows: a single leaf when the
/// records fit one, otherwise internal nodes above leaves. Records are
/// spread evenly over each node's children, so every node but the root is
/// at least about half full (above libhdf5's merge threshold, which is below
/// half), and each node holds at most its depth's maximum.
pub(crate) fn build_btree_v2(
    p: BTreeV2Params,
    records: &[u8],
    addr: u64,
    offset_size: u8,
    length_size: u8,
) -> Result<Vec<u8>, FormatError> {
    let rs = usize::from(p.record_size);
    if rs == 0 || !records.len().is_multiple_of(rs) {
        return Err(FormatError::SerializationError(format!(
            "B-tree v2 records are {} bytes, not a multiple of the record size {rs}",
            records.len()
        )));
    }
    let n = (records.len() / rs) as u64;
    let hdr_len = header_size(offset_size, length_size);

    // The shallowest depth whose subtree can hold every record.
    let mut info = node_info(p.node_size, p.record_size, offset_size, 0);
    let max_leaf = info[0].max_nrec;
    if max_leaf == 0 || max_leaf > u64::from(u16::MAX) {
        return Err(FormatError::SerializationError(format!(
            "a {}-byte B-tree v2 node holds {max_leaf} {}-byte records; \
             a node holds 1 to 65535",
            p.node_size, p.record_size
        )));
    }
    let mut depth = 0u16;
    while info[usize::from(depth)].cum_max_nrec < n {
        depth += 1;
        if depth > MAX_WRITE_DEPTH {
            return Err(FormatError::SerializationError(format!(
                "{n} records do not fit a B-tree v2 of {}-byte nodes",
                p.node_size
            )));
        }
        info = node_info(p.node_size, p.record_size, offset_size, depth);
        let max = info[usize::from(depth)].max_nrec;
        if max == 0 || max > u64::from(u16::MAX) {
            return Err(FormatError::SerializationError(format!(
                "a {}-byte B-tree v2 internal node holds {max} records; \
                 a node holds 1 to 65535",
                p.node_size
            )));
        }
    }

    let mut w = TreeWriter {
        p,
        records,
        info: &info,
        nrec_width: bytes_for_max_records(max_leaf),
        offset_size,
        first_node: addr + hdr_len as u64,
        nodes: Vec::new(),
    };
    let root = (n > 0)
        .then(|| w.node(depth, 0, saturating_usize(n)))
        .transpose()?;

    let mut out = Vec::with_capacity(hdr_len + w.nodes.len() * p.node_size as usize);
    out.extend_from_slice(b"BTHD");
    out.push(0); // version
    out.push(p.tree_type);
    out.extend_from_slice(&p.node_size.to_le_bytes());
    out.extend_from_slice(&p.record_size.to_le_bytes());
    out.extend_from_slice(&depth.to_le_bytes());
    out.push(p.split_percent);
    out.push(p.merge_percent);
    match root {
        Some(r) => push_uint(&mut out, r.addr, offset_size as usize),
        None => out.extend(core::iter::repeat_n(0xFF, offset_size as usize)),
    }
    let root_nrec = root.map_or(0, |r| r.nrec);
    out.extend_from_slice(&(root_nrec as u16).to_le_bytes());
    push_uint(&mut out, n, length_size as usize);
    let sum = jenkins_lookup3(&out);
    out.extend_from_slice(&sum.to_le_bytes());
    debug_assert_eq!(out.len(), hdr_len);
    for node in &w.nodes {
        out.extend_from_slice(node);
    }
    Ok(out)
}

/// A written node, as its parent points at it.
#[derive(Debug, Clone, Copy)]
struct NodeRef {
    addr: u64,
    /// Records in the node itself.
    nrec: u64,
    /// Records in the subtree it roots.
    all_nrec: u64,
}

struct TreeWriter<'a> {
    p: BTreeV2Params,
    records: &'a [u8],
    info: &'a [NodeInfo],
    /// Width of a child's record count: what a leaf's maximum needs.
    nrec_width: usize,
    offset_size: u8,
    /// Address of the first node (right after the header).
    first_node: u64,
    /// Nodes in file order (children before their parent).
    nodes: Vec<Vec<u8>>,
}

impl TreeWriter<'_> {
    fn record(&self, i: usize) -> &[u8] {
        let rs = usize::from(self.p.record_size);
        &self.records[i * rs..(i + 1) * rs]
    }

    fn push_node(&mut self, mut node: Vec<u8>) -> u64 {
        // The checksum covers the node up to it, not the padding after.
        let sum = jenkins_lookup3(&node);
        node.extend_from_slice(&sum.to_le_bytes());
        debug_assert!(node.len() <= self.p.node_size as usize);
        node.resize(self.p.node_size as usize, 0);
        let addr = self.first_node + self.nodes.len() as u64 * u64::from(self.p.node_size);
        self.nodes.push(node);
        addr
    }

    /// Write the subtree of `depth` holding records `first..first + n`.
    fn node(&mut self, depth: u16, first: usize, n: usize) -> Result<NodeRef, FormatError> {
        let rs = usize::from(self.p.record_size);
        let mut node = Vec::with_capacity(self.p.node_size as usize);
        if depth == 0 {
            debug_assert!(n as u64 <= self.info[0].max_nrec);
            node.extend_from_slice(b"BTLF");
            node.push(0); // version
            node.push(self.p.tree_type);
            node.extend_from_slice(&self.records[first * rs..(first + n) * rs]);
            let addr = self.push_node(node);
            return Ok(NodeRef {
                addr,
                nrec: n as u64,
                all_nrec: n as u64,
            });
        }

        // As few children as hold the records, at least two, with the
        // records spread evenly: `k` children and `k - 1` records between
        // them.
        let below = self.info[usize::from(depth) - 1].cum_max_nrec;
        let k = (n as u64 + 1).div_ceil(below + 1).max(2);
        let max = self.info[usize::from(depth)].max_nrec;
        if k - 1 > max || (n as u64) < k - 1 + k {
            return Err(FormatError::SerializationError(format!(
                "cannot spread {n} B-tree v2 records over {k} children at depth {depth}"
            )));
        }
        let k = saturating_usize(k);
        let in_children = n - (k - 1);
        let (base, extra) = (in_children / k, in_children % k);

        let mut children = Vec::with_capacity(k);
        let mut separators = Vec::with_capacity(k - 1);
        let mut next = first;
        for c in 0..k {
            let m = base + usize::from(c < extra);
            children.push(self.node(depth - 1, next, m)?);
            next += m;
            if c + 1 < k {
                separators.push(next);
                next += 1;
            }
        }
        debug_assert_eq!(next, first + n);

        node.extend_from_slice(b"BTIN");
        node.push(0); // version
        node.push(self.p.tree_type);
        for &s in &separators {
            node.extend_from_slice(self.record(s));
        }
        let total_width = if depth > 1 {
            self.info[usize::from(depth) - 1].cum_max_nrec_size
        } else {
            0
        };
        for c in &children {
            push_uint(&mut node, c.addr, self.offset_size as usize);
            push_uint(&mut node, c.nrec, self.nrec_width);
            if depth > 1 {
                push_uint(&mut node, c.all_nrec, total_width);
            }
        }
        let addr = self.push_node(node);
        Ok(NodeRef {
            addr,
            nrec: (k - 1) as u64,
            all_nrec: n as u64,
        })
    }
}

/// Append `v` as a `width`-byte little-endian integer.
fn push_uint(buf: &mut Vec<u8>, v: u64, width: usize) {
    let bytes = v.to_le_bytes();
    buf.extend_from_slice(&bytes[..width.min(8)]);
    buf.extend(vec![0u8; width.saturating_sub(8)]);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::btree_v2::{BTreeV2Header, collect_btree_v2_records};

    fn params(node_size: u32, record_size: u16) -> BTreeV2Params {
        BTreeV2Params {
            tree_type: 5,
            node_size,
            record_size,
            split_percent: 100,
            merge_percent: 40,
        }
    }

    /// `n` 11-byte records: a big-endian counter, so byte order is key order.
    fn records(n: usize, rs: usize) -> Vec<u8> {
        let mut out = Vec::with_capacity(n * rs);
        for i in 0..n {
            let mut r = vec![0u8; rs];
            r[..8].copy_from_slice(&(i as u64).to_be_bytes());
            out.extend_from_slice(&r);
        }
        out
    }

    fn roundtrip(node_size: u32, rs: u16, n: usize, os: u8, ls: u8) -> BTreeV2Header {
        let recs = records(n, usize::from(rs));
        let base = 4096u64;
        let tree = build_btree_v2(params(node_size, rs), &recs, base, os, ls).unwrap();
        let mut file = vec![0u8; base as usize];
        file.extend_from_slice(&tree);
        let hdr = BTreeV2Header::parse(&file, base as usize, os, ls).unwrap();
        assert_eq!(hdr.total_records, n as u64);
        let got = collect_btree_v2_records(&file, &hdr, os, ls).unwrap();
        assert_eq!(got.len(), n);
        let flat: Vec<u8> = got.into_iter().flat_map(|r| r.data).collect();
        assert_eq!(flat, recs, "node {node_size} rs {rs} n {n}");
        hdr
    }

    #[test]
    fn one_leaf_then_deeper_trees_read_back_in_order() {
        // 512-byte nodes of 11-byte records: 45 per leaf, 1149 at depth 1,
        // 26 449 at depth 2.
        let info = node_info(512, 11, 8, 3);
        assert_eq!(
            info.iter().map(|i| i.cum_max_nrec).collect::<Vec<_>>(),
            [45, 1149, 26_449, 608_349]
        );
        for (n, depth) in [
            (0, 0),
            (1, 0),
            (45, 0),
            (46, 1),
            (1149, 1),
            (1150, 2),
            (26_449, 2),
            (26_450, 3),
            (100_000, 3),
        ] {
            let hdr = roundtrip(512, 11, n, 8, 8);
            assert_eq!(hdr.depth, depth, "{n} records");
        }
    }

    #[test]
    fn pointer_widths_follow_the_offset_and_length_sizes() {
        for (os, ls) in [(4, 4), (8, 4), (4, 8), (2, 2)] {
            roundtrip(512, 11, 5000, os, ls);
        }
        // Wide counts: a leaf of 2048 bytes / 9-byte records (226, one byte)
        // and deeper subtree totals of three bytes.
        roundtrip(2048, 9, 300_000, 8, 8);
    }

    #[test]
    fn every_node_is_within_its_capacity_and_above_the_merge_threshold() {
        let rs = 17u16;
        let n = 70_000usize;
        let info = node_info(512, rs, 8, 3);
        let recs = records(n, usize::from(rs));
        let tree = build_btree_v2(params(512, rs), &recs, 0, 8, 8).unwrap();
        let hdr_len = header_size(8, 8);
        let nodes = (tree.len() - hdr_len) / 512;
        for i in 0..nodes {
            let node = &tree[hdr_len + i * 512..hdr_len + (i + 1) * 512];
            let sig = &node[..4];
            if sig == b"BTLF" {
                continue; // counts checked through the parents below
            }
            assert_eq!(sig, b"BTIN");
        }
        // Walk from the header: each child's count within [40%, 100%].
        let hdr = BTreeV2Header::parse(&tree, 0, 8, 8).unwrap();
        assert_eq!(hdr.depth, 3);
        assert!(u64::from(hdr.num_records_in_root) <= info[3].max_nrec);
        fn walk(tree: &[u8], addr: usize, nrec: usize, depth: usize, info: &[NodeInfo], rs: usize) {
            if depth == 0 {
                return;
            }
            let nrec_w = bytes_for_max_records(info[0].max_nrec);
            let tot_w = if depth > 1 {
                info[depth - 1].cum_max_nrec_size
            } else {
                0
            };
            let mut pos = addr + 6 + nrec * rs;
            for _ in 0..=nrec {
                let a = u64::from_le_bytes(tree[pos..pos + 8].try_into().unwrap()) as usize;
                pos += 8;
                let mut c = 0usize;
                for b in 0..nrec_w {
                    c |= usize::from(tree[pos + b]) << (8 * b);
                }
                pos += nrec_w + tot_w;
                let max = info[depth - 1].max_nrec as usize;
                assert!(c <= max && c * 100 > max * 40, "{c} of {max}");
                walk(tree, a, c, depth - 1, info, rs);
            }
        }
        walk(
            &tree,
            hdr.root_node_address as usize,
            usize::from(hdr.num_records_in_root),
            3,
            &info,
            usize::from(rs),
        );
        assert!(nodes > 0);
    }

    /// Descending to a key range finds exactly the records a full read
    /// holds in it — runs of equal keys that straddle node boundaries
    /// included — at every depth, and nothing for keys not in the tree.
    #[test]
    fn a_key_range_search_matches_a_full_scan() {
        use crate::btree_v2::find_btree_v2_records;
        use core::cmp::Ordering;
        let rs = 11usize;
        // Keys 0, 0, 0, 2, 2, 2, 4, ...: runs of three, odd keys missing.
        for n in [1usize, 45, 46, 1150, 30_000] {
            let mut recs = Vec::with_capacity(n * rs);
            for i in 0..n {
                let mut r = vec![0u8; rs];
                r[..8].copy_from_slice(&((i / 3 * 2) as u64).to_be_bytes());
                r[8..].copy_from_slice(&[(i % 3) as u8, 0, 0]);
                recs.extend_from_slice(&r);
            }
            let base = 4096u64;
            let tree = build_btree_v2(params(512, 11), &recs, base, 8, 8).unwrap();
            let mut file = vec![0u8; base as usize];
            file.extend_from_slice(&tree);
            let hdr = BTreeV2Header::parse(&file, base as usize, 8, 8).unwrap();
            let all = collect_btree_v2_records(&file, &hdr, 8, 8).unwrap();
            let key = |r: &[u8]| u64::from_be_bytes(r[..8].try_into().unwrap());
            let last = key(&all[n - 1].data);
            let probes = (0..=last + 1).step_by(if n > 1000 { 37 } else { 1 });
            for k in probes.chain([last, last + 1, u64::MAX]) {
                let found =
                    find_btree_v2_records(&file, &hdr, 8, &mut |r: &[u8]| key(r).cmp(&k)).unwrap();
                let want: Vec<&[u8]> = all
                    .iter()
                    .map(|r| r.data.as_slice())
                    .filter(|r| key(r) == k)
                    .collect();
                let got: Vec<&[u8]> = found.iter().map(|r| r.data.as_slice()).collect();
                assert_eq!(got, want, "n {n} key {k}");
                assert_eq!(
                    got.len(),
                    if k % 2 == 0 && k <= last {
                        want.len()
                    } else {
                        0
                    }
                );
            }
            // Every record, or none, when the whole tree is in or out of range.
            let every = find_btree_v2_records(&file, &hdr, 8, &mut |_| Ordering::Equal).unwrap();
            assert_eq!(every.len(), n);
            let none = find_btree_v2_records(&file, &hdr, 8, &mut |_| Ordering::Less).unwrap();
            assert!(none.is_empty());
        }
    }

    /// A two-level tree read through a `read_at`-only storage gives what
    /// the slice gives — records, descents and errors — whole, truncated
    /// at every length, and with each byte of its nodes flipped, and each
    /// node costs one read.
    #[test]
    fn storage_reads_match_slice_reads() {
        use crate::btree_v2::{
            collect_btree_v2_records_in, find_btree_v2_records, find_btree_v2_records_in,
        };
        use crate::storage::CountingStorage;
        let (rs, n, base) = (11usize, 120usize, 64usize);
        let recs = records(n, rs);
        let tree = build_btree_v2(params(128, 11), &recs, base as u64, 8, 8).unwrap();
        let mut whole = vec![0u8; base];
        whole.extend_from_slice(&tree);
        let hdr = BTreeV2Header::parse(&whole, base, 8, 8).unwrap();
        assert!(hdr.depth >= 1, "{hdr:?}");
        let key = |r: &[u8]| u64::from_be_bytes(r[..8].try_into().unwrap());
        let mut files = Vec::new();
        for cut in base..=whole.len() {
            files.push(whole[..cut].to_vec());
        }
        for at in base..whole.len() {
            let mut bad = whole.clone();
            bad[at] ^= 0x5a;
            files.push(bad);
        }
        let mut ok = 0;
        for f in &files {
            let st = CountingStorage::new(f.clone());
            let want_h = BTreeV2Header::parse(f, base, 8, 8);
            let got_h = BTreeV2Header::parse_in(&st, base as u64, 8, 8);
            assert_eq!(format!("{got_h:?}"), format!("{want_h:?}"));
            // The nodes of the intact header, over each damaged file.
            let want = collect_btree_v2_records(f, &hdr, 8, 8);
            st.reset();
            let got = collect_btree_v2_records_in(&st, &hdr, 8, 8);
            assert_eq!(format!("{got:?}"), format!("{want:?}"));
            if want.is_ok() {
                ok += 1;
                assert!(st.reads() <= 1 + n as u64 / 3, "{} reads", st.reads());
            }
            for k in [0u64, 7, 60, 119, 500] {
                let want = find_btree_v2_records(f, &hdr, 8, &mut |r: &[u8]| key(r).cmp(&k));
                let got = find_btree_v2_records_in(&st, &hdr, 8, &mut |r: &[u8]| key(r).cmp(&k));
                assert_eq!(format!("{got:?}"), format!("{want:?}"));
            }
        }
        assert!(ok > 1);
    }

    #[test]
    fn a_node_too_small_or_too_big_is_an_error() {
        assert!(build_btree_v2(params(16, 11), &records(1, 11), 0, 8, 8).is_err());
        // A leaf with room for more than 65 535 records.
        assert!(build_btree_v2(params(1 << 20, 11), &records(1, 11), 0, 8, 8).is_err());
        // Records that are not whole.
        assert!(build_btree_v2(params(512, 11), &[0u8; 12], 0, 8, 8).is_err());
    }
}