//! Extensible Array writer: serialize EA layout messages and build EA structures.

#[cfg(not(feature = "std"))]
extern crate alloc;

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

use crate::checksum::jenkins_lookup3;
use crate::chunked_write::{
    WrittenChunk, filtered_chunk_size_len, push_addr, push_index_element, push_v4_chunk_dims,
};

/// Serialize a v4 Extensible Array layout message.
pub(crate) fn serialize_v4_extensible_array(
    chunk_dims: &[u32],
    ea_address: u64,
    offset_size: u8,
    element_size: u32,
) -> Vec<u8> {
    let mut buf = Vec::new();
    buf.push(4); // version
    buf.push(2); // class = chunked
    buf.push(0x00); // flags

    let ndims = chunk_dims.len() as u8 + 1;
    buf.push(ndims);

    push_v4_chunk_dims(&mut buf, chunk_dims, element_size);

    // chunk index type = 4 (Extensible Array)
    buf.push(4);

    // EA creation parameters (must match AEHD and HDF5 C library defaults)
    buf.push(MAX_NELMTS_BITS);
    buf.push(IDX_BLK_ELMTS);
    buf.push(SUP_BLK_MIN_DATA_PTRS);
    buf.push(DATA_BLK_MIN_ELMTS);
    buf.push(MAX_DBLK_PAGE_NELMTS_BITS);

    // EA header address
    match offset_size {
        4 => buf.extend_from_slice(&(ea_address as u32).to_le_bytes()),
        8 => buf.extend_from_slice(&ea_address.to_le_bytes()),
        _ => unreachable!("unexpected offset size: {offset_size}"),
    }

    buf
}

// EA creation parameters — the HDF5 library's defaults for chunk indexes
// (`H5D_EARRAY_*`); the layout message above and the header must agree.
const MAX_NELMTS_BITS: u8 = 32;
const IDX_BLK_ELMTS: u8 = 4;
const SUP_BLK_MIN_DATA_PTRS: u8 = 4;
const DATA_BLK_MIN_ELMTS: u8 = 16;
const MAX_DBLK_PAGE_NELMTS_BITS: u8 = 10;

/// One data block of the array: its first element (relative to the end of
/// the index block's own elements), element count, and address when it is
/// allocated.
struct DataBlock {
    start: usize,
    nelmts: usize,
    addr: Option<u64>,
}

/// Build a complete Extensible Array at a known absolute address.
///
/// `slots[i]` is the element at linear index `i` (see `chunk_grid`); `None`
/// marks an unallocated chunk. The first `IDX_BLK_ELMTS` elements live in
/// the index block, the rest in data blocks grouped by super block level
/// exactly as `H5EA__hdr_init` sizes them: level `u` has `2^(u/2)` data
/// blocks of `DATA_BLK_MIN_ELMTS * 2^ceil(u/2)` elements. The data blocks of
/// the first levels are addressed straight from the index block; later
/// levels go through a super block (EASB). Data blocks larger than a page
/// (`2^MAX_DBLK_PAGE_NELMTS_BITS` elements) are paged, with their page-init
/// bits kept in the owning super block. Only blocks holding a defined element
/// are allocated; the rest keep the undefined address, as in a file the
/// library wrote.
pub fn build_extensible_array_at(
    slots: &[Option<WrittenChunk>],
    offset_size: u8,
    length_size: u8,
    has_filters: bool,
    ea_base_address: u64,
) -> Vec<u8> {
    let os = offset_size as usize;
    let chunk_size_bytes = has_filters.then(|| filtered_chunk_size_len(slots));
    let elem_size = os + chunk_size_bytes.map_or(0, |n| n + 4);
    let client_id: u8 = if has_filters { 1 } else { 0 };
    let arr_off_size = (MAX_NELMTS_BITS as usize).div_ceil(8);
    let page_nelmts = 1usize << MAX_DBLK_PAGE_NELMTS_BITS;
    let idx_blk = IDX_BLK_ELMTS as usize;

    // Elements past the last defined one are never realised
    // (`max_idx_set` is one past the highest index ever set).
    let max_idx_set = slots.iter().rposition(Option::is_some).map_or(0, |i| i + 1);
    let slots = &slots[..max_idx_set];
    let defined_in = |start: usize, n: usize| -> bool {
        let lo = idx_blk.saturating_add(start).min(slots.len());
        let hi = idx_blk
            .saturating_add(start)
            .saturating_add(n)
            .min(slots.len());
        slots[lo..hi].iter().any(Option::is_some)
    };

    // Super block levels: (ndblks, dblk_nelmts, first element).
    let log2_dmin = (DATA_BLK_MIN_ELMTS as u32).trailing_zeros() as usize;
    let nsblks = 1 + MAX_NELMTS_BITS as usize - log2_dmin;
    let ndblk_addrs = 2 * (SUP_BLK_MIN_DATA_PTRS as usize - 1);
    let mut levels: Vec<(usize, usize, usize)> = Vec::with_capacity(nsblks);
    let mut start = 0usize;
    for u in 0..nsblks {
        let ndblks = 1usize << (u / 2);
        let nelmts = (DATA_BLK_MIN_ELMTS as usize) << u.div_ceil(2);
        levels.push((ndblks, nelmts, start));
        // Saturate: on 32-bit targets the last levels only need to compare
        // as "beyond the end".
        start = start.saturating_add(ndblks.saturating_mul(nelmts));
    }
    // Levels whose data blocks the index block addresses directly.
    let mut direct_levels = 0;
    let mut n = 0;
    while n < ndblk_addrs {
        n += levels[direct_levels].0;
        direct_levels += 1;
    }
    let nsblk_addrs = nsblks - direct_levels;

    let dblk_size = |nelmts: usize| -> usize {
        let prefix = 4 + 1 + 1 + os + arr_off_size + 4;
        if nelmts > page_nelmts {
            prefix + (nelmts / page_nelmts) * (page_nelmts * elem_size + 4)
        } else {
            prefix + nelmts * elem_size
        }
    };
    let sblk_bitmap_len = |ndblks: usize, nelmts: usize| -> usize {
        if nelmts > page_nelmts {
            ndblks * (nelmts / page_nelmts).div_ceil(8)
        } else {
            0
        }
    };

    // Plan addresses: header, index block, the direct data blocks, then each
    // allocated super block followed by its allocated data blocks.
    let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * length_size as usize + os + 4;
    let aeib_address = ea_base_address + aehd_size as u64;
    let aeib_size = 4 + 1 + 1 + os + idx_blk * elem_size + ndblk_addrs * os + nsblk_addrs * os + 4;
    let mut cursor = aeib_address + aeib_size as u64;

    let mut ndata_blks = 0u64;
    let mut data_blk_size = 0u64;
    let mut nsuper_blks = 0u64;
    let mut super_blk_size = 0u64;
    let mut realized = idx_blk as u64;

    let mut plan_dblk = |cursor: &mut u64, start: usize, nelmts: usize| -> DataBlock {
        let addr = defined_in(start, nelmts).then(|| {
            let a = *cursor;
            let size = dblk_size(nelmts) as u64;
            *cursor += size;
            ndata_blks += 1;
            data_blk_size += size;
            realized += nelmts as u64;
            a
        });
        DataBlock {
            start,
            nelmts,
            addr,
        }
    };

    let mut direct: Vec<DataBlock> = Vec::with_capacity(ndblk_addrs);
    for &(ndblks, nelmts, first) in &levels[..direct_levels] {
        for k in 0..ndblks {
            direct.push(plan_dblk(&mut cursor, first + k * nelmts, nelmts));
        }
    }
    // (super block address, level, its data blocks)
    let mut supers: Vec<(Option<u64>, usize, Vec<DataBlock>)> = Vec::with_capacity(nsblk_addrs);
    for (u, &(ndblks, nelmts, first)) in levels.iter().enumerate().skip(direct_levels) {
        if !defined_in(first, ndblks.saturating_mul(nelmts)) {
            supers.push((None, u, Vec::new()));
            continue;
        }
        let sb_size =
            4 + 1 + 1 + os + arr_off_size + sblk_bitmap_len(ndblks, nelmts) + ndblks * os + 4;
        let sb_addr = cursor;
        cursor += sb_size as u64;
        nsuper_blks += 1;
        super_blk_size += sb_size as u64;
        let dblks = (0..ndblks)
            .map(|k| plan_dblk(&mut cursor, first + k * nelmts, nelmts))
            .collect();
        supers.push((Some(sb_addr), u, dblks));
    }

    let slot = |i: usize| slots.get(i).and_then(Option::as_ref);
    let write_length = |buf: &mut Vec<u8>, val: u64| match length_size {
        4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
        _ => buf.extend_from_slice(&val.to_le_bytes()),
    };
    let write_addr_opt = |buf: &mut Vec<u8>, addr: Option<u64>| match addr {
        Some(a) => push_addr(buf, a, offset_size),
        None => buf.extend(core::iter::repeat_n(0xFF, os)),
    };
    let block_prefix = |buf: &mut Vec<u8>, sig: &[u8; 4], block_off: usize| {
        buf.extend_from_slice(sig);
        buf.push(0); // version
        buf.push(client_id);
        push_addr(buf, ea_base_address, offset_size);
        buf.extend_from_slice(&(block_off as u64).to_le_bytes()[..arr_off_size]);
    };
    // Serialise one data block (paged or not) onto `out`.
    let write_dblk = |out: &mut Vec<u8>, db: &DataBlock| {
        let at = out.len();
        block_prefix(out, b"EADB", db.start);
        let first = idx_blk + db.start;
        if db.nelmts > page_nelmts {
            // Paged: the prefix carries only its own checksum; each page
            // follows with one of its own.
            let sum = jenkins_lookup3(&out[at..]);
            out.extend_from_slice(&sum.to_le_bytes());
            for p in 0..db.nelmts / page_nelmts {
                let page_at = out.len();
                for e in 0..page_nelmts {
                    let i = first + p * page_nelmts + e;
                    push_index_element(out, slot(i), offset_size, chunk_size_bytes);
                }
                let sum = jenkins_lookup3(&out[page_at..]);
                out.extend_from_slice(&sum.to_le_bytes());
            }
        } else {
            for i in first..first + db.nelmts {
                push_index_element(out, slot(i), offset_size, chunk_size_bytes);
            }
            let sum = jenkins_lookup3(&out[at..]);
            out.extend_from_slice(&sum.to_le_bytes());
        }
        debug_assert_eq!(out.len() - at, dblk_size(db.nelmts));
    };

    // Header (EAHD). The six statistics are, in order: super blocks, their
    // bytes, data blocks, their bytes, max index set, elements realised.
    let mut out = Vec::with_capacity(saturating_usize(cursor - ea_base_address));
    out.extend_from_slice(b"EAHD");
    out.push(0); // version
    out.push(client_id);
    out.push(elem_size as u8);
    out.push(MAX_NELMTS_BITS);
    out.push(IDX_BLK_ELMTS);
    out.push(DATA_BLK_MIN_ELMTS);
    out.push(SUP_BLK_MIN_DATA_PTRS);
    out.push(MAX_DBLK_PAGE_NELMTS_BITS);
    write_length(&mut out, nsuper_blks);
    write_length(&mut out, super_blk_size);
    write_length(&mut out, ndata_blks);
    write_length(&mut out, data_blk_size);
    write_length(&mut out, max_idx_set as u64);
    write_length(&mut out, realized);
    push_addr(&mut out, aeib_address, offset_size);
    let sum = jenkins_lookup3(&out);
    out.extend_from_slice(&sum.to_le_bytes());
    debug_assert_eq!(out.len(), aehd_size);

    // Index block (EAIB): inline elements, data block and super block
    // addresses.
    let ib_start = out.len();
    out.extend_from_slice(b"EAIB");
    out.push(0);
    out.push(client_id);
    push_addr(&mut out, ea_base_address, offset_size);
    for i in 0..idx_blk {
        push_index_element(&mut out, slot(i), offset_size, chunk_size_bytes);
    }
    for db in &direct {
        write_addr_opt(&mut out, db.addr);
    }
    for (sb_addr, _, _) in &supers {
        write_addr_opt(&mut out, *sb_addr);
    }
    let sum = jenkins_lookup3(&out[ib_start..]);
    out.extend_from_slice(&sum.to_le_bytes());
    debug_assert_eq!(out.len() - ib_start, aeib_size);

    for db in direct.iter().filter(|d| d.addr.is_some()) {
        write_dblk(&mut out, db);
    }
    for (sb_addr, u, dblks) in &supers {
        if sb_addr.is_none() {
            continue;
        }
        let (ndblks, nelmts, first) = levels[*u];
        let sb_start = out.len();
        block_prefix(&mut out, b"EASB", first);
        if nelmts > page_nelmts {
            // Page-init bits, `npages` per data block, packed MSB-first
            // (`H5VM_bit_set`): every page of an allocated data block is
            // written.
            let npages = nelmts / page_nelmts;
            let mut bitmap = vec![0u8; sblk_bitmap_len(ndblks, nelmts)];
            for (k, db) in dblks.iter().enumerate() {
                if db.addr.is_some() {
                    for p in 0..npages {
                        let bit = k * npages + p;
                        bitmap[bit / 8] |= 0x80 >> (bit % 8);
                    }
                }
            }
            out.extend_from_slice(&bitmap);
        }
        for db in dblks {
            write_addr_opt(&mut out, db.addr);
        }
        let sum = jenkins_lookup3(&out[sb_start..]);
        out.extend_from_slice(&sum.to_le_bytes());
        for db in dblks.iter().filter(|d| d.addr.is_some()) {
            write_dblk(&mut out, db);
        }
    }
    debug_assert_eq!(out.len() as u64, cursor - ea_base_address);
    out
}