//! Walk every block of a fractal heap and verify it: signatures, the
//! back-pointer to the heap header, each block's heap offset, and the
//! checksums (always present on indirect blocks; on direct blocks when the
//! heap header's flag says so). The library reads only the blocks an object
//! lives in and does not verify block checksums, so `check` does it here.

use std::collections::HashSet;

use clawhdf5_format::checksum::jenkins_lookup3;
use clawhdf5_format::fractal_heap::FractalHeapHeader;

use crate::h5::{Error, H5};

/// Heap header flag bit 1: direct blocks carry a checksum.
const FLAG_CHECKSUM_DBLOCKS: u8 = 0x02;
const MAX_DEPTH: u32 = 16;
const MAX_BLOCKS: usize = 1 << 20;

#[derive(Default, Debug)]
pub struct HeapReport {
    pub direct_blocks: usize,
    pub indirect_blocks: usize,
    /// Blocks whose checksum was verified.
    pub checksums: usize,
    pub problems: Vec<Error>,
}

struct Walk<'a> {
    data: &'a [u8],
    heap: u64,
    fh: FractalHeapHeader,
    checksum_dblocks: bool,
    boff_bytes: usize,
    os: usize,
    ls: usize,
    seen: HashSet<u64>,
    r: HeapReport,
}

fn le(b: &[u8]) -> u64 {
    b.iter()
        .take(8)
        .enumerate()
        .fold(0u64, |a, (i, &x)| a | (u64::from(x) << (8 * i)))
}

fn undefined(v: u64, os: usize) -> bool {
    if os >= 8 {
        v == u64::MAX
    } else {
        v == (1u64 << (8 * os)) - 1
    }
}

fn log2(v: u64) -> u32 {
    63u32.saturating_sub(v.max(1).leading_zeros())
}

/// Verify the fractal heap whose header is at `heap`. The header itself is
/// parsed (and its checksum verified) by the library; an error there is
/// returned as the only problem.
pub fn verify(h5: &H5, heap: u64) -> HeapReport {
    let data = h5.data();
    let Ok(off) = usize::try_from(heap) else {
        return HeapReport {
            problems: vec![Error::at(heap, "fractal heap address out of range")],
            ..Default::default()
        };
    };
    let fh = match FractalHeapHeader::parse(data, off, h5.os(), h5.ls()) {
        Ok(f) => f,
        Err(e) => {
            return HeapReport {
                problems: vec![Error::at(heap, format!("fractal heap header: {e}"))],
                ..Default::default()
            };
        }
    };
    // Flags: signature(4) version(1) heap ID length(2) filter length(2) flags(1).
    let flags = data.get(off + 9).copied().unwrap_or(0);
    let mut w = Walk {
        data,
        heap,
        checksum_dblocks: flags & FLAG_CHECKSUM_DBLOCKS != 0 && fh.filter_pipeline.is_none(),
        boff_bytes: usize::from(fh.max_heap_size).div_ceil(8),
        os: usize::from(h5.os()),
        ls: usize::from(h5.ls()),
        fh,
        seen: HashSet::new(),
        r: HeapReport::default(),
    };
    if w.fh.table_width == 0 || w.fh.starting_block_size == 0 || !w.fh.table_width.is_power_of_two()
    {
        w.problem(heap, "fractal heap header: invalid doubling table geometry");
        return w.r;
    }
    let root = w.fh.root_block_address;
    if !undefined(root, w.os) {
        if w.fh.current_rows_in_root_indirect_block == 0 {
            let size = w.fh.starting_block_size;
            w.direct(root, size, 0);
        } else {
            let rows = w.fh.current_rows_in_root_indirect_block;
            w.indirect(root, rows, 0, 0);
        }
    }
    w.r
}

impl Walk<'_> {
    fn problem(&mut self, addr: u64, msg: impl Into<String>) {
        self.r.problems.push(Error::at(addr, msg));
    }

    fn row_size(&self, row: usize) -> Option<u64> {
        let s = self.fh.starting_block_size;
        if row <= 1 {
            Some(s)
        } else {
            let sh = u32::try_from(row - 1).ok()?;
            s.checked_mul(1u64.checked_shl(sh)?)
        }
    }

    fn max_direct_rows(&self) -> usize {
        let ratio = (self.fh.max_direct_block_size / self.fh.starting_block_size).max(1);
        log2(ratio) as usize + 2
    }

    fn rows_for_size(&self, size: u64) -> u16 {
        let first = log2(self.fh.starting_block_size) + log2(u64::from(self.fh.table_width));
        (log2(size).saturating_sub(first) + 1) as u16
    }

    /// Common block prefix: signature, version, heap header address and
    /// block offset. Returns the position after it, or `None` after
    /// recording a problem.
    fn prefix(&mut self, addr: u64, sig: &[u8; 4], what: &str, heap_offset: u64) -> Option<usize> {
        if !self.seen.insert(addr) {
            self.problem(
                addr,
                format!("fractal heap {what} block reached twice (cycle)"),
            );
            return None;
        }
        if self.seen.len() > MAX_BLOCKS {
            self.problem(self.heap, "fractal heap has too many blocks; stopped");
            return None;
        }
        let Ok(start) = usize::try_from(addr) else {
            self.problem(
                addr,
                format!("fractal heap {what} block address out of range"),
            );
            return None;
        };
        let hdr_len = 5 + self.os + self.boff_bytes;
        let Some(b) = start
            .checked_add(hdr_len)
            .and_then(|e| self.data.get(start..e))
        else {
            self.problem(
                addr,
                format!("fractal heap {what} block lies past the end of the file"),
            );
            return None;
        };
        if &b[..4] != sig {
            self.problem(addr, format!("fractal heap {what} block: bad signature"));
            return None;
        }
        if b[4] != 0 {
            self.problem(addr, format!("fractal heap {what} block: version {}", b[4]));
            return None;
        }
        let back = le(&b[5..5 + self.os]);
        if back != self.heap {
            self.problem(
                addr,
                format!(
                    "fractal heap {what} block points at heap header {back:#x}, not {:#x}",
                    self.heap
                ),
            );
        }
        let boff = le(&b[5 + self.os..hdr_len]);
        if boff != heap_offset {
            self.problem(
                addr,
                format!("fractal heap {what} block has heap offset {boff}, expected {heap_offset}"),
            );
        }
        Some(start + hdr_len)
    }

    fn direct(&mut self, addr: u64, size: u64, heap_offset: u64) {
        let Some(pos) = self.prefix(addr, b"FHDB", "direct", heap_offset) else {
            return;
        };
        self.r.direct_blocks += 1;
        if self.fh.filter_pipeline.is_some() {
            return; // stored filtered: its size on disk is not the block size
        }
        let start = pos - (5 + self.os + self.boff_bytes);
        let Some(end) = usize::try_from(size)
            .ok()
            .and_then(|s| start.checked_add(s))
        else {
            self.problem(addr, "fractal heap direct block size out of range");
            return;
        };
        let Some(block) = self.data.get(start..end) else {
            self.problem(
                addr,
                "fractal heap direct block extends past the end of the file",
            );
            return;
        };
        if self.checksum_dblocks {
            let Some(stored) = block.get(pos - start..pos - start + 4) else {
                self.problem(addr, "fractal heap direct block too small for its checksum");
                return;
            };
            let stored = u32::from_le_bytes([stored[0], stored[1], stored[2], stored[3]]);
            let mut copy = block.to_vec();
            copy[pos - start..pos - start + 4].fill(0);
            let computed = jenkins_lookup3(&copy);
            self.r.checksums += 1;
            if computed != stored {
                self.problem(
                    addr,
                    format!(
                        "fractal heap direct block: checksum mismatch: stored {stored:#010x}, computed {computed:#010x}"
                    ),
                );
            }
        }
    }

    fn indirect(&mut self, addr: u64, nrows: u16, heap_offset: u64, depth: u32) {
        if depth > MAX_DEPTH {
            self.problem(addr, "fractal heap indirect blocks nested too deeply");
            return;
        }
        let Some(mut pos) = self.prefix(addr, b"FHIB", "indirect", heap_offset) else {
            return;
        };
        self.r.indirect_blocks += 1;
        let start = pos - (5 + self.os + self.boff_bytes);
        let width = usize::from(self.fh.table_width);
        let filtered = self.fh.filter_pipeline.is_some();
        let direct_rows = self.max_direct_rows();
        let mut children: Vec<(u64, bool, u64, u64)> = Vec::new(); // addr, direct, size/rows, offset
        let mut off = heap_offset;
        for row in 0..usize::from(nrows) {
            let Some(rs) = self.row_size(row) else {
                self.problem(addr, "fractal heap row size overflows");
                return;
            };
            let direct = row < direct_rows;
            for _ in 0..width {
                let Some(b) = self.data.get(pos..pos + self.os) else {
                    self.problem(
                        addr,
                        "fractal heap indirect block extends past the end of the file",
                    );
                    return;
                };
                let child = le(b);
                pos += self.os;
                if direct && filtered {
                    pos += self.ls + 4;
                }
                if !undefined(child, self.os) {
                    children.push((child, direct, rs, off));
                }
                off = off.saturating_add(rs);
            }
        }
        let Some(stored) = self.data.get(pos..pos + 4) else {
            self.problem(
                addr,
                "fractal heap indirect block extends past the end of the file",
            );
            return;
        };
        let stored = u32::from_le_bytes([stored[0], stored[1], stored[2], stored[3]]);
        let computed = jenkins_lookup3(&self.data[start..pos]);
        self.r.checksums += 1;
        if computed != stored {
            self.problem(
                addr,
                format!(
                    "fractal heap indirect block: checksum mismatch: stored {stored:#010x}, computed {computed:#010x}"
                ),
            );
            return; // its child pointers cannot be trusted
        }
        for (child, direct, size, off) in children {
            if direct {
                self.direct(child, size, off);
            } else {
                let rows = self.rows_for_size(size);
                self.indirect(child, rows, off, depth + 1);
            }
        }
    }
}