carbonado 0.7.1

use std::io::{Cursor, Read};

#[cfg(feature = "backend-rust")]
use bao_tree::io::{outboard::PostOrderMemOutboard, sync::keyed_valid_ranges};
use bao_tree::{
    BaoTree, ChunkNum, ChunkRanges,
    io::{
        DecodeError,
        outboard::EmptyOutboard,
        sync::{ReadAt, WriteAt, keyed_decode_ranges},
    },
    iter::BaoChunk,
};

use crate::{
    constants::{BAO_BLOCK_SIZE, FEC_M, SLICE_LEN},
    crypto::carbonado_verification_key,
    error::CarbonadoError,
    utils::decode_bao_hash,
};

/// Blake3 chunks (1 KiB each) covered by one 4 KiB Carbonado slice / Bao leaf.
const CHUNKS_PER_SLICE: u64 = 1 << BAO_BLOCK_SIZE.chunk_log();

/// Map a 4 KiB inboard FEC leaf index to `(stripe_index, symbol_index)`.
///
/// Symbols `0..4` are data leaves; `4..8` are parity leaves. Inboard body order is
/// stripe 0's eight leaves, then stripe 1, and so on.
pub fn leaf_index_to_stripe_symbol(leaf_index: u32) -> (u32, u8) {
    (leaf_index / FEC_M as u32, (leaf_index % FEC_M as u32) as u8)
}

/// Inverse of [`leaf_index_to_stripe_symbol`].
pub fn stripe_symbol_to_leaf_index(stripe_index: u32, symbol: u8) -> u32 {
    debug_assert!((symbol as usize) < FEC_M);
    stripe_index * FEC_M as u32 + u32::from(symbol)
}

/// Map a contiguous run of 4 KiB slice indices to keyed-bao [`ChunkRanges`].
pub fn slice_to_chunk_ranges(index: u32, count: u32) -> ChunkRanges {
    let start = ChunkNum(u64::from(index) * CHUNKS_PER_SLICE);
    let end = ChunkNum(u64::from(index + count) * CHUNKS_PER_SLICE);
    ChunkRanges::from(start..end)
}

/// Map bao-tree [`DecodeError`] to [`CarbonadoError`] (shared by full decode and slice verify).
pub(crate) fn map_decode_error(err: DecodeError) -> CarbonadoError {
    match err {
        DecodeError::ParentHashMismatch(_) | DecodeError::LeafHashMismatch(_) => {
            CarbonadoError::AuthenticationFailed
        }
        DecodeError::ParentNotFound(node) => CarbonadoError::BaoResponseTruncated(format!(
            "parent hash pair missing at tree node {:?}",
            node
        )),
        DecodeError::LeafNotFound(chunk) => CarbonadoError::BaoResponseTruncated(format!(
            "leaf data missing at chunk offset {}",
            chunk.to_bytes()
        )),
        DecodeError::Io(e) => CarbonadoError::StdIoError(e),
    }
}

fn map_valid_ranges_read_error(err: std::io::Error) -> CarbonadoError {
    CarbonadoError::OutboardVerificationFailed(format!(
        "bao outboard data read during slice validation: {err}"
    ))
}

#[cfg(feature = "backend-rust")]
fn chunk_count(ranges: &ChunkRanges) -> u64 {
    ranges
        .boundaries()
        .windows(2)
        .map(|w| (w[1] - w[0]).0)
        .sum()
}

/// Returns `(slice_byte_start, slice_byte_end, actual_len)` or [`CarbonadoError::InvalidSliceIndex`].
fn slice_byte_range(
    index: u32,
    count: u32,
    content_len: u64,
) -> Result<(u64, u64, u64), CarbonadoError> {
    let slice_byte_start = u64::from(index) * u64::from(SLICE_LEN);
    if slice_byte_start >= content_len {
        return Err(CarbonadoError::InvalidSliceIndex { index, content_len });
    }
    let slice_byte_len = u64::from(count) * u64::from(SLICE_LEN);
    let slice_byte_end = slice_byte_start
        .saturating_add(slice_byte_len)
        .min(content_len);
    let actual_len = slice_byte_end.saturating_sub(slice_byte_start);
    Ok((slice_byte_start, slice_byte_end, actual_len))
}

/// In-memory [`WriteAt`] target that retains only the requested byte sub-range.
///
/// Used with a full-layout (`ChunkRanges::all()`) keyed decode over inboard responses;
/// discards writes outside the slice window so **retained output** stays O(slice).
/// Peak RSS still includes the caller-owned full inboard body when that blob is resident.
struct SliceRegionWriter {
    region_start: u64,
    region_end: u64,
    buf: Vec<u8>,
}

impl SliceRegionWriter {
    fn for_region(region_start: u64, region_len: u64) -> Self {
        Self {
            region_start,
            region_end: region_start.saturating_add(region_len),
            buf: vec![0u8; region_len as usize],
        }
    }
}

impl WriteAt for SliceRegionWriter {
    fn write_at(&mut self, offset: u64, data: &[u8]) -> std::io::Result<usize> {
        let write_start = offset.max(self.region_start);
        let write_end = offset
            .saturating_add(data.len() as u64)
            .min(self.region_end);
        if write_start >= write_end {
            return Ok(data.len());
        }
        let skip = (write_start - offset) as usize;
        let rel = (write_start - self.region_start) as usize;
        let len = (write_end - write_start) as usize;
        self.buf[rel..rel + len].copy_from_slice(&data[skip..skip + len]);
        Ok(data.len())
    }

    fn write_all_at(&mut self, offset: u64, data: &[u8]) -> std::io::Result<()> {
        self.write_at(offset, data)?;
        Ok(())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

/// Verified read of `count` contiguous 4 KiB slices at `index` from an inboard bao
/// response (`[u64le content_len | response_bytes]`).
///
/// **Retained output:** O(slice) via [`SliceRegionWriter`].
///
/// **Input:** full inboard body (`input: &[u8]`) — not streaming `ReadAt`. Peak RSS is
/// O(body) whenever the caller already holds the blob (same honesty class as W4a C input).
///
/// **Time / I/O:** O(N) over the embedded bao response bytes. Inboard artifacts store a
/// full `ChunkRanges::all()` response; partial keyed decode desyncs the sequential reader,
/// so verification walks the entire encoded stream even when only one slice is requested.
///
/// **`count == 0`:** empty success immediately (no auth) — pure-Rust and dual
/// `lean::verify_slice` short-circuit. Pure Lean C `carbonado_verify_slice` is auth-first.
pub fn verify_slice_inboard_seekable(
    input: &[u8],
    index: u32,
    count: u32,
    hash: &[u8],
    format: u8,
) -> Result<Vec<u8>, CarbonadoError> {
    if count == 0 {
        return Ok(vec![]);
    }
    let content_len = crate::stream::bao::inboard_bao_content_len_prefix(input)?;
    if content_len == 0 {
        return Err(CarbonadoError::InvalidSliceIndex { index, content_len });
    }
    let response = &input[8..];
    let root = decode_bao_hash(hash)?;
    let tree = BaoTree::new(content_len, BAO_BLOCK_SIZE);
    let key = carbonado_verification_key(format);
    // Inboard artifacts embed a full bao response (encode uses ChunkRanges::all()); partial
    // keyed decode ranges only match partial responses, so verify walks the full layout.
    let ranges = ChunkRanges::all();

    let (slice_byte_start, _slice_byte_end, actual_len) =
        slice_byte_range(index, count, content_len)?;

    let mut writer = SliceRegionWriter::for_region(slice_byte_start, actual_len);
    let mut ob = EmptyOutboard { tree, root };
    keyed_decode_ranges(Cursor::new(response), &ranges, &mut writer, &mut ob, &key)
        .map_err(map_decode_error)?;

    Ok(writer.buf)
}

/// Unvalidated inboard slice extraction (kept for Bao-response walks).
#[allow(dead_code)]
///
/// P1-SCRUB: pre-order walk over full inboard response layout is allowed here; must not
/// allocate an O(N) logical buffer (only the requested shard bytes are retained).
///
/// Walks the bao response sequentially (early-stop once the slice window is filled).
/// Does not perform keyed hash checks; RS + re-bao oracle in scrub filters bad candidates.
/// Returns [`CarbonadoError::BaoResponseTruncated`] if the response ends before the slice
/// window is fully populated.
pub(crate) fn extract_slice_inboard_for_scrub(
    input: &[u8],
    index: u32,
    count: u32,
) -> Result<Vec<u8>, CarbonadoError> {
    if count == 0 {
        return Ok(vec![]);
    }
    let content_len = crate::stream::bao::inboard_bao_content_len_prefix(input)?;
    if content_len == 0 {
        return Err(CarbonadoError::InvalidSliceIndex { index, content_len });
    }
    let response = &input[8..];
    let tree = BaoTree::new(content_len, BAO_BLOCK_SIZE);

    let (slice_byte_start, slice_byte_end, actual_len) =
        slice_byte_range(index, count, content_len)?;

    let mut out = vec![0u8; actual_len as usize];
    let mut cursor = Cursor::new(response);
    let mut logical_offset = 0u64;
    let mut filled = 0usize;
    // Inboard blobs always embed a full (ChunkRanges::all()) bao response; walk that layout
    // sequentially but only retain bytes for the requested slice (no full-stream alloc).
    let ranges = ChunkRanges::all();

    for item in tree.ranges_pre_order_chunks_iter_ref(&ranges, 0) {
        match item {
            BaoChunk::Parent { .. } => {
                let mut skip = [0u8; 64];
                cursor
                    .read_exact(&mut skip)
                    .map_err(|e| CarbonadoError::BaoResponseTruncated(e.to_string()))?;
            }
            BaoChunk::Leaf { size, .. } => {
                let mut sz = size as u64;
                let remain = content_len.saturating_sub(logical_offset);
                if sz > remain {
                    sz = remain;
                }
                let leaf_start = logical_offset;
                let leaf_end = logical_offset.saturating_add(sz);
                logical_offset = leaf_end;

                let mut leaf = vec![0u8; sz as usize];
                cursor
                    .read_exact(&mut leaf)
                    .map_err(|e| CarbonadoError::BaoResponseTruncated(e.to_string()))?;

                if leaf_end > slice_byte_start && leaf_start < slice_byte_end {
                    let copy_start = leaf_start.max(slice_byte_start);
                    let copy_end = leaf_end.min(slice_byte_end);
                    let src_off = (copy_start - leaf_start) as usize;
                    let dst_off = (copy_start - slice_byte_start) as usize;
                    let len = (copy_end - copy_start) as usize;
                    out[dst_off..dst_off + len].copy_from_slice(&leaf[src_off..src_off + len]);
                    filled += len;
                }

                if logical_offset >= slice_byte_end {
                    break;
                }
            }
        }
    }

    if filled < actual_len as usize {
        return Err(CarbonadoError::BaoResponseTruncated(format!(
            "scrub slice extract incomplete: got {filled} of {actual_len} bytes at index {index}"
        )));
    }
    Ok(out)
}

/// Verified read of `count` contiguous 4 KiB slices at `index` from bare data plus a
/// post-order outboard sidecar.
///
/// **Memory and time:** O(slice) — validates only the requested chunk ranges via
/// `keyed_valid_ranges`, then reads the corresponding bare bytes.
pub fn verify_slice_outboard<D: ReadAt>(
    data: D,
    outboard_bytes: &[u8],
    data_len: u64,
    index: u32,
    count: u32,
    hash: &[u8],
    format: u8,
) -> Result<Vec<u8>, CarbonadoError> {
    if count == 0 {
        return Ok(vec![]);
    }
    if data_len == 0 {
        return Err(CarbonadoError::InvalidSliceIndex {
            index,
            content_len: data_len,
        });
    }
    let root = decode_bao_hash(hash)?;
    let tree = BaoTree::new(data_len, BAO_BLOCK_SIZE);
    let ob = PostOrderMemOutboard {
        root,
        tree,
        data: outboard_bytes,
    };
    let key = carbonado_verification_key(format);
    let ranges = slice_to_chunk_ranges(index, count);
    // Cap expected chunks at content length (partial last leaf / short files).
    let content_chunks = data_len.div_ceil(1024);
    let expected_chunks = (u64::from(count) * CHUNKS_PER_SLICE)
        .min(content_chunks.saturating_sub(u64::from(index) * CHUNKS_PER_SLICE));

    let mut validated = ChunkRanges::empty();
    for item in keyed_valid_ranges(&ob, &data, &ranges, &key) {
        let range = item.map_err(map_valid_ranges_read_error)?;
        validated |= ChunkRanges::from(range);
    }
    if chunk_count(&validated) < expected_chunks {
        return Err(CarbonadoError::AuthenticationFailed);
    }

    let (slice_byte_start, _slice_byte_end, actual_len) = slice_byte_range(index, count, data_len)?;

    let mut out = vec![0u8; actual_len as usize];
    data.read_exact_at(slice_byte_start, &mut out)
        .map_err(map_valid_ranges_read_error)?;
    Ok(out)
}

/// Byte range of each 4 KiB Bao leaf's payload inside an inboard blob
/// (`[u64le content_len | response]`). Parent hash pairs are not included.
///
/// Used by scrub tests to nick a single leaf without touching Bao parent nodes.
pub fn inboard_leaf_data_ranges(
    input: &[u8],
) -> Result<Vec<std::ops::Range<usize>>, CarbonadoError> {
    let content_len = crate::stream::bao::inboard_bao_content_len_prefix(input)?;
    if content_len == 0 {
        return Ok(vec![]);
    }
    let response = &input[8..];
    let tree = BaoTree::new(content_len, BAO_BLOCK_SIZE);
    let ranges = ChunkRanges::all();
    let mut cursor = 0usize;
    let mut out = Vec::new();
    let mut logical_offset = 0u64;

    for item in tree.ranges_pre_order_chunks_iter_ref(&ranges, 0) {
        match item {
            BaoChunk::Parent { .. } => {
                cursor = cursor.saturating_add(64);
                if cursor > response.len() {
                    return Err(CarbonadoError::BaoResponseTruncated(
                        "inboard leaf-range walk: parent pair past end of response".to_string(),
                    ));
                }
            }
            BaoChunk::Leaf { size, .. } => {
                let mut sz = size as u64;
                let remain = content_len.saturating_sub(logical_offset);
                if sz > remain {
                    sz = remain;
                }
                let start = 8 + cursor;
                let end = start.saturating_add(sz as usize);
                if end > input.len() {
                    return Err(CarbonadoError::BaoResponseTruncated(format!(
                        "inboard leaf-range walk: leaf bytes {start}..{end} past encoded len {}",
                        input.len()
                    )));
                }
                out.push(start..end);
                cursor += sz as usize;
                logical_offset = logical_offset.saturating_add(sz);
            }
        }
    }
    Ok(out)
}

/// Bao-verify each 4 KiB leaf of an inboard blob in one pre-order walk.
///
/// `Some(bytes)` is a leaf that matches its expected keyed hash. `None` is an
/// erasure (leaf hash mismatch, truncated leaf, or unauthenticated parent).
/// Scrub treats `None` as an RS erasure in that stripe.
pub fn classify_inboard_leaves(
    input: &[u8],
    hash: &[u8],
    format: u8,
) -> Result<Vec<Option<Vec<u8>>>, CarbonadoError> {
    let content_len = crate::stream::bao::inboard_bao_content_len_prefix(input)?;
    if content_len == 0 {
        return Ok(vec![]);
    }
    let n_leaves = content_len.div_ceil(u64::from(SLICE_LEN)) as usize;
    let mut leaves = vec![None; n_leaves];
    let root = decode_bao_hash(hash)?;
    let key = carbonado_verification_key(format);
    let tree = BaoTree::new(content_len, BAO_BLOCK_SIZE);
    let response = &input[8..];
    let mut cursor = Cursor::new(response);
    let mut stack = vec![blake3::Hash::from(*root.as_bytes())];
    let ranges = ChunkRanges::all();

    for item in tree.ranges_pre_order_chunks_iter_ref(&ranges, 0) {
        match item {
            BaoChunk::Parent { left, right, .. } => {
                let mut pair = [0u8; 64];
                if cursor.read_exact(&mut pair).is_err() {
                    break;
                }
                let l_hash =
                    blake3::Hash::from(<[u8; 32]>::try_from(&pair[..32]).map_err(|_| {
                        CarbonadoError::BaoResponseTruncated(
                            "inboard parent pair: left hash".to_string(),
                        )
                    })?);
                let r_hash =
                    blake3::Hash::from(<[u8; 32]>::try_from(&pair[32..]).map_err(|_| {
                        CarbonadoError::BaoResponseTruncated(
                            "inboard parent pair: right hash".to_string(),
                        )
                    })?);
                let _expected = stack.pop();
                // Continue with the on-disk pair so later leaves still classify.
                if right {
                    stack.push(r_hash);
                }
                if left {
                    stack.push(l_hash);
                }
            }
            BaoChunk::Leaf {
                size,
                is_root,
                start_chunk,
                ..
            } => {
                let mut buf = vec![0u8; size];
                if cursor.read_exact(&mut buf).is_err() {
                    break;
                }
                let remain = content_len.saturating_sub(start_chunk.to_bytes());
                if (buf.len() as u64) > remain {
                    buf.truncate(remain as usize);
                }
                let actual = bao_tree::keyed_hash_subtree(start_chunk.0, &buf, is_root, &key);
                let expected = stack.pop();
                let leaf_index = (start_chunk.0 / CHUNKS_PER_SLICE) as usize;
                if leaf_index < leaves.len() && expected == Some(actual) {
                    leaves[leaf_index] = Some(buf);
                }
            }
        }
    }
    Ok(leaves)
}

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

    #[test]
    fn leaf_index_maps_to_stripe_and_symbol() {
        assert_eq!(leaf_index_to_stripe_symbol(0), (0, 0));
        assert_eq!(leaf_index_to_stripe_symbol(3), (0, 3));
        assert_eq!(leaf_index_to_stripe_symbol(4), (0, 4));
        assert_eq!(leaf_index_to_stripe_symbol(7), (0, 7));
        assert_eq!(leaf_index_to_stripe_symbol(8), (1, 0));
        assert_eq!(leaf_index_to_stripe_symbol(15), (1, 7));
        for stripe in 0u32..5 {
            for symbol in 0u8..FEC_M as u8 {
                let leaf = stripe_symbol_to_leaf_index(stripe, symbol);
                assert_eq!(leaf_index_to_stripe_symbol(leaf), (stripe, symbol));
            }
        }
    }
}