//! Selection reads that cost what the selection costs, not what the dataset
//! costs.
//!
//! [`crate::data_read::read_raw_data_selection`] used to decode the *entire*
//! dataset and then pick elements out of it, so reading a 64x64 window of a
//! large dataset took about as long as reading all of it. A contiguous
//! dataset's selection is now copied straight out of the file, one `memcpy`
//! per contiguous run of selected elements (`crate::gather`). For chunked
//! data the selection's bounding box is materialised — only the chunks that
//! overlap it — and the extractor runs over that small buffer with the
//! selection translated to the box's origin. Extraction semantics are
//! therefore exactly the full-read ones.

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

use crate::chunked_read::{alloc_output, checked_byte_len, list_chunks_for_read_in};
use crate::data_layout::DataLayout;
use crate::data_read::extract_selection_from_buffer;
use crate::dataspace::Dataspace;
use crate::error::FormatError;
use crate::filter_pipeline::FilterPipeline;
use crate::filters::{all_filters_skipped, decompress_chunk_exact_with};
use crate::selection::Selection;
use crate::storage::{ExtentReq, Storage, for_each_extent_batch};

/// The smallest axis-aligned box containing every selected element, as
/// `(start, extent)` per dimension. `None` when there is nothing to gain or
/// the selection is not valid for `dims` (the caller's full path then reports
/// the error exactly as before).
fn bounding_box(selection: &Selection, dims: &[u64]) -> Option<(Vec<u64>, Vec<u64>)> {
    match selection {
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => {
            let rank = dims.len();
            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
                return None;
            }
            let mut extent = Vec::with_capacity(rank);
            for d in 0..rank {
                if count[d] == 0 || block[d] == 0 {
                    return None;
                }
                // Last selected index + 1, relative to start.
                let span = (count[d] - 1)
                    .checked_mul(stride[d])?
                    .checked_add(block[d])?;
                if start[d].checked_add(span)? > dims[d] {
                    return None;
                }
                extent.push(span);
            }
            Some((start.clone(), extent))
        }
        Selection::Points(points) => {
            let rank = dims.len();
            let first = points.first()?;
            if first.len() != rank {
                return None;
            }
            let (mut lo, mut hi) = (first.clone(), first.clone());
            for p in points {
                if p.len() != rank {
                    return None;
                }
                for d in 0..rank {
                    if p[d] >= dims[d] {
                        return None;
                    }
                    lo[d] = lo[d].min(p[d]);
                    hi[d] = hi[d].max(p[d]);
                }
            }
            let extent = lo.iter().zip(&hi).map(|(l, h)| h - l + 1).collect();
            Some((lo, extent))
        }
        Selection::All | Selection::None => None,
    }
}

/// Check that `selection` addresses only elements that exist in a dataset of
/// shape `dims`. Without this an out-of-range selection read *something*: a
/// hyperslab past the edge came back padded with zeros, and a point whose
/// column was out of range wrapped into the next row.
pub fn validate(selection: &Selection, dims: &[u64]) -> Result<(), FormatError> {
    let rank = dims.len();
    let bad = |msg: alloc_or_std::String| Err(FormatError::SelectionOutOfBounds(msg));
    match selection {
        Selection::All | Selection::None => Ok(()),
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => {
            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
                return bad(format!("hyperslab rank does not match dataset rank {rank}"));
            }
            for d in 0..rank {
                if count[d] == 0 || block[d] == 0 {
                    continue; // selects nothing along this dimension
                }
                let end = (count[d] - 1)
                    .checked_mul(stride[d])
                    .and_then(|v| v.checked_add(block[d]))
                    .and_then(|v| v.checked_add(start[d]));
                if !end.is_some_and(|end| end <= dims[d]) {
                    return bad(format!(
                        "dimension {d}: start {} stride {} count {} block {} exceeds extent {}",
                        start[d], stride[d], count[d], block[d], dims[d]
                    ));
                }
                if block[d] > stride[d] && count[d] > 1 {
                    return bad(format!(
                        "dimension {d}: block {} larger than stride {} (overlapping blocks)",
                        block[d], stride[d]
                    ));
                }
            }
            Ok(())
        }
        Selection::Points(points) => {
            for p in points {
                if p.len() != rank {
                    return bad(format!("point {p:?} does not match dataset rank {rank}"));
                }
                if let Some(d) = (0..rank).find(|&d| p[d] >= dims[d]) {
                    return bad(format!(
                        "point {p:?}: coordinate {} exceeds extent {} of dimension {d}",
                        p[d], dims[d]
                    ));
                }
            }
            Ok(())
        }
    }
}

/// The same selection expressed relative to `origin`.
fn translate(selection: &Selection, origin: &[u64]) -> Selection {
    match selection {
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => Selection::Hyperslab {
            start: start.iter().zip(origin).map(|(s, o)| s - o).collect(),
            stride: stride.clone(),
            count: count.clone(),
            block: block.clone(),
        },
        Selection::Points(points) => Selection::Points(
            points
                .iter()
                .map(|p| p.iter().zip(origin).map(|(c, o)| c - o).collect())
                .collect(),
        ),
        other => other.clone(),
    }
}

/// Copy the part of a source region that overlaps the box into `out` (which
/// is the box, row-major).
///
/// The source region starts at `src_origin` in dataset coordinates, has shape
/// `src_shape`, and its elements are in `src` row-major. One `memcpy` per
/// overlapping row of the last dimension.
#[allow(clippy::too_many_arguments)]
fn copy_overlap(
    src: &[u8],
    src_origin: &[u64],
    src_shape: &[u64],
    out: &mut [u8],
    box_start: &[u64],
    box_extent: &[u64],
    elem_size: usize,
) {
    let rank = box_start.len();
    // Overlap in dataset coordinates.
    let mut lo = vec![0u64; rank];
    let mut hi = vec![0u64; rank];
    for d in 0..rank {
        lo[d] = src_origin[d].max(box_start[d]);
        hi[d] = (src_origin[d] + src_shape[d]).min(box_start[d] + box_extent[d]);
        if lo[d] >= hi[d] {
            return;
        }
    }
    let strides = |shape: &[u64]| {
        let mut s = vec![1u64; rank];
        for d in (0..rank.saturating_sub(1)).rev() {
            s[d] = s[d + 1] * shape[d + 1];
        }
        s
    };
    let (src_strides, out_strides) = (strides(src_shape), strides(box_extent));
    let last = rank - 1;
    // Byte offsets into the in-memory buffers; one that does not fit `usize`
    // (a 32-bit target) is out of both buffers, like one past their ends.
    let bytes = |elements: u64| usize::try_from(elements).ok()?.checked_mul(elem_size);
    let Some(run) = bytes(hi[last] - lo[last]) else {
        return;
    };

    let mut idx = lo.clone();
    loop {
        let src_at: u64 = (0..rank)
            .map(|d| (idx[d] - src_origin[d]) * src_strides[d])
            .sum();
        let out_at: u64 = (0..rank)
            .map(|d| (idx[d] - box_start[d]) * out_strides[d])
            .sum();
        if let (Some(s), Some(o)) = (bytes(src_at), bytes(out_at))
            && let (Some(from), Some(to)) = (
                src.get(s..s.saturating_add(run)),
                out.get_mut(o..o.saturating_add(run)),
            )
        {
            to.copy_from_slice(from);
        }
        // Advance over every dimension but the last.
        let mut d = last;
        loop {
            if d == 0 {
                return;
            }
            d -= 1;
            idx[d] += 1;
            if idx[d] < hi[d] {
                break;
            }
            idx[d] = lo[d];
        }
    }
}

/// Read `selection` without materialising the whole dataset, when that is
/// possible and worthwhile. `Ok(None)` means "use the full-read path": an
/// `All`/`None`/invalid selection, a layout this doesn't handle (compact,
/// virtual, storage-less), or a bounding box covering most of the dataset.
#[allow(clippy::too_many_arguments)]
pub fn read_selection(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    selection: &Selection,
) -> Result<Option<Vec<u8>>, FormatError> {
    read_selection_in(
        file_data,
        layout,
        dataspace,
        elem_size,
        pipeline,
        offset_size,
        length_size,
        selection,
    )
}

/// [`read_selection`] over any [`Storage`].
#[allow(clippy::too_many_arguments)]
pub fn read_selection_in<S: Storage + ?Sized>(
    file_data: &S,
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    selection: &Selection,
) -> Result<Option<Vec<u8>>, FormatError> {
    let dims = &dataspace.dimensions;
    if dims.is_empty() || elem_size == 0 {
        return Ok(None);
    }
    let total = dataspace.checked_num_elements()?;
    // Contiguous data is addressable in place: copy the selection's runs
    // straight out of it, whatever fraction of the dataset it covers, with no
    // intermediate box (and no full copy for a large selection).
    if let (
        DataLayout::Contiguous {
            address: Some(address),
            ..
        },
        Selection::Hyperslab { .. } | Selection::Points(_),
    ) = (layout, selection)
    {
        validate(selection, dims)?;
        let base = usize::try_from(*address)
            .map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
        let file_len = crate::storage::len_usize(file_data);
        let eof = FormatError::UnexpectedEof {
            expected: base,
            available: file_len,
        };
        if let Some(all) = file_data.as_contiguous() {
            let data = all
                .get(base..)
                .and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
                .ok_or(eof)?;
            return crate::gather::gather::<u8>(data, dims, elem_size, selection).map(Some);
        }
        // Not in memory: the same bounds check, then only the selected runs
        // are read.
        let len = checked_byte_len(total, elem_size)
            .ok()
            .filter(|&len| base <= file_len && len <= file_len - base)
            .ok_or(eof)?;
        return crate::gather::gather_storage(
            file_data,
            base as u64,
            len,
            dims,
            elem_size,
            selection,
        )
        .map(Some);
    }
    let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
        return Ok(None);
    };
    let box_elements = box_extent
        .iter()
        .try_fold(1u64, |acc, &e| acc.checked_mul(e))
        .ok_or_else(|| FormatError::Overflow("selection bounding box overflows".into()))?;
    // A box covering most of the dataset gains nothing over the full path.
    if box_elements.saturating_mul(2) > total {
        return Ok(None);
    }
    let mut boxed = alloc_output(checked_byte_len(box_elements, elem_size)?)?;

    match layout {
        DataLayout::Chunked {
            btree_address: Some(_),
            ..
        } => {
            let (chunks, chunk_dims) = list_chunks_for_read_in(
                file_data,
                layout,
                dataspace,
                elem_size,
                pipeline,
                offset_size,
                length_size,
            )?;
            let rank = dims.len();
            let chunk_shape: Vec<u64> = chunk_dims.iter().map(|&d| d as u64).collect();
            let chunk_bytes = crate::chunked_read::checked_chunk_byte_len(&chunk_dims, elem_size)?;
            // The chunks overlapping the box, in index order.
            let wanted: Vec<&crate::chunked_read::ChunkInfo> = chunks
                .iter()
                .filter(|chunk| {
                    if chunk.offsets.len() < rank || chunk.address == u64::MAX {
                        return false;
                    }
                    let origin = &chunk.offsets[..rank];
                    (0..rank).all(|d| {
                        origin[d] < box_start[d] + box_extent[d]
                            && origin[d].saturating_add(chunk_shape[d]) > box_start[d]
                    })
                })
                .collect();
            // Their stored bytes, batch by batch when the file is not in
            // memory; each batch's chunks are decoded into this thread's
            // reusable buffers before the next batch is fetched.
            let reqs: Vec<ExtentReq> = wanted
                .iter()
                .map(|c| crate::chunked_read::chunk_req(c, pipeline, chunk_bytes, true))
                .collect();
            for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
                crate::chunked_read::with_scratch(|scratch| -> Result<(), FormatError> {
                    for i in batch {
                        let chunk = wanted[i];
                        let origin = &chunk.offsets[..rank];
                        usize::try_from(chunk.address).map_err(|_| {
                            FormatError::Overflow("chunk address exceeds usize".into())
                        })?;
                        let raw = raw_bytes.get(i, &reqs[i])?;
                        // Mirrors the full-read path: filter-mask bit i set
                        // means filter i was not applied to this chunk.
                        let data: &[u8] = match pipeline {
                            Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => {
                                decompress_chunk_exact_with(
                                    raw,
                                    pl,
                                    chunk_bytes,
                                    elem_size as u32,
                                    chunk.filter_mask,
                                    &chunk.offsets[..rank],
                                    scratch,
                                )?
                            }
                            _ => raw,
                        };
                        copy_overlap(
                            data,
                            origin,
                            &chunk_shape,
                            &mut boxed,
                            &box_start,
                            &box_extent,
                            elem_size,
                        );
                    }
                    Ok(())
                })
            })?;
        }
        _ => return Ok(None),
    }

    extract_selection_from_buffer(
        &boxed,
        &box_extent,
        elem_size,
        &translate(selection, &box_start),
    )
    .map(Some)
}