//! PyDataset — h5py-style read access to HDF5 datasets.
//!
//! `ds[key]` parses the key into hyperslab selections (see `select`) and
//! reads them through the facade's `read_selection`, which decodes only the
//! chunks a small selection touches (see its docs for when it decodes the
//! whole dataset instead); the
//! bytes it returns become the numpy array's buffer without a copy (see
//! `convert`). All file access and decoding runs with the GIL released, so
//! Python threads reading the same or different datasets run in parallel,
//! and a remote file's network reads never hold the GIL.

use std::sync::{Arc, Mutex, PoisonError};

use clawhdf5_format::datatype::Datatype;
use clawhdf5_format::object_header::ObjectHeader;
use clawhdf5_rs::File;
use pyo3::exceptions::{PyNotImplementedError, PyOSError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyList, PyTuple};

use crate::attrs::PyAttrs;
use crate::convert::{Converter, Elements, VlError, resolve_vl};
use crate::handle::Handle;
use crate::select::{self, Plan};
use crate::{PyEmpty, edit, node, to_py_err};

/// What opening a dataset reads from the file (without the GIL).
pub(crate) struct DatasetMeta {
    /// `None` for a dataset with a null dataspace (h5py's `Empty`).
    shape: Option<Vec<u64>>,
    chunks: Option<Vec<u64>>,
    datatype: Datatype,
}

impl DatasetMeta {
    pub(crate) fn load(f: &File, addr: u64, hdr: &ObjectHeader, path: &str) -> PyResult<Self> {
        let null = node::is_null(&node::dataspace(f, hdr, path)?);
        let ds = f.dataset_at(addr).map_err(to_py_err)?;
        let shape = if null {
            None
        } else {
            Some(ds.shape().map_err(to_py_err)?)
        };
        let datatype = ds.raw_datatype().map_err(to_py_err)?;
        let chunks = shape
            .as_ref()
            .and_then(|s| node::chunk_shape(f, hdr, s.len()));
        Ok(Self {
            shape,
            chunks,
            datatype,
        })
    }
}

/// A dataset in a file opened for reading.
///
/// ```python
/// ds = f['group/dataset']
/// ds.shape, ds.dtype, ds.attrs['units']
/// block = ds[10:20, ::2]      # a small selection reads only its chunks
/// ```
#[pyclass(name = "Dataset")]
pub struct PyDataset {
    handle: Arc<Handle>,
    path: String,
    /// Where the dataset's object header is: reads open it from here rather
    /// than resolve `path` again.
    addr: u64,
    /// The shape (`None` for a null dataspace, h5py's `Empty`), with the
    /// file generation it was read at: an edit (a resize) may change it.
    shape: Mutex<(u64, Option<Vec<u64>>)>,
    /// The chunk shape, for a chunked dataset.
    chunks: Option<Vec<u64>>,
    datatype: Datatype,
    /// Why the datatype cannot be read into numpy, if it cannot.
    conv: Result<Converter, String>,
}

impl PyDataset {
    pub(crate) fn new(
        py: Python<'_>,
        handle: Arc<Handle>,
        path: String,
        addr: u64,
        meta: DatasetMeta,
    ) -> Self {
        let generation = handle.generation();
        let conv = crate::no_panic(|| Converter::new(py, &meta.datatype, handle.offset_size))
            .map_err(|e| e.value(py).to_string());
        Self {
            handle,
            path,
            addr,
            shape: Mutex::new((generation, meta.shape)),
            chunks: meta.chunks,
            datatype: meta.datatype,
            conv,
        }
    }

    fn converter(&self) -> PyResult<&Converter> {
        self.conv
            .as_ref()
            .map_err(|msg| PyTypeError::new_err(format!("{}: {msg}", node::name(&self.path))))
    }

    /// The current shape: the one read at open, or re-read after an edit.
    fn dims(&self, py: Python<'_>) -> PyResult<Option<Vec<u64>>> {
        let generation = self.handle.generation();
        {
            let cached = self.shape.lock().unwrap_or_else(PoisonError::into_inner);
            if cached.0 == generation {
                return Ok(cached.1.clone());
            }
        }
        let addr = self.addr;
        let null = self
            .shape
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .1
            .is_none();
        let shape = if null {
            None
        } else {
            Some(self.handle.with(py, |f| {
                f.dataset_at(addr)
                    .and_then(|ds| ds.shape())
                    .map_err(to_py_err)
            })?)
        };
        *self.shape.lock().unwrap_or_else(PoisonError::into_inner) = (generation, shape.clone());
        Ok(shape)
    }

    fn check_writable(&self) -> PyResult<()> {
        if self.handle.is_writable() {
            Ok(())
        } else {
            Err(PyOSError::new_err(format!(
                "{}: the file is open read-only; open it with mode 'r+' to change it",
                node::name(&self.path)
            )))
        }
    }

    /// Read the selection described by `plan` into a numpy array.
    fn read_plan<'py>(
        &self,
        py: Python<'py>,
        plan: &Plan,
        dims: &[u64],
    ) -> PyResult<Bound<'py, PyAny>> {
        let conv = self.converter()?;
        let out_shape = plan.out_shape();

        let arr = if plan.is_empty() {
            conv.empty(py, &out_shape)?
        } else {
            let (vl, elem_size, unit) = (conv.is_vl(), conv.elem_size, conv.vl_unit);
            let list_axis = plan.list_axis();
            let chunk_len = match (&self.chunks, list_axis) {
                (Some(c), Some(a)) => c.get(a).copied(),
                _ => None,
            };
            let (reads, list_axis) = plan.reads(dims, chunk_len, elem_size);
            let read_shape = plan.read_shape();
            let handle = &*self.handle;
            let addr = self.addr;
            // Everything below touches only Rust data: release the GIL.
            let read = |file: &File| -> Result<Elements, ReadError> {
                let ds = file.dataset_at(addr)?;
                let mut blocks = Vec::with_capacity(reads.len());
                for read in reads {
                    let raw = ds.read_selection(&read.sel)?;
                    let mut shape = read.shape;
                    let want = shape.iter().product::<usize>() * elem_size;
                    if raw.len() != want {
                        return Err(ReadError::Other(format!(
                            "read {} bytes, expected {want}",
                            raw.len()
                        )));
                    }
                    let raw = match (&read.pick, list_axis) {
                        (Some(pick), Some(axis)) => {
                            let kept = select::gather_along(&raw, &shape, axis, pick, elem_size);
                            shape[axis] = pick.len();
                            kept
                        }
                        _ => raw,
                    };
                    blocks.push((raw, shape));
                }
                // Several blocks only for a list index: join their bytes
                // (every byte of every element, padding included) along
                // that axis before anything becomes numpy.
                let raw = match (blocks.len(), list_axis) {
                    (1, _) => blocks.pop().expect("one block").0,
                    (_, Some(axis)) => select::join_along(&blocks, axis, elem_size),
                    _ => {
                        return Err(ReadError::Other(
                            "several reads without an index list".into(),
                        ));
                    }
                };
                if !vl {
                    return Ok(Elements::Bytes(raw));
                }
                let sb = file.superblock();
                let n = read_shape.iter().product();
                resolve_vl(
                    file.storage(),
                    &raw,
                    n,
                    sb.offset_size,
                    sb.length_size,
                    unit,
                )
                .map(Elements::Vl)
                .map_err(ReadError::Vl)
            };
            let data = py
                .detach(|| {
                    handle
                        .with_detached(|f| {
                            Ok(
                                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| read(f)))
                                    .unwrap_or_else(|p| {
                                        Err(ReadError::Panic(crate::panic_text(&*p)))
                                    }),
                            )
                        })
                        .unwrap_or_else(|e| Err(ReadError::Py(e)))
                })
                .map_err(|e| e.into_py(&self.path))?;
            let joined = conv.to_array(py, data, &read_shape, false)?;
            // Drop the axes indexed by an integer (length 1 in the blocks).
            let mut shape = out_shape.clone();
            if let crate::convert::Layout::Subarray(sub) = &conv.layout {
                shape.extend_from_slice(sub);
            }
            joined.call_method1("reshape", (PyTuple::new(py, shape)?,))?
        };

        let arr = select_fields(py, arr, &plan.fields)?;
        if plan.scalar {
            return arr.get_item(PyTuple::empty(py));
        }
        Ok(arr)
    }
}

/// An error from the read closure, turned into a Python error with the GIL.
enum ReadError {
    Lib(clawhdf5_rs::Error),
    Vl(VlError),
    Py(PyErr),
    Other(String),
    Panic(String),
}

impl From<clawhdf5_rs::Error> for ReadError {
    fn from(e: clawhdf5_rs::Error) -> Self {
        ReadError::Lib(e)
    }
}

impl ReadError {
    fn into_py(self, path: &str) -> PyErr {
        match self {
            ReadError::Lib(e) => to_py_err(e),
            ReadError::Vl(e) => e.into_py(&node::name(path)),
            ReadError::Py(e) => e,
            ReadError::Other(msg) => PyValueError::new_err(format!("{}: {msg}", node::name(path))),
            ReadError::Panic(msg) => crate::InternalError::new_err(format!(
                "{}: clawhdf5 internal error (please report it): {msg}",
                node::name(path)
            )),
        }
    }
}

/// Keep only the named compound fields, as h5py's `ds['x']` / `ds['x', 'y']`.
fn select_fields<'py>(
    py: Python<'py>,
    arr: Bound<'py, PyAny>,
    fields: &[String],
) -> PyResult<Bound<'py, PyAny>> {
    if fields.is_empty() {
        return Ok(arr);
    }
    let names = arr.getattr("dtype")?.getattr("names")?;
    if names.is_none() {
        return Err(PyValueError::new_err(
            "Field names only allowed for compound types",
        ));
    }
    let names: Vec<String> = names.extract()?;
    for f in fields {
        if !names.contains(f) {
            return Err(PyValueError::new_err(format!(
                "Field {f} does not appear in this type."
            )));
        }
    }
    let np = py.import("numpy")?;
    if let [one] = fields {
        return np.call_method1("ascontiguousarray", (arr.get_item(one)?,));
    }
    let picked = arr.get_item(PyList::new(py, fields)?)?;
    py.import("numpy.lib.recfunctions")?
        .call_method1("repack_fields", (picked,))
}

#[pymethods]
impl PyDataset {
    /// The shape of the dataset (`None` for an empty/null dataspace).
    #[getter]
    fn shape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        match self.dims(py)? {
            Some(s) => Ok(PyTuple::new(py, s)?.into_any()),
            None => Ok(py.None().into_bound(py)),
        }
    }

    /// The maximum shape (`None` per unlimited dimension), like h5py.
    #[getter]
    fn maxshape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let Some(shape) = self.dims(py)? else {
            return Ok(py.None().into_bound(py));
        };
        let addr = self.addr;
        let max = self
            .handle
            .with(py, |f| {
                f.dataset_at(addr)
                    .and_then(|ds| ds.max_dimensions())
                    .map_err(to_py_err)
            })?
            .unwrap_or(shape);
        let items: Vec<Option<u64>> = max
            .into_iter()
            .map(|d| (d != u64::MAX).then_some(d))
            .collect();
        Ok(PyTuple::new(py, items)?.into_any())
    }

    /// The chunk shape, or `None` for a dataset that is not chunked.
    #[getter]
    fn chunks<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        match &self.chunks {
            Some(c) => Ok(PyTuple::new(py, c)?.into_any()),
            None => Ok(py.None().into_bound(py)),
        }
    }

    /// The dataset's numpy dtype, as h5py reports it.
    #[getter]
    fn dtype<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        Ok(self.converter()?.dtype.bind(py).clone())
    }

    #[getter]
    fn ndim(&self, py: Python<'_>) -> PyResult<usize> {
        Ok(self.dims(py)?.map_or(0, |s| s.len()))
    }

    /// Number of elements (`None` for an empty/null dataspace, as h5py).
    #[getter]
    fn size(&self, py: Python<'_>) -> PyResult<Option<u64>> {
        Ok(self.dims(py)?.map(|s| s.iter().product()))
    }

    /// The dataset's full name, e.g. `/group/data`.
    #[getter]
    fn name(&self) -> String {
        node::name(&self.path)
    }

    /// The dataset's attributes (dict-like; writable in a file opened with
    /// `'r+'`).
    #[getter]
    fn attrs(&self, py: Python<'_>) -> PyResult<PyAttrs> {
        PyAttrs::read(py, Arc::clone(&self.handle), self.addr, &self.path)
    }

    /// Read with h5py indexing: integers, slices with positive steps,
    /// `...`, one increasing list of integers, and compound field names.
    /// A selection whose bounding box covers at most half the dataset reads
    /// only the chunks (or contiguous rows) it overlaps.
    fn __getitem__<'py>(
        &self,
        py: Python<'py>,
        key: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let Some(dims) = self.dims(py)? else {
            let is_empty_tuple = key.cast::<PyTuple>().is_ok_and(|t| t.is_empty());
            let is_ellipsis = key.is_instance_of::<pyo3::types::PyEllipsis>();
            if is_empty_tuple || is_ellipsis {
                let empty = PyEmpty::new(self.converter()?.dtype.clone_ref(py));
                return Ok(empty.into_pyobject(py)?.into_any());
            }
            return Err(PyValueError::new_err("Empty datasets cannot be sliced"));
        };
        let plan = select::parse(key, &dims)?;
        self.read_plan(py, &plan, &dims)
    }

    /// Write with h5py indexing (file opened with `'r+'`): `ds[key] = value`.
    ///
    /// The key is what `ds[key]` reads (without compound field names). The
    /// value is converted to the dataset's dtype as h5py converts it (a
    /// numpy array as libhdf5 does, clipping out-of-range numbers; anything
    /// else through `numpy.asarray(value, dtype=ds.dtype)`), and broadcast
    /// to the selection as h5py broadcasts. The edit is written and synced
    /// before this returns; what the in-place editor cannot write raises
    /// `NotImplementedError` and leaves the file as it was.
    fn __setitem__(
        &self,
        py: Python<'_>,
        key: &Bound<'_, PyAny>,
        value: &Bound<'_, PyAny>,
    ) -> PyResult<()> {
        self.check_writable()?;
        let Some(dims) = self.dims(py)? else {
            return Err(PyNotImplementedError::new_err(
                "writing to an empty (null dataspace) dataset is not supported",
            ));
        };
        let plan = select::parse(key, &dims)?;
        if !plan.fields.is_empty() {
            return Err(PyNotImplementedError::new_err(
                "writing compound fields by name is not supported by clawhdf5's in-place editor; \
                 write whole elements",
            ));
        }
        let category = edit::category(&self.datatype)?;
        let conv = self.converter()?;
        let bytes = edit::dataset_bytes(
            py,
            value,
            conv.dtype.bind(py),
            category,
            &plan,
            self.chunks.as_deref(),
        )?;
        if plan.is_empty() {
            return Ok(());
        }
        let sel = edit::selection(&plan, &dims)?;
        let path = node::name(&self.path);
        self.handle
            .edit(py, |ed| ed.write_selection(&path, &sel, &bytes))
    }

    /// Change the dataset's shape (file opened with `'r+'`), as h5py's
    /// `Dataset.resize`: `ds.resize((100, 20))`, or `ds.resize(100, axis=0)`.
    /// Only chunked datasets, within their maximum shape; new elements read
    /// as the fill value.
    #[pyo3(signature = (size, axis=None))]
    fn resize(&self, py: Python<'_>, size: &Bound<'_, PyAny>, axis: Option<isize>) -> PyResult<()> {
        self.check_writable()?;
        let Some(dims) = self.dims(py)? else {
            return Err(PyTypeError::new_err("Empty datasets cannot be resized"));
        };
        if self.chunks.is_none() {
            return Err(PyTypeError::new_err("Only chunked datasets can be resized"));
        }
        let shape: Vec<u64> = match axis {
            Some(axis) => {
                let rank = dims.len();
                let a = usize::try_from(axis)
                    .ok()
                    .filter(|&a| a < rank)
                    .ok_or_else(|| {
                        PyValueError::new_err(format!(
                            "Invalid axis (0 to {} allowed)",
                            rank.saturating_sub(1)
                        ))
                    })?;
                let n: u64 = size.extract().map_err(|_| {
                    PyTypeError::new_err("Argument must be a single int if axis is specified")
                })?;
                let mut s = dims.clone();
                s[a] = n;
                s
            }
            // As h5py: without `axis` the size is a sequence (`tuple(size)`).
            None => size.extract().map_err(|_| {
                PyTypeError::new_err(format!(
                    "'{}' object is not iterable",
                    size.get_type()
                        .name()
                        .map(|n| n.to_string())
                        .unwrap_or_default()
                ))
            })?,
        };
        if shape.len() != dims.len() {
            return Err(PyValueError::new_err(format!(
                "new shape {shape:?} has {} dimensions, the dataset {}",
                shape.len(),
                dims.len()
            )));
        }
        let path = node::name(&self.path);
        self.handle.edit(py, |ed| ed.resize(&path, &shape))
    }

    /// `numpy.asarray(ds)` reads the whole dataset.
    #[pyo3(signature = (dtype=None, copy=None))]
    fn __array__<'py>(
        &self,
        py: Python<'py>,
        dtype: Option<&Bound<'py, PyAny>>,
        copy: Option<bool>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let _ = copy; // every read is a fresh array
        let Some(dims) = self.dims(py)? else {
            return Err(PyValueError::new_err("an empty dataset has no array value"));
        };
        let ellipsis = pyo3::types::PyEllipsis::get(py).to_owned().into_any();
        let plan = select::parse(&ellipsis, &dims)?;
        let arr = self.read_plan(py, &plan, &dims)?;
        match dtype {
            Some(dt) => arr.call_method1("astype", (dt,)),
            None => Ok(arr),
        }
    }

    fn __len__(&self, py: Python<'_>) -> PyResult<usize> {
        match self.dims(py)?.as_deref() {
            Some([first, ..]) => Ok(*first as usize),
            _ => Err(PyTypeError::new_err(
                "Attempt to take len() of scalar dataset",
            )),
        }
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let dtype = match &self.conv {
            Ok(c) => c
                .dtype
                .bind(py)
                .str()
                .map(|s| s.to_string())
                .unwrap_or_default(),
            Err(_) => format!("{:?}", self.datatype),
        };
        let shape = match self.dims(py) {
            Ok(Some(s)) => format!("{s:?}"),
            Ok(None) => "None".to_string(),
            Err(_) => "?".to_string(),
        };
        format!(
            "<HDF5 dataset \"{}\": shape {shape}, type \"{dtype}\">",
            node::name(&self.path)
        )
    }
}