//! Crafted Blosc2 frames and chunks cannot make the decoder allocate out of
//! proportion to the HDF5 chunk it decodes.
//!
//! A frame's header, its offsets chunk and its chunk headers all declare
//! sizes, and the decoder used to allocate what they declared: a 173-byte
//! frame whose offsets chunk claimed 2 GiB was decoded in full before any
//! check failed. Every allocation is now bounded by the output limit (the
//! HDF5 chunk's size) and the input's length.
//!
//! Peak heap use is measured with a counting global allocator; the tests
//! share it, so each holds `SERIAL` for its whole run.
#![cfg(feature = "blosc2")]

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};

use clawhdf5_format::filters_blosc2::{blosc2_decompress, blosc2_decompress_chunk};

struct Counting;

static CURRENT: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);
static SERIAL: Mutex<()> = Mutex::new(());

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(layout) };
        if !p.is_null() {
            let now = CURRENT.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
            PEAK.fetch_max(now, Ordering::Relaxed);
        }
        p
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc_zeroed(layout) };
        if !p.is_null() {
            let now = CURRENT.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
            PEAK.fetch_max(now, Ordering::Relaxed);
        }
        p
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) };
        CURRENT.fetch_sub(layout.size(), Ordering::Relaxed);
    }
}

#[global_allocator]
static ALLOC: Counting = Counting;

/// Bytes allocated at the peak of `f`, above what was live when it started.
fn peak_during<T>(f: impl FnOnce() -> T) -> (T, usize) {
    let base = CURRENT.load(Ordering::Relaxed);
    PEAK.store(base, Ordering::Relaxed);
    let out = f();
    (out, PEAK.load(Ordering::Relaxed).saturating_sub(base))
}

/// What decoding one HDF5 chunk of `limit` bytes from `input` may hold at
/// once: the output, a few blocks of scratch (each no larger than the
/// output), the offsets table, and the Zstandard decoder's state, which has
/// a fixed ceiling: a window of at most 128 KiB (or twice the stream) and a
/// block's table of sequences (up to 98,303 of 12 bytes, 1.2 MB).
fn bound(limit: usize, input: &[u8]) -> usize {
    6 * limit + 2 * input.len() + (2 << 20)
}

fn lock() -> std::sync::MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(|e| e.into_inner())
}

/// A 32-byte (extended) Blosc2 chunk header.
fn chunk_header(ts: u8, nbytes: i32, blocksize: i32, cbytes: i32, special: u8) -> Vec<u8> {
    let mut c = vec![5u8, 1, 0x05, ts];
    for v in [nbytes, blocksize, cbytes] {
        c.extend_from_slice(&v.to_le_bytes());
    }
    c.resize(32, 0);
    c[31] = special << 4;
    c
}

/// A chunk of `nbytes` bytes that repeats one value (special type 3).
fn repeated(value: &[u8], nbytes: i32, blocksize: i32) -> Vec<u8> {
    let mut c = chunk_header(
        value.len() as u8,
        nbytes,
        blocksize,
        32 + value.len() as i32,
        3,
    );
    c.extend_from_slice(value);
    c
}

/// A frame offset recording a special chunk of `kind` (1 zeros, 2 NaN).
fn special_offset(kind: u8) -> [u8; 8] {
    (((0x80 | kind) as i64) << 56).to_le_bytes()
}

/// A B2ND metalayer.
fn nd_meta(shape: &[i64], chunks: &[i32], blocks: &[i32]) -> Vec<u8> {
    let n = shape.len() as u8;
    let mut m = vec![0x95, 0, n, 0x90 | n];
    for s in shape {
        m.push(0xd3);
        m.extend_from_slice(&s.to_be_bytes());
    }
    for dims in [chunks, blocks] {
        m.push(0x90 | n);
        for d in dims {
            m.push(0xd2);
            m.extend_from_slice(&d.to_be_bytes());
        }
    }
    m
}

/// A contiguous frame: header (with a `b2nd` metalayer if given), the data
/// chunks, then the offsets chunk.
fn frame(
    meta: Option<&[u8]>,
    nbytes: i64,
    typesize: i32,
    chunksize: i32,
    data: &[u8],
    offsets: &[u8],
) -> Vec<u8> {
    let mut h = vec![0u8; 91];
    h[0] = 0x9e;
    h[1] = 0xa8;
    h[2..10].copy_from_slice(b"b2frame\0");
    h[25] = 2;
    match meta {
        Some(m) => {
            h.extend_from_slice(&[0xde, 0, 1, 0xa4]);
            h.extend_from_slice(b"b2nd");
            let at = h.len() as i32 + 5;
            h.push(0xd2);
            h.extend_from_slice(&at.to_be_bytes());
            h.push(0xc6);
            h.extend_from_slice(&(m.len() as u32).to_be_bytes());
            h.extend_from_slice(m);
        }
        None => h.extend_from_slice(&[0xde, 0, 0]),
    }
    let header_len = h.len() as i32;
    h[11..15].copy_from_slice(&header_len.to_be_bytes());
    h[30..38].copy_from_slice(&nbytes.to_be_bytes());
    h[39..47].copy_from_slice(&(data.len() as i64).to_be_bytes());
    h[48..52].copy_from_slice(&typesize.to_be_bytes());
    h[58..62].copy_from_slice(&chunksize.to_be_bytes());
    h.extend_from_slice(data);
    h.extend_from_slice(offsets);
    let len = h.len() as u64;
    h[16..24].copy_from_slice(&len.to_be_bytes());
    h
}

/// The frame header's own sizes must not size the offsets chunk: a frame
/// declaring 32 Mi chunks of 4 bytes, whose offsets chunk (40 bytes) says
/// "one repeated offset, 256 MiB of them", made the decoder build all
/// 256 MiB of offsets for a 1 MiB HDF5 chunk and then return 4 bytes.
#[test]
fn offsets_chunk_is_bounded_by_the_output_limit() {
    let _g = lock();
    let limit = 1 << 20;
    let offsets_len: i32 = 256 << 20;
    let nchunks = offsets_len as i64 / 8;
    let offsets = repeated(&special_offset(1), offsets_len, 64 << 20);
    let f = frame(None, nchunks * 4, 4, 4, &[], &offsets);
    let (r, peak) = peak_during(|| blosc2_decompress(&f, limit));
    assert!(r.is_err(), "decoded {:?} bytes", r.map(|v| v.len()));
    assert!(
        peak <= bound(limit, &f),
        "peak {peak} bytes for a {}-byte frame",
        f.len()
    );
    // The same frame with a variable chunk size (0): the offsets chunk
    // alone says how many chunks there are.
    let f = frame(None, nchunks * 4, 4, 0, &[], &offsets);
    let (r, peak) = peak_during(|| blosc2_decompress(&f, limit));
    assert!(r.is_err());
    assert!(peak <= bound(limit, &f), "chunksize 0: peak {peak} bytes");
}

/// A legitimate frame of this shape (one chunk, its offset special) still
/// decodes.
#[test]
fn small_frames_still_decode() {
    let _g = lock();
    let offsets = repeated(&special_offset(1), 8, 8);
    let f = frame(None, 64, 4, 64, &[], &offsets);
    assert_eq!(blosc2_decompress(&f, 64).unwrap(), vec![0; 64]);
    let _ = blosc2_decompress_chunk;
}

/// A chunk that decodes to nothing kept its declared block size (up to
/// 512 MiB) and allocated two scratch blocks of it: about 1 GiB for a
/// 20-byte chunk.
#[test]
fn empty_chunk_does_not_allocate_its_block_size() {
    let _g = lock();
    let mut c = vec![5u8, 1, 0x01, 1];
    for v in [0i32, 0x1FFF_F000, 20] {
        c.extend_from_slice(&v.to_le_bytes());
    }
    c.resize(20, 0);
    let (r, peak) = peak_during(|| blosc2_decompress_chunk(&c, 1 << 20));
    assert_eq!(r.map(|v| v.len()).unwrap_or(0), 0);
    assert!(
        peak <= bound(0, &c),
        "peak {peak} bytes for a 20-byte chunk"
    );
    // Inside a frame for a non-empty HDF5 chunk it is an error, not data.
    let offsets = repeated(&0i64.to_le_bytes(), 8, 8);
    let f = frame(None, 64, 4, 64, &c, &offsets);
    let (r, peak) = peak_during(|| blosc2_decompress(&f, 64));
    assert!(r.is_err(), "decoded {:?}", r.map(|v| v.len()));
    assert!(peak <= bound(64, &f), "in a frame: peak {peak} bytes");
}

/// B2ND chunks were decoded whole, padding included, with up to 16x the
/// HDF5 chunk size as their limit. Blocks are now placed as they are
/// decoded, so the padding is never held.
///
/// Ten dimensions: nine of 3 split into blocks of 2 (padded to 4) and one
/// of 4, so each chunk is 13x the array. One chunk, stored three ways: as a
/// NaN chunk in the frame's offsets, as a repeated-value chunk, and as a
/// chunk of stored (uncompressed) blocks.
#[test]
fn b2nd_padding_is_never_held() {
    let _g = lock();
    let ts = 4usize;
    let mut shape = vec![3i64; 9];
    shape.push(4);
    let chunks: Vec<i32> = shape.iter().map(|&s| s as i32).collect();
    let mut blocks = vec![2i32; 9];
    blocks.push(4);
    let meta = nd_meta(&shape, &chunks, &blocks);
    let items: usize = shape.iter().product::<i64>() as usize;
    let limit = items * ts;
    let block_bytes = ts * blocks.iter().product::<i32>() as usize;
    let ext_bytes = ts * 4usize.pow(9) * 4;
    assert!(ext_bytes > 13 * limit);
    let offsets = |off: [u8; 8]| repeated(&off, 8, 8);

    let value = 1.5f32.to_le_bytes();
    let stored = {
        // Every block stored raw: block k holds the value k.
        let mut c = chunk_header(4, ext_bytes as i32, block_bytes as i32, 0, 0);
        c[2] = 0x02 | 0x10; // memcpyed, not split
        c.truncate(16);
        for k in 0..ext_bytes / block_bytes {
            c.extend((k as f32).to_le_bytes().repeat(block_bytes / 4));
        }
        let n = c.len() as i32;
        c[12..16].copy_from_slice(&n.to_le_bytes());
        c
    };
    let cases: Vec<(&str, Vec<u8>)> = vec![
        (
            "NaN offset",
            frame(
                Some(&meta),
                ext_bytes as i64,
                4,
                ext_bytes as i32,
                &[],
                &offsets(special_offset(2)),
            ),
        ),
        (
            "repeated value",
            frame(
                Some(&meta),
                ext_bytes as i64,
                4,
                ext_bytes as i32,
                &repeated(&value, ext_bytes as i32, block_bytes as i32),
                &offsets(0i64.to_le_bytes()),
            ),
        ),
        (
            "stored blocks",
            frame(
                Some(&meta),
                ext_bytes as i64,
                4,
                ext_bytes as i32,
                &stored,
                &offsets(0i64.to_le_bytes()),
            ),
        ),
    ];
    for (name, f) in cases {
        let (r, peak) = peak_during(|| blosc2_decompress(&f, limit));
        let out = r.unwrap_or_else(|e| panic!("{name}: {e}"));
        assert_eq!(out.len(), limit, "{name}");
        match name {
            "NaN offset" => assert!(
                out.chunks(4)
                    .all(|v| f32::from_le_bytes(v.try_into().unwrap()).is_nan())
            ),
            "repeated value" => assert!(out.chunks(4).all(|v| v == value)),
            _ => {
                // Element (i0..i9) lies in block (i0/2, .., i8/2), numbered
                // in C order over a 2x..x2x1 grid of blocks.
                let mut idx = [0usize; 10];
                for (e, v) in out.chunks(4).enumerate() {
                    let mut n = e;
                    for d in (0..10).rev() {
                        idx[d] = n % shape[d] as usize;
                        n /= shape[d] as usize;
                    }
                    let k = idx[..9].iter().fold(0, |k, &i| k * 2 + i / 2);
                    assert_eq!(
                        f32::from_le_bytes(v.try_into().unwrap()),
                        k as f32,
                        "{name} {e}"
                    );
                }
            }
        }
        assert!(
            peak <= bound(limit, &f),
            "{name}: peak {peak} bytes for a {limit}-byte chunk ({}-byte frame)",
            f.len()
        );
    }
}

/// A B2ND chunk larger than the array (here 16x, the old cap) is refused,
/// or at least never allocated.
#[test]
fn b2nd_chunk_larger_than_the_array_is_not_allocated() {
    let _g = lock();
    let limit = 1 << 20;
    let c = 16 * limit as i32;
    let meta = nd_meta(&[limit as i64], &[c], &[c]);
    let offsets = repeated(&special_offset(1), 8, 8);
    let f = frame(Some(&meta), c as i64, 1, c, &[], &offsets);
    let (r, peak) = peak_during(|| blosc2_decompress(&f, limit));
    assert!(
        peak <= bound(limit, &f),
        "peak {peak} bytes ({:?})",
        r.map(|v| v.len())
    );
}

/// ruzstd reserves a frame's declared window (up to 100 MiB) before it
/// decodes a frame with a decoder it has used before: a Blosc2 chunk of
/// two 16-byte Zstandard streams, each declaring a 96 MiB window,
/// allocated 96 MiB. c-blosc2 compresses each block with its size known,
/// so its windows never exceed the block.
#[test]
fn zstd_window_is_bounded_by_the_output() {
    let _g = lock();
    let mut z = 0xfd2f_b528u32.to_le_bytes().to_vec();
    // No single segment, no checksum; window 2^26 + 4/8 of it = 96 MiB.
    z.extend_from_slice(&[0x00, (16 << 3) | 4]);
    // One raw block, last, of 16 bytes.
    let h = 1 | (16 << 3);
    z.extend_from_slice(&[h as u8, (h >> 8) as u8, 0]);
    z.extend_from_slice(&[7; 16]);
    // Two blocks of 16 bytes, one stream each (not split), Zstandard
    // (codec 4).
    let chunk = |z: &[u8]| {
        let mut c = vec![5u8, 1, 0x10 | (4 << 5), 1];
        for v in [32i32, 16, 0] {
            c.extend_from_slice(&v.to_le_bytes());
        }
        let first = 24 + 4 + z.len();
        c.extend_from_slice(&24i32.to_le_bytes());
        c.extend_from_slice(&(first as i32).to_le_bytes());
        for _ in 0..2 {
            c.extend_from_slice(&(z.len() as i32).to_le_bytes());
            c.extend_from_slice(z);
        }
        let n = c.len() as i32;
        c[12..16].copy_from_slice(&n.to_le_bytes());
        c
    };
    let c = chunk(&z);
    let (r, peak) = peak_during(|| blosc2_decompress_chunk(&c, 32));
    assert!(peak <= bound(32, &c), "peak {peak} bytes ({r:?})");
    assert!(r.is_err(), "{r:?}");
    // The same streams with a window they can use read.
    z[5] = 0;
    assert_eq!(
        blosc2_decompress_chunk(&chunk(&z), 32).unwrap(),
        vec![7; 32]
    );
}