//! Copying a selection out of a row-major buffer one contiguous run at a time.
//!
//! A selection's elements, in output order, fall into runs that are adjacent
//! in the source: a whole block along the last dimension, blocks that touch
//! (`stride == block`), and whole rows when the inner dimensions are selected
//! in full. Copying run by run turns a 256 x 256 hyperslab of a 1024-wide
//! dataset into 256 `memcpy`s of 1 KiB, where the old extractor recursed and
//! bounds-checked once per element.

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

use crate::data_read::NativeElement;
use crate::error::FormatError;
use crate::selection::Selection;
use crate::storage::{ExtentBytes, ExtentReq, Storage, raw_batches};

/// Row-major element strides of `dims` (the last dimension has stride 1).
fn strides(dims: &[u64]) -> Vec<u64> {
    let mut s = vec![1u64; dims.len()];
    for d in (0..dims.len().saturating_sub(1)).rev() {
        s[d] = s[d + 1].wrapping_mul(dims[d + 1]);
    }
    s
}

/// Merges adjacent runs before handing them on.
struct Coalesce<F: FnMut(u64, u64)> {
    start: u64,
    len: u64,
    emit: F,
}

impl<F: FnMut(u64, u64)> Coalesce<F> {
    #[inline]
    fn push(&mut self, start: u64, len: u64) {
        if len == 0 {
            return;
        }
        if self.len > 0 && self.start.wrapping_add(self.len) == start {
            self.len += len;
            return;
        }
        self.flush();
        self.start = start;
        self.len = len;
    }

    fn flush(&mut self) {
        if self.len > 0 {
            (self.emit)(self.start, self.len);
            self.len = 0;
        }
    }
}

/// Call `emit(first_element, element_count)` for each run of a hyperslab's
/// elements that is contiguous in a row-major dataset of shape `dims`, in
/// the order the selection returns them. Adjacent runs are merged.
///
/// Coordinates at or past a dimension's extent are skipped, as the
/// element-wise extractor always did; callers that want them to be an error
/// validate the selection first. The four vectors must have `dims.len()`
/// entries.
pub(crate) fn hyperslab_runs(
    dims: &[u64],
    start: &[u64],
    stride: &[u64],
    count: &[u64],
    block: &[u64],
    emit: impl FnMut(u64, u64),
) {
    let rank = dims.len();
    let mut out = Coalesce {
        start: 0,
        len: 0,
        emit,
    };
    if rank == 0 {
        out.push(0, 1);
        out.flush();
        return;
    }
    if (0..rank).any(|d| count[d] == 0 || block[d] == 0) {
        return;
    }
    let strides = strides(dims);
    let last = rank - 1;
    // Odometer over the outer dimensions: (block index, offset in block).
    let mut ci = vec![0u64; last];
    let mut bi = vec![0u64; last];
    'outer: loop {
        // Base offset of this row, or skip it if a coordinate is out of range.
        let mut base = 0u64;
        let mut in_range = true;
        for d in 0..last {
            let coord = start[d]
                .saturating_add(ci[d].saturating_mul(stride[d]))
                .saturating_add(bi[d]);
            if coord >= dims[d] {
                in_range = false;
                break;
            }
            base = base.wrapping_add(coord.wrapping_mul(strides[d]));
        }
        if in_range && (stride[last] == block[last] || count[last] == 1) {
            // Blocks that touch (the common unit-stride case: block 1,
            // stride 1) are one range; don't split it into per-element runs.
            let s = start[last];
            let e = s
                .saturating_add(count[last].saturating_mul(block[last]))
                .min(dims[last]);
            if s < e {
                out.push(base.wrapping_add(s), e - s);
            }
        } else if in_range {
            for c in 0..count[last] {
                let s = start[last].saturating_add(c.saturating_mul(stride[last]));
                if s >= dims[last] {
                    continue;
                }
                let e = s.saturating_add(block[last]).min(dims[last]);
                out.push(base.wrapping_add(s), e - s);
            }
        }
        // Advance the odometer, last outer dimension fastest.
        let mut d = last;
        loop {
            if d == 0 {
                break 'outer;
            }
            d -= 1;
            bi[d] += 1;
            if bi[d] < block[d] {
                break;
            }
            bi[d] = 0;
            ci[d] += 1;
            if ci[d] < count[d] {
                break;
            }
            ci[d] = 0;
        }
    }
    out.flush();
}

/// The selected elements of `src` — a row-major dataset of shape `dims` and
/// `elem_size`-byte elements — copied into a fresh `Vec<T>`, one `memcpy` per
/// contiguous run, with no zero-filling of the output first.
///
/// For `T` other than `u8`, `elem_size` must equal `size_of::<T>()`. The
/// selection must be a validated hyperslab, point list or `None` (`All` is the
/// caller's to handle); `src` must hold exactly the dataset. Anything that
/// would read outside `src` is an error, never a partial result.
pub(crate) fn gather<T: NativeElement>(
    src: &[u8],
    dims: &[u64],
    elem_size: usize,
    selection: &Selection,
) -> Result<Vec<T>, FormatError> {
    let t_size = core::mem::size_of::<T>();
    if elem_size == 0 || (t_size != 1 && t_size != elem_size) {
        return Err(FormatError::DataSizeMismatch {
            expected: t_size,
            actual: elem_size,
        });
    }
    let n_elements = match selection {
        Selection::None => 0,
        Selection::Hyperslab { count, block, .. } => count
            .iter()
            .zip(block)
            .try_fold(1u64, |acc, (&c, &b)| acc.checked_mul(c.checked_mul(b)?))
            .ok_or_else(|| FormatError::Overflow("hyperslab count x block overflows".into()))?,
        Selection::Points(points) => points.len() as u64,
        Selection::All => {
            return Err(FormatError::SelectionOutOfBounds(
                "gather does not take Selection::All".into(),
            ));
        }
    };
    let out_bytes = crate::chunked_read::checked_byte_len(n_elements, elem_size)?;
    let out_len = out_bytes / t_size;
    let mut out: Vec<T> = crate::bulk_alloc::vec_for_bulk(out_len);
    let dst = out.as_mut_ptr().cast::<u8>();
    let mut written = 0usize;
    let mut failed = false;
    let mut copy_run = |first: u64, n: u64| {
        if failed {
            return;
        }
        let range = usize::try_from(first)
            .ok()
            .and_then(|f| f.checked_mul(elem_size))
            .zip(
                usize::try_from(n)
                    .ok()
                    .and_then(|n| n.checked_mul(elem_size)),
            )
            .and_then(|(at, len)| Some((at, len, at.checked_add(len)?)));
        match range {
            Some((at, len, end)) if end <= src.len() && written + len <= out_bytes => {
                // SAFETY: `src[at..end]` is in bounds (checked above), and
                // `dst + written .. + len` lies within `out`'s capacity of
                // `out_bytes` bytes (checked above); `out` is a fresh
                // allocation, so the regions do not overlap.
                unsafe {
                    core::ptr::copy_nonoverlapping(src.as_ptr().add(at), dst.add(written), len)
                };
                written += len;
            }
            _ => failed = true,
        }
    };
    let mut bad_point = false;
    match selection {
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => {
            let rank = dims.len();
            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
                return Err(FormatError::SelectionOutOfBounds(
                    "hyperslab rank does not match dataset rank".into(),
                ));
            }
            hyperslab_runs(dims, start, stride, count, block, &mut copy_run);
        }
        Selection::Points(points) => {
            let strides = strides(dims);
            let mut runs = Coalesce {
                start: 0,
                len: 0,
                emit: &mut copy_run,
            };
            for p in points {
                if p.len() != dims.len() || p.iter().zip(dims).any(|(c, n)| c >= n) {
                    bad_point = true;
                    break;
                }
                let at = p
                    .iter()
                    .zip(&strides)
                    .fold(0u64, |acc, (c, s)| acc.wrapping_add(c.wrapping_mul(*s)));
                runs.push(at, 1);
            }
            runs.flush();
        }
        Selection::None | Selection::All => {}
    }
    if failed || bad_point || written != out_bytes {
        return Err(FormatError::SelectionOutOfBounds(
            "selection addresses elements outside the dataset".into(),
        ));
    }
    // SAFETY: all `out_bytes` bytes, i.e. `out_len` values of `T`, were
    // written above, and every bit pattern is a valid `T` (`NativeElement`).
    unsafe { out.set_len(out_len) };
    Ok(out)
}

/// Largest gap between two of a selection's runs that [`gather_storage`]
/// reads through rather than asking for the runs separately: skipping a
/// few KiB costs a remote backend far less than another request (and a
/// local one less than another call and allocation).
pub(crate) const GATHER_GAP_BYTES: usize = 4 << 10;

/// Largest single read [`gather_storage`] makes of a selection's runs: runs
/// are merged into reads up to this size, and a longer run is split.
pub(crate) const GATHER_SPAN_BYTES: usize = 8 << 20;

/// Call `emit(first_element, element_count)` for each run of a validated
/// hyperslab or point selection (in output order; see [`hyperslab_runs`]),
/// or the error for a hyperslab of the wrong rank or a point outside `dims`
/// (runs before that point have been emitted).
fn selection_runs(
    dims: &[u64],
    selection: &Selection,
    emit: &mut dyn FnMut(u64, u64),
) -> Result<(), FormatError> {
    match selection {
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => {
            let rank = dims.len();
            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
                return Err(FormatError::SelectionOutOfBounds(
                    "hyperslab rank does not match dataset rank".into(),
                ));
            }
            hyperslab_runs(dims, start, stride, count, block, emit);
        }
        Selection::Points(points) => {
            let strides = strides(dims);
            let mut coalesce = Coalesce {
                start: 0,
                len: 0,
                emit,
            };
            for p in points {
                if p.len() != dims.len() || p.iter().zip(dims).any(|(c, n)| c >= n) {
                    return Err(FormatError::SelectionOutOfBounds(
                        "selection addresses elements outside the dataset".into(),
                    ));
                }
                let at = p
                    .iter()
                    .zip(&strides)
                    .fold(0u64, |acc, (c, s)| acc.wrapping_add(c.wrapping_mul(*s)));
                coalesce.push(at, 1);
            }
            coalesce.flush();
        }
        Selection::None | Selection::All => {}
    }
    Ok(())
}

/// One read of [`gather_storage`]: bytes `[start, end)` of the dataset,
/// which hold the output's bytes up to `out_end` (from where the previous
/// span's end left off).
#[derive(Clone, Copy)]
struct Span {
    start: usize,
    end: usize,
    out_end: usize,
}

/// [`gather`] of bytes (`T = u8`) from a dataset that is not in memory: the
/// dataset's `src_len` bytes start at `base` in `file`, which must hold all
/// of them (the caller checks). Same checks and errors as [`gather`].
///
/// The selection's runs are walked twice. The first walk checks them and
/// plans the reads: runs in increasing order with at most
/// [`GATHER_GAP_BYTES`] between them are read as one span (the gap is read
/// and dropped), up to [`GATHER_SPAN_BYTES`] per span. So a strided
/// selection is a few large reads, not one per element, and nothing is
/// allocated per run. The spans are fetched batch by batch (one
/// [`Storage::read_ranges`] call per [`crate::storage::RAW_BATCH_BYTES`])
/// while the second walk copies each run out of its span.
pub(crate) fn gather_storage<S: Storage + ?Sized>(
    file: &S,
    base: u64,
    src_len: usize,
    dims: &[u64],
    elem_size: usize,
    selection: &Selection,
) -> Result<Vec<u8>, FormatError> {
    if elem_size == 0 {
        return Err(FormatError::DataSizeMismatch {
            expected: 1,
            actual: elem_size,
        });
    }
    let n_elements = match selection {
        Selection::None => 0,
        Selection::Hyperslab { count, block, .. } => count
            .iter()
            .zip(block)
            .try_fold(1u64, |acc, (&c, &b)| acc.checked_mul(c.checked_mul(b)?))
            .ok_or_else(|| FormatError::Overflow("hyperslab count x block overflows".into()))?,
        Selection::Points(points) => points.len() as u64,
        Selection::All => {
            return Err(FormatError::SelectionOutOfBounds(
                "gather does not take Selection::All".into(),
            ));
        }
    };
    let out_bytes = crate::chunked_read::checked_byte_len(n_elements, elem_size)?;
    let outside = || {
        FormatError::SelectionOutOfBounds("selection addresses elements outside the dataset".into())
    };

    // First walk: check every run and plan the spans.
    let mut spans: Vec<Span> = Vec::new();
    let mut total = 0usize;
    let mut failed = false;
    selection_runs(dims, selection, &mut |first: u64, n: u64| {
        if failed {
            return;
        }
        let range = usize::try_from(first)
            .ok()
            .and_then(|f| f.checked_mul(elem_size))
            .zip(
                usize::try_from(n)
                    .ok()
                    .and_then(|n| n.checked_mul(elem_size)),
            )
            .and_then(|(at, len)| Some((at, len, at.checked_add(len)?)));
        let Some((mut at, mut len)) = range
            .filter(|&(_, len, end)| end <= src_len && len <= out_bytes - total)
            .map(|(at, len, _)| (at, len))
        else {
            failed = true;
            return;
        };
        while len > 0 {
            let room = match spans.last_mut() {
                Some(s)
                    if at >= s.end
                        && at - s.end <= GATHER_GAP_BYTES
                        && at - s.start < GATHER_SPAN_BYTES =>
                {
                    let take = len.min(GATHER_SPAN_BYTES - (at - s.start));
                    s.end = at + take;
                    s.out_end += take;
                    take
                }
                _ => {
                    let take = len.min(GATHER_SPAN_BYTES);
                    spans.push(Span {
                        start: at,
                        end: at + take,
                        out_end: total + take,
                    });
                    take
                }
            };
            total += room;
            at += room;
            len -= room;
        }
    })?;
    if failed || total != out_bytes {
        return Err(outside());
    }

    // The spans' reads, and the batches they are fetched in.
    let reqs: Vec<ExtentReq> = spans
        .iter()
        .map(|s| ExtentReq {
            addr: base + s.start as u64,
            len: s.end - s.start,
            fetch: Some(s.end - s.start),
        })
        .collect();
    let batches = raw_batches(reqs.len(), false, |i| reqs[i].len);

    // Second walk: copy each run out of its span, fetching each batch of
    // spans when the walk reaches it (and dropping the previous one).
    let mut out = crate::bulk_alloc::vec_for_bulk(out_bytes);
    let mut span = 0usize;
    let mut batch = 0usize;
    let mut fetched: Option<ExtentBytes<'_>> = None;
    let mut error: Option<FormatError> = None;
    selection_runs(dims, selection, &mut |first: u64, n: u64| {
        if error.is_some() {
            return;
        }
        // Checked by the first walk (these cannot saturate or wrap).
        let mut at = crate::addr::saturating_usize(first).wrapping_mul(elem_size);
        let mut len = crate::addr::saturating_usize(n).wrapping_mul(elem_size);
        while len > 0 {
            while spans.get(span).is_some_and(|s| s.out_end <= out.len()) {
                span += 1;
            }
            if fetched.is_none() || span >= batches[batch].end {
                fetched = None;
                while batches.get(batch).is_some_and(|b| span >= b.end) {
                    batch += 1;
                }
                let (Some(b), Some(_)) = (batches.get(batch).cloned(), spans.get(span)) else {
                    // The second walk emitted more than the first.
                    error = Some(outside());
                    return;
                };
                match ExtentBytes::fetch(file, &reqs[b.clone()], b.start) {
                    Ok(f) => fetched = Some(f),
                    Err(e) => {
                        error = Some(e);
                        return;
                    }
                }
            }
            let s = spans[span];
            let take = len.min(s.out_end - out.len());
            let bytes = match fetched
                .as_ref()
                .map(|f| f.get(span, &reqs[span]))
                .unwrap_or_else(|| Err(outside()))
            {
                Ok(b) => b,
                Err(e) => {
                    error = Some(e);
                    return;
                }
            };
            match at
                .checked_sub(s.start)
                .and_then(|o| bytes.get(o..o.checked_add(take)?))
            {
                Some(b) => out.extend_from_slice(b),
                None => {
                    error = Some(outside());
                    return;
                }
            }
            at += take;
            len -= take;
        }
    })?;
    if let Some(e) = error {
        return Err(e);
    }
    if out.len() != out_bytes {
        return Err(outside());
    }
    Ok(out)
}

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

    fn runs(dims: &[u64], sel: [&[u64]; 4]) -> Vec<(u64, u64)> {
        let mut v = Vec::new();
        hyperslab_runs(dims, sel[0], sel[1], sel[2], sel[3], |s, n| v.push((s, n)));
        v
    }

    #[test]
    fn runs_merge_blocks_and_whole_rows() {
        // A box: one run per row.
        assert_eq!(
            runs(&[4, 10], [&[1, 2], &[1, 1], &[2, 3], &[1, 1]]),
            vec![(12, 3), (22, 3)]
        );
        // Whole rows: one run.
        assert_eq!(
            runs(&[4, 10], [&[1, 0], &[1, 1], &[3, 10], &[1, 1]]),
            vec![(10, 30)]
        );
        // stride == block: blocks merge.
        assert_eq!(
            runs(&[1, 10], [&[0, 1], &[1, 2], &[1, 4], &[1, 2]]),
            vec![(1, 8)]
        );
        // Strided with blocks along both dimensions.
        assert_eq!(
            runs(&[6, 10], [&[0, 1], &[3, 4], &[2, 2], &[2, 2]]),
            vec![
                (1, 2),
                (5, 2),
                (11, 2),
                (15, 2),
                (31, 2),
                (35, 2),
                (41, 2),
                (45, 2)
            ]
        );
        // Empty.
        assert!(runs(&[4, 10], [&[0, 0], &[1, 1], &[0, 3], &[1, 1]]).is_empty());
        // Scalar.
        assert_eq!(runs(&[], [&[], &[], &[], &[]]), vec![(0, 1)]);
    }

    #[test]
    fn gather_matches_element_order_and_rejects_out_of_range() {
        let dims = [3u64, 4];
        let src: Vec<u8> = (0..12u16).flat_map(|v| v.to_le_bytes()).collect();
        let sel = Selection::Hyperslab {
            start: vec![0, 1],
            stride: vec![2, 2],
            count: vec![2, 2],
            block: vec![1, 1],
        };
        let got: Vec<u8> = gather(&src, &dims, 2, &sel).unwrap();
        let want: Vec<u8> = [1u16, 3, 9, 11]
            .iter()
            .flat_map(|v| v.to_le_bytes())
            .collect();
        assert_eq!(got, want);
        let pts = Selection::Points(vec![vec![2, 3], vec![0, 0], vec![0, 1]]);
        let got: Vec<u8> = gather(&src, &dims, 2, &pts).unwrap();
        let want: Vec<u8> = [11u16, 0, 1].iter().flat_map(|v| v.to_le_bytes()).collect();
        assert_eq!(got, want);
        // Past the extent, or a source shorter than the dataset: an error.
        let bad = Selection::Points(vec![vec![3, 0]]);
        assert!(gather::<u8>(&src, &dims, 2, &bad).is_err());
        let past = Selection::Hyperslab {
            start: vec![2, 0],
            stride: vec![1, 1],
            count: vec![2, 4],
            block: vec![1, 1],
        };
        assert!(gather::<u8>(&src, &dims, 2, &past).is_err());
        assert!(gather::<u8>(&src[..20], &dims, 2, &pts).is_err());
    }
}