//! HDF5 Local Heap parsing.

#[cfg(not(feature = "std"))]
use alloc::string::String;

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

/// Parsed HDF5 Local Heap header.
#[derive(Debug, Clone)]
pub struct LocalHeap {
    /// Size of the data segment in bytes.
    pub data_segment_size: u64,
    /// Offset of the free list head within the data segment.
    pub free_list_head_offset: u64,
    /// File address of the data segment.
    pub data_segment_address: u64,
}

fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
    let s = size as usize;
    if pos.checked_add(s).is_none_or(|end| end > data.len()) {
        return Err(FormatError::UnexpectedEof {
            expected: pos.saturating_add(s),
            available: data.len(),
        });
    }
    let slice = &data[pos..pos + s];
    Ok(match 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::InvalidOffsetSize(size)),
    })
}

/// First read of a name on a backend without the file in memory: most link
/// names are shorter than this.
const NAME_READ_START: usize = 64;

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

    /// [`Self::parse`] over any [`Storage`]: one read of the header.
    pub fn parse_in<S: Storage + ?Sized>(
        file: &S,
        offset: u64,
        offset_size: u8,
        length_size: u8,
    ) -> Result<LocalHeap, FormatError> {
        // signature(4) + version(1) + reserved(3) = 8, then length_size*2 + offset_size
        let ls = length_size as usize;
        let os = offset_size as usize;
        let total = 8 + ls * 2 + os;
        let header = read_exact_at(file, offset, total)?;
        let file_data: &[u8] = &header;
        let offset = 0usize;

        if &file_data[offset..offset + 4] != b"HEAP" {
            return Err(FormatError::InvalidLocalHeapSignature);
        }

        let version = file_data[offset + 4];
        if version != 0 {
            return Err(FormatError::InvalidLocalHeapVersion(version));
        }

        let mut pos = offset + 8;
        let data_segment_size = read_offset(file_data, pos, length_size)?;
        pos += ls;
        let free_list_head_offset = read_offset(file_data, pos, length_size)?;
        pos += ls;
        let data_segment_address = read_offset(file_data, pos, offset_size)?;

        Ok(LocalHeap {
            data_segment_size,
            free_list_head_offset,
            data_segment_address,
        })
    }

    /// Walk the free list the way libhdf5 does when it loads a heap's data
    /// (`H5HL__fl_deserialize`), rejecting a heap whose free list points
    /// outside the data segment. libhdf5 refuses such a heap ("bad heap free
    /// list"), and names read from it would be garbage.
    ///
    /// libhdf5 only loads a heap when it needs a name from it (an empty
    /// group's broken heap goes unnoticed), so call this before the first
    /// [`Self::read_string`], not on parse.
    ///
    /// The end of the list is `H5HL_FREE_NULL` (1); an all-ones value (the
    /// undefined address) is accepted as "no free list" too.
    pub fn validate_free_list(&self, file_data: &[u8], length_size: u8) -> Result<(), FormatError> {
        self.validate_free_list_in(file_data, length_size)
    }

    /// [`Self::validate_free_list`] over any [`Storage`]: two small reads
    /// per free block.
    pub fn validate_free_list_in<S: Storage + ?Sized>(
        &self,
        file: &S,
        length_size: u8,
    ) -> Result<(), FormatError> {
        const FREE_NULL: u64 = 1;
        let ls = length_size as usize;
        let undefined = if ls >= 8 {
            u64::MAX
        } else {
            (1u64 << (8 * ls)) - 1
        };
        let size = self.data_segment_size;
        let seg = self.data_segment_address;
        let mut next = self.free_list_head_offset;
        // Each free block holds two lengths, so a list longer than this
        // revisits a block: a cycle.
        let max_blocks = size / (2 * ls as u64) + 1;
        let mut walked = 0u64;
        while next != FREE_NULL && next != undefined {
            if next >= size || walked >= max_blocks {
                return Err(FormatError::InvalidLocalHeapFreeList);
            }
            walked += 1;
            let at = seg
                .checked_add(next)
                .and_then(|a| usize::try_from(a).ok())
                .ok_or(FormatError::InvalidLocalHeapFreeList)?;
            let block_offset = next;
            next = read_offset(&read_exact_at(file, at as u64, ls)?, 0, length_size)?;
            if next == 0 {
                return Err(FormatError::InvalidLocalHeapFreeList);
            }
            let block_size =
                read_offset(&read_exact_at(file, (at + ls) as u64, ls)?, 0, length_size)?;
            if block_offset
                .checked_add(block_size)
                .is_none_or(|end| end > size)
            {
                return Err(FormatError::InvalidLocalHeapFreeList);
            }
        }
        Ok(())
    }

    /// Read a null-terminated string from the heap's data segment at the given byte offset.
    pub fn read_string(&self, file_data: &[u8], string_offset: u64) -> Result<String, FormatError> {
        self.read_string_in(file_data, string_offset)
    }

    /// [`Self::read_string`] over any [`Storage`]: one read of up to 64
    /// bytes for a short name, more (each four times the last) up to the end
    /// of the data segment for a longer one.
    pub fn read_string_in<S: Storage + ?Sized>(
        &self,
        file: &S,
        string_offset: u64,
    ) -> Result<String, FormatError> {
        let file_len = len_usize(file);
        let seg_addr = to_usize(self.data_segment_address)?;
        let str_start =
            seg_addr
                .checked_add(to_usize(string_offset)?)
                .ok_or(FormatError::Overflow(
                    "local heap seg_addr + string_offset overflow".into(),
                ))?;
        let seg_end = seg_addr
            .checked_add(to_usize(self.data_segment_size)?)
            .ok_or(FormatError::Overflow(
                "local heap seg_addr + data_segment_size overflow".into(),
            ))?;

        if str_start >= file_len || str_start >= seg_end {
            return Err(FormatError::UnexpectedEof {
                expected: str_start + 1,
                available: file_len,
            });
        }

        // Find the null terminator, which lies before the end of the data
        // segment (or of the file). In memory that is one borrowed slice;
        // otherwise the bytes are read in growing pieces, so a name costs a
        // read of about its own length, not of the rest of the segment
        // (whose size is an untrusted header field).
        let search_end = seg_end.min(file_len);
        let total = search_end - str_start;
        let mut want = if file.as_contiguous().is_some() {
            total
        } else {
            total.min(NAME_READ_START)
        };
        loop {
            let rest = read_exact_at(file, str_start as u64, want)?;
            if let Some(len) = rest.iter().position(|&b| b == 0) {
                let s = core::str::from_utf8(&rest[..len])
                    .map_err(|_| FormatError::InvalidLocalHeapSignature)?;
                return Ok(String::from(s));
            }
            if want == total {
                return Err(FormatError::UnexpectedEof {
                    expected: search_end + 1,
                    available: search_end,
                });
            }
            want = want.saturating_mul(4).min(total);
        }
    }
}

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

    fn build_heap_file(
        heap_offset: usize,
        data_seg_offset: usize,
        strings: &[&str],
        offset_size: u8,
        length_size: u8,
    ) -> Vec<u8> {
        // Build data segment
        let mut data_seg = Vec::new();
        for s in strings {
            data_seg.extend_from_slice(s.as_bytes());
            data_seg.push(0); // null terminator
        }
        let data_seg_size = data_seg.len();

        let total_size = data_seg_offset + data_seg_size + 64;
        let mut file = vec![0u8; total_size];

        // Write heap header at heap_offset
        let mut pos = heap_offset;
        file[pos..pos + 4].copy_from_slice(b"HEAP");
        pos += 4;
        file[pos] = 0; // version
        pos += 1;
        // reserved 3
        pos += 3;
        // data_segment_size
        write_val(&mut file, pos, data_seg_size as u64, length_size);
        pos += length_size as usize;
        // free_list_head_offset: H5HL_FREE_NULL (no free space)
        write_val(&mut file, pos, 1, length_size);
        pos += length_size as usize;
        // data_segment_address
        write_val(&mut file, pos, data_seg_offset as u64, offset_size);

        // Write data segment
        file[data_seg_offset..data_seg_offset + data_seg_size].copy_from_slice(&data_seg);

        file
    }

    fn write_val(buf: &mut [u8], pos: usize, val: u64, size: u8) {
        match size {
            4 => buf[pos..pos + 4].copy_from_slice(&(val as u32).to_le_bytes()),
            8 => buf[pos..pos + 8].copy_from_slice(&val.to_le_bytes()),
            _ => panic!("test"),
        }
    }

    #[test]
    fn parse_rejects_near_usize_max_offset_without_panicking() {
        // Found by fuzzing: `offset + total` overflowed for a crafted
        // near-usize::MAX offset.
        let file = build_heap_file(0, 100, &["hello"], 8, 8);
        let result = LocalHeap::parse(&file, usize::MAX - 4, 8, 8);
        assert!(
            matches!(result, Err(FormatError::UnexpectedEof { .. })),
            "expected a clean UnexpectedEof, got {result:?}"
        );
    }

    #[test]
    fn parse_heap_header() {
        let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        assert_eq!(heap.data_segment_address, 100);
        assert_eq!(heap.data_segment_size, 12); // "hello\0world\0"
    }

    #[test]
    fn read_string_at_offset_0() {
        let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        let s = heap.read_string(&file, 0).unwrap();
        assert_eq!(s, "hello");
    }

    #[test]
    fn read_string_at_offset_6() {
        let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        let s = heap.read_string(&file, 6).unwrap();
        assert_eq!(s, "world");
    }

    #[test]
    fn invalid_signature() {
        let mut file = build_heap_file(0, 100, &["x"], 8, 8);
        file[0] = b'X';
        let err = LocalHeap::parse(&file, 0, 8, 8).unwrap_err();
        assert_eq!(err, FormatError::InvalidLocalHeapSignature);
    }

    #[test]
    fn read_string_past_segment() {
        let file = build_heap_file(0, 100, &["hi"], 8, 8);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        let err = heap.read_string(&file, 100).unwrap_err();
        assert!(matches!(err, FormatError::UnexpectedEof { .. }));
    }

    #[test]
    fn parse_heap_4byte_offsets() {
        let file = build_heap_file(0, 80, &["test"], 4, 4);
        let heap = LocalHeap::parse(&file, 0, 4, 4).unwrap();
        assert_eq!(heap.data_segment_address, 80);
        let s = heap.read_string(&file, 0).unwrap();
        assert_eq!(s, "test");
    }

    /// Heap with data segment `[a, b, c, 0-padding]` whose free list starts
    /// at `head` and has one block `(next, size)` at offset 8.
    fn heap_with_free_block(head: u64, next: u64, size: u64) -> Vec<u8> {
        let mut file = build_heap_file(0, 100, &["abcdefg"], 8, 8);
        file.resize(200, 0);
        write_val(&mut file, 8, 32, 8); // data segment size
        write_val(&mut file, 16, head, 8);
        write_val(&mut file, 108, next, 8);
        write_val(&mut file, 116, size, 8);
        file
    }

    #[test]
    fn free_list_inside_the_segment_is_accepted() {
        let file = heap_with_free_block(8, 1, 24);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        heap.validate_free_list(&file, 8).unwrap();
        assert_eq!(heap.read_string(&file, 0).unwrap(), "abcdefg");
        // An all-ones head is "no free list" too.
        let file = heap_with_free_block(u64::MAX, 0, 0);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        assert!(heap.validate_free_list(&file, 8).is_ok());
    }

    #[test]
    fn bad_free_list_is_rejected_like_libhdf5() {
        for (head, next, size, why) in [
            (40, 1, 8, "head past the segment"),
            (8, 1, 25, "block runs past the segment"),
            (8, 0, 8, "next offset of zero"),
            (8, 8, 8, "cycle"),
            (8, 999, 8, "next past the segment"),
        ] {
            let file = heap_with_free_block(head, next, size);
            // The header itself parses; the free list is checked on use.
            let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
            assert_eq!(
                heap.validate_free_list(&file, 8).unwrap_err(),
                FormatError::InvalidLocalHeapFreeList,
                "{why}"
            );
        }
    }

    #[test]
    fn invalid_version() {
        let mut file = build_heap_file(0, 100, &["x"], 8, 8);
        file[4] = 1; // bad version
        let err = LocalHeap::parse(&file, 0, 8, 8).unwrap_err();
        assert_eq!(err, FormatError::InvalidLocalHeapVersion(1));
    }

    /// Header, free list and strings read identically through a
    /// `read_at`-only storage, for every truncation of the file.
    #[test]
    fn storage_reads_match_slice_reads() {
        use crate::storage::CountingStorage;
        let plain = build_heap_file(0, 64, &["", "alpha", "beta"], 8, 8);
        // A free block of 16 bytes at segment offset 12, ending the list.
        let mut free = build_heap_file(0, 64, &["", "alpha", "beta", &"x".repeat(20)], 8, 8);
        free[16..24].copy_from_slice(&12u64.to_le_bytes());
        free[64 + 12..64 + 20].copy_from_slice(&1u64.to_le_bytes());
        free[64 + 20..64 + 28].copy_from_slice(&16u64.to_le_bytes());
        let mut bad_free = free.clone();
        bad_free[64 + 20..64 + 28].copy_from_slice(&99u64.to_le_bytes());
        for full in [plain, free, bad_free] {
            for cut in 0..=full.len() {
                let f = &full[..cut];
                let storage = CountingStorage::new(f.to_vec());
                let want = LocalHeap::parse(f, 0, 8, 8);
                let got = LocalHeap::parse_in(&storage, 0, 8, 8);
                assert_eq!(format!("{got:?}"), format!("{want:?}"));
                let Ok(heap) = want else { continue };
                assert_eq!(
                    heap.validate_free_list_in(&storage, 8),
                    heap.validate_free_list(f, 8)
                );
                for off in [0u64, 1, 2, 6, 7, 11, 100] {
                    assert_eq!(heap.read_string_in(&storage, off), heap.read_string(f, off));
                }
            }
        }
    }

    /// Names of every length around the first read's size, and one with no
    /// terminator, read identically through a `read_at`-only storage; a
    /// short name in a heap whose header claims a huge data segment costs
    /// one small read, not a read of the rest of the file.
    #[test]
    fn long_names_and_hostile_segment_sizes() {
        use crate::storage::CountingStorage;
        let names: Vec<String> = [0usize, 1, 63, 64, 65, 255, 256, 257, 1000, 5000]
            .iter()
            .map(|&n| "n".repeat(n))
            .collect();
        let refs: Vec<&str> = names.iter().map(String::as_str).collect();
        let mut file = build_heap_file(0, 64, &refs, 8, 8);
        let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
        let storage = CountingStorage::new(file.clone());
        let mut off = 0u64;
        for name in &names {
            let got = heap.read_string_in(&storage, off);
            assert_eq!(got, heap.read_string(&file, off));
            assert_eq!(got.unwrap(), *name);
            off += name.len() as u64 + 1;
        }
        // The last name loses its terminator: both report the same error.
        let seg_end = 64 + heap.data_segment_size as usize;
        file[seg_end - 1] = b'n';
        let storage = CountingStorage::new(file.clone());
        let last = off - names[names.len() - 1].len() as u64 - 1;
        let want = heap.read_string(&file, last);
        assert!(want.is_err());
        assert_eq!(heap.read_string_in(&storage, last), want);

        // A 64 MiB file whose heap claims a data segment reaching its end.
        let mut big = build_heap_file(0, 64, &["short", "names"], 8, 8);
        big.resize(64 << 20, 0);
        big[8..16].copy_from_slice(&((64u64 << 20) - 64).to_le_bytes());
        let heap = LocalHeap::parse(&big, 0, 8, 8).unwrap();
        let storage = CountingStorage::new(big.clone());
        assert_eq!(heap.read_string_in(&storage, 6).unwrap(), "names");
        assert_eq!((storage.reads(), storage.bytes_read()), (1, 64));
    }
}