//! 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. Here the selection's
//! bounding box is materialised instead — only the rows of a contiguous
//! dataset, or only the chunks, that overlap it — and the existing 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};
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::decompress_chunk;
use crate::selection::Selection;

/// 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;
    let run = ((hi[last] - lo[last]) as usize) * elem_size;

    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();
        let (s, o) = (src_at as usize * elem_size, out_at as usize * elem_size);
        if let (Some(from), Some(to)) = (src.get(s..s + run), out.get_mut(o..o + 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> {
    let dims = &dataspace.dimensions;
    if dims.is_empty() || elem_size == 0 {
        return Ok(None);
    }
    let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
        return Ok(None);
    };
    let total = dataspace.checked_num_elements()?;
    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::Contiguous {
            address: Some(address),
            ..
        } => {
            let base = usize::try_from(*address)
                .map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
            let data = file_data
                .get(base..)
                .and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
                .ok_or(FormatError::UnexpectedEof {
                    expected: base,
                    available: file_data.len(),
                })?;
            let origin = vec![0u64; dims.len()];
            copy_overlap(
                data,
                &origin,
                dims,
                &mut boxed,
                &box_start,
                &box_extent,
                elem_size,
            );
        }
        DataLayout::Chunked {
            btree_address: Some(_),
            ..
        } => {
            let (chunks, chunk_dims) = list_chunks(
                file_data,
                layout,
                dataspace,
                elem_size,
                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)?;
            for chunk in &chunks {
                if chunk.offsets.len() < rank || chunk.address == u64::MAX {
                    continue;
                }
                let origin = &chunk.offsets[..rank];
                let overlaps = (0..rank).all(|d| {
                    origin[d] < box_start[d] + box_extent[d]
                        && origin[d].saturating_add(chunk_shape[d]) > box_start[d]
                });
                if !overlaps {
                    continue;
                }
                let at = usize::try_from(chunk.address)
                    .map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?;
                let raw = at
                    .checked_add(chunk.chunk_size as usize)
                    .and_then(|end| file_data.get(at..end))
                    .ok_or(FormatError::UnexpectedEof {
                        expected: at.saturating_add(chunk.chunk_size as usize),
                        available: file_data.len(),
                    })?;
                // Mirrors the full-read path: a non-zero filter mask means the
                // chunk was stored unfiltered.
                let decoded;
                let data: &[u8] = match pipeline {
                    Some(pl) if chunk.filter_mask == 0 => {
                        decoded = decompress_chunk(raw, pl, chunk_bytes, elem_size as u32)?;
                        &decoded
                    }
                    _ => raw,
                };
                copy_overlap(
                    data,
                    origin,
                    &chunk_shape,
                    &mut boxed,
                    &box_start,
                    &box_extent,
                    elem_size,
                );
            }
        }
        _ => return Ok(None),
    }

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