//! HDF5 Global Heap collection parsing.

#[cfg(not(feature = "std"))]
use alloc::{borrow::Cow, format, string::String, vec::Vec};
#[cfg(feature = "std")]
use std::borrow::Cow;

use crate::error::FormatError;
use crate::storage::{Storage, len_usize, read_exact_at};

/// Magic signature for global heap collections.
const GCOL_SIGNATURE: [u8; 4] = *b"GCOL";

/// A parsed global heap collection.
#[derive(Debug, Clone)]
pub struct GlobalHeapCollection {
    /// Total size of this collection including header.
    pub collection_size: u64,
    /// Objects within this collection.
    pub objects: Vec<GlobalHeapObject>,
}

/// A single object within a global heap collection.
#[derive(Debug, Clone)]
pub struct GlobalHeapObject {
    /// Object index (1-based; 0 is the free space marker).
    pub index: u16,
    /// Reference count.
    pub reference_count: u16,
    /// Object data.
    pub data: Vec<u8>,
}

/// Checks that `[offset, offset + needed)` ends by `data_len`.
fn ensure_len(data_len: usize, offset: usize, needed: usize) -> Result<(), FormatError> {
    match offset.checked_add(needed) {
        Some(end) if end <= data_len => Ok(()),
        _ => Err(FormatError::UnexpectedEof {
            expected: offset.saturating_add(needed),
            available: data_len,
        }),
    }
}

fn read_length(data: &[u8], offset: usize, length_size: u8) -> Result<u64, FormatError> {
    let s = length_size as usize;
    ensure_len(data.len(), offset, s)?;
    let slice = &data[offset..offset + s];
    Ok(match length_size {
        2 => u16::from_le_bytes([slice[0], slice[1]]) as u64,
        4 => u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) as u64,
        8 => u64::from_le_bytes([
            slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
        ]),
        _ => return Err(FormatError::InvalidLengthSize(length_size)),
    })
}

fn object_overrun_msg(index: u16, size: usize, collection_size: u64) -> String {
    format!(
        "global heap object {index} ({size} bytes) runs past the end of its \
         {collection_size}-byte collection"
    )
}

/// Round up to next multiple of 8.
fn pad8(x: usize) -> usize {
    (x + 7) & !7
}

/// Where one object of a global heap collection lies in the file, without
/// its data: see [`GlobalHeapCollection::parse_index`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GlobalHeapObjectRef {
    /// Object index (1-based; 0 is the free space marker).
    pub index: u16,
    /// Reference count.
    pub reference_count: u16,
    /// Offset of the object's data in the file data the collection was
    /// parsed from.
    pub offset: usize,
    /// Size of the object's data in bytes.
    pub size: usize,
}

/// A global heap collection's objects, located but not copied.
#[derive(Debug, Clone)]
pub struct GlobalHeapIndex {
    /// Total size of this collection including header.
    pub collection_size: u64,
    /// The objects, in file order.
    pub objects: Vec<GlobalHeapObjectRef>,
}

impl GlobalHeapCollection {
    /// Parse a global heap collection at the given offset in the file data.
    pub fn parse(
        file_data: &[u8],
        offset: usize,
        length_size: u8,
    ) -> Result<GlobalHeapCollection, FormatError> {
        Self::parse_in(&file_data, offset as u64, length_size)
    }

    /// [`Self::parse`] over any [`Storage`]: one read of the header, one of
    /// the collection.
    pub fn parse_in(
        file: &dyn Storage,
        offset: u64,
        length_size: u8,
    ) -> Result<GlobalHeapCollection, FormatError> {
        let (bytes, base, index) = Self::read_collection(file, offset, length_size)?;
        Ok(GlobalHeapCollection {
            collection_size: index.collection_size,
            objects: index
                .objects
                .iter()
                .map(|o| GlobalHeapObject {
                    index: o.index,
                    reference_count: o.reference_count,
                    data: bytes[o.offset - base..o.offset - base + o.size].to_vec(),
                })
                .collect(),
        })
    }

    /// Locate the objects of the global heap collection at `offset` without
    /// copying their data, so a caller can keep many collections indexed
    /// for the cost of their object headers.
    ///
    /// The collection must lie inside `file_data`, and every object inside
    /// the collection, as libhdf5 lays them out; an object that runs past
    /// its collection is an error.
    pub fn parse_index(
        file_data: &[u8],
        offset: usize,
        length_size: u8,
    ) -> Result<GlobalHeapIndex, FormatError> {
        Self::parse_index_in(&file_data, offset as u64, length_size)
    }

    /// [`Self::parse_index`] over any [`Storage`]: one read of the header,
    /// one of the collection. The object offsets are file offsets.
    pub fn parse_index_in(
        file: &dyn Storage,
        offset: u64,
        length_size: u8,
    ) -> Result<GlobalHeapIndex, FormatError> {
        Ok(Self::read_collection(file, offset, length_size)?.2)
    }

    /// Read the collection at `offset` and index its objects: the
    /// collection's bytes, its offset as a `usize`, and the index (with
    /// file offsets).
    fn read_collection(
        file: &dyn Storage,
        offset: u64,
        length_size: u8,
    ) -> Result<(Cow<'_, [u8]>, usize, GlobalHeapIndex), FormatError> {
        let file_len = len_usize(file);
        // signature(4) + version(1) + reserved(3) + collection_size(length_size),
        // padded to a multiple of 8 as libhdf5 lays it out (`H5HG_SIZEOF_HDR`).
        // With 8-byte lengths the padding is 0; with 4-byte lengths it is 4,
        // and reading without it put every object 4 bytes early.
        let header_size = pad8(8 + length_size as usize);
        let header = read_exact_at(file, offset, header_size)?;
        let offset = usize::try_from(offset).map_err(|_| FormatError::UnexpectedEof {
            expected: usize::MAX,
            available: file_len,
        })?;

        if header[..4] != GCOL_SIGNATURE {
            return Err(FormatError::InvalidGlobalHeapSignature);
        }

        let version = header[4];
        if version != 1 {
            return Err(FormatError::InvalidGlobalHeapVersion(version));
        }

        let collection_size = read_length(&header, 8, length_size)?;
        let collection_end = usize::try_from(collection_size)
            .ok()
            .and_then(|size| offset.checked_add(size))
            .ok_or(FormatError::UnexpectedEof {
                expected: usize::MAX,
                available: file_len,
            })?;
        if collection_end > file_len {
            return Err(FormatError::UnexpectedEof {
                expected: collection_end,
                available: file_len,
            });
        }
        let collection = read_exact_at(file, offset as u64, collection_end - offset)?;
        // Positions below are file offsets; `file_data(p)` is the byte at `p`.
        let file_data = |p: usize| collection[p - offset];

        let mut pos = offset + header_size;
        let mut objects = Vec::new();

        // Parse objects until we hit index 0 (free space) or run out of space
        while pos + 2 <= collection_end {
            let object_index = u16::from_le_bytes([file_data(pos), file_data(pos + 1)]);

            if object_index == 0 {
                // Free space marker — done
                break;
            }

            // object_index(2) + reference_count(2) + reserved(4) +
            // object_size(length_size), padded to 8 (`H5HG_SIZEOF_OBJHDR`).
            let obj_header_size = pad8(8 + length_size as usize);
            ensure_len(collection_end, pos, obj_header_size)?;

            let reference_count = u16::from_le_bytes([file_data(pos + 2), file_data(pos + 3)]);
            let object_size =
                usize::try_from(read_length(&collection[pos - offset..], 8, length_size)?)
                    .map_err(|_| FormatError::Overflow("global heap object size".into()))?;

            pos += obj_header_size;
            if pos
                .checked_add(object_size)
                .is_none_or(|end| end > collection_end)
            {
                return Err(FormatError::VlDataError(object_overrun_msg(
                    object_index,
                    object_size,
                    collection_size,
                )));
            }

            objects.push(GlobalHeapObjectRef {
                index: object_index,
                reference_count,
                offset: pos,
                size: object_size,
            });

            // Advance past data + padding to 8-byte boundary
            pos = pos.saturating_add(pad8(object_size));
        }

        let index = GlobalHeapIndex {
            collection_size,
            objects,
        };
        Ok((collection, offset, index))
    }

    /// Get an object by its index.
    pub fn get_object(&self, index: u16) -> Option<&GlobalHeapObject> {
        self.objects.iter().find(|o| o.index == index)
    }
}

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

    /// Build a global heap collection with given objects.
    fn build_collection(
        objects: &[(u16, u16, &[u8])], // (index, ref_count, data)
        length_size: u8,
    ) -> Vec<u8> {
        let ls = length_size as usize;

        // Calculate total size
        // libhdf5 pads both headers to a multiple of 8.
        let header_size = pad8(8 + ls);
        let mut obj_size_total = 0usize;
        for (_, _, data) in objects {
            let obj_header = pad8(8 + ls);
            obj_size_total += obj_header + pad8(<[u8]>::len(data));
        }
        // Free space marker (2 bytes for index 0)
        obj_size_total += 2;
        let collection_size = header_size + obj_size_total;

        let mut buf = Vec::new();
        buf.extend_from_slice(&GCOL_SIGNATURE);
        buf.push(1); // version
        buf.extend_from_slice(&[0u8; 3]); // reserved

        // collection_size
        match length_size {
            4 => buf.extend_from_slice(&(collection_size as u32).to_le_bytes()),
            8 => buf.extend_from_slice(&(collection_size as u64).to_le_bytes()),
            _ => panic!("unsupported length_size"),
        }
        buf.resize(header_size, 0);

        // Objects
        for (index, ref_count, data) in objects {
            buf.extend_from_slice(&index.to_le_bytes());
            buf.extend_from_slice(&ref_count.to_le_bytes());
            buf.extend_from_slice(&[0u8; 4]); // reserved
            match length_size {
                4 => buf.extend_from_slice(&(data.len() as u32).to_le_bytes()),
                8 => buf.extend_from_slice(&(data.len() as u64).to_le_bytes()),
                _ => panic!("unsupported"),
            }
            buf.resize(buf.len() + (pad8(8 + ls) - (8 + ls)), 0);
            buf.extend_from_slice(data);
            // Pad to 8 bytes
            let padded = pad8(<[u8]>::len(data));
            buf.resize(buf.len() + (padded - <[u8]>::len(data)), 0);
        }

        // Free space marker
        buf.extend_from_slice(&0u16.to_le_bytes());

        buf
    }

    #[test]
    fn parse_collection_two_objects() {
        let data = build_collection(&[(1, 1, b"hello"), (2, 1, b"world!!!")], 8);
        let coll = GlobalHeapCollection::parse(&data, 0, 8).unwrap();
        assert_eq!(coll.objects.len(), 2);
        assert_eq!(coll.objects[0].index, 1);
        assert_eq!(coll.objects[0].data, b"hello");
        assert_eq!(coll.objects[1].index, 2);
        assert_eq!(coll.objects[1].data, b"world!!!");
    }

    #[test]
    fn get_object_by_index() {
        let data = build_collection(&[(1, 1, b"aaa"), (3, 2, b"bbb")], 8);
        let coll = GlobalHeapCollection::parse(&data, 0, 8).unwrap();
        let obj = coll.get_object(3).unwrap();
        assert_eq!(obj.data, b"bbb");
        assert_eq!(obj.reference_count, 2);
        assert!(coll.get_object(99).is_none());
    }

    #[test]
    fn free_space_terminates_parsing() {
        // Build collection with free space marker immediately
        let mut data = Vec::new();
        data.extend_from_slice(&GCOL_SIGNATURE);
        data.push(1);
        data.extend_from_slice(&[0u8; 3]);
        let size = 8u64 + 8 + 2; // header + length_size + free space marker
        data.extend_from_slice(&size.to_le_bytes());
        data.extend_from_slice(&0u16.to_le_bytes()); // free space

        let coll = GlobalHeapCollection::parse(&data, 0, 8).unwrap();
        assert_eq!(coll.objects.len(), 0);
    }

    #[test]
    fn invalid_signature_error() {
        let mut data = build_collection(&[(1, 1, b"x")], 8);
        data[0] = b'X'; // corrupt
        let err = GlobalHeapCollection::parse(&data, 0, 8).unwrap_err();
        assert_eq!(err, FormatError::InvalidGlobalHeapSignature);
    }

    #[test]
    fn invalid_version_error() {
        let mut data = build_collection(&[(1, 1, b"x")], 8);
        data[4] = 2; // wrong version
        let err = GlobalHeapCollection::parse(&data, 0, 8).unwrap_err();
        assert_eq!(err, FormatError::InvalidGlobalHeapVersion(2));
    }

    #[test]
    fn parse_with_4byte_length() {
        let data = build_collection(&[(1, 1, b"test")], 4);
        let coll = GlobalHeapCollection::parse(&data, 0, 4).unwrap();
        assert_eq!(coll.objects.len(), 1);
        assert_eq!(coll.objects[0].data, b"test");
    }

    /// Collections, and every truncation of them, index and parse
    /// identically through a `read_at`-only storage: two reads each.
    #[test]
    fn storage_parse_matches_slice_parse() {
        use crate::storage::CountingStorage;
        let objs: &[(u16, u16, &[u8])] = &[(1, 1, b"hello"), (2, 3, b"a longer object")];
        for ls in [4u8, 8] {
            let coll = build_collection(objs, ls);
            let mut corrupt = coll.clone();
            corrupt[8] = 200; // collection size past the end of the file
            let mut overrun = coll.clone();
            let size_at = pad8(8 + ls as usize) + 8;
            overrun[size_at] = 250; // first object runs past the collection
            for full in [coll, corrupt, overrun] {
                for at in [0usize, 5] {
                    for cut in 0..=full.len() {
                        let mut f = vec![0u8; at];
                        f.extend_from_slice(&full[..cut]);
                        let storage = CountingStorage::new(f.clone());
                        let want = GlobalHeapCollection::parse(&f, at, ls);
                        let got = GlobalHeapCollection::parse_in(&storage, at as u64, ls);
                        assert_eq!(format!("{got:?}"), format!("{want:?}"));
                        let want = GlobalHeapCollection::parse_index(&f, at, ls);
                        let got = GlobalHeapCollection::parse_index_in(&storage, at as u64, ls);
                        assert_eq!(format!("{got:?}"), format!("{want:?}"));
                    }
                }
            }
        }
        let storage = CountingStorage::new(build_collection(objs, 8));
        assert_eq!(
            GlobalHeapCollection::parse_in(&storage, 0, 8)
                .unwrap()
                .objects
                .len(),
            2
        );
        assert_eq!(storage.reads(), 2);
    }
}