lasprs 0.14.2

use super::error::*;
use hdf5_metno::{Dataset, Extents, File, H5Type, SimpleExtents};
// Leave this in place. The unused import is a false positive. It is used in a
// macro.
use ndarray::{ArrayView, ShapeBuilder};
use reinterpret::reinterpret_slice;
use snafu::prelude::*;
use std::path::{Path, PathBuf};
type Result<T> = std::result::Result<T, MeasurementError>;

use crate::{
    MeasurementMetadata,
    daq::{DataType, RawStreamData},
    file_attr_names::{
        INPUT_SIGNAL_CLIPPED_ATTR, MEASURED_INPUT_DATASET_NAME, MEASURED_PRECAPTURE_DATASET_NAME,
        MEASUREMENT_STILL_RUNNING,
    },
    measurement::{Measurement, MeasurementError, SharedMeasurement},
    tools::h5::{write_h5_attr_scalar, write_h5_attr_scalar_overwrite},
};

pub struct EmptyMeasurement;
pub struct MeasurementWithMeta {
    /// The main recorded dataset
    dset: Dataset,
    /// The pre-capture dataset, if applicable (negative time pre-capture)
    precapture_dset: Option<Dataset>,
    recorded_frames_ctr: usize,
    precapture_frames_ctr: usize,
    meta: MeasurementMetadata,
}

pub struct MeasurementWriter<S> {
    file: File,
    state: S,
}

impl MeasurementWriter<EmptyMeasurement> {
    /// Creates a new `MeasurementWriter` that writes to the given file path.
    ///
    /// # Arguments
    ///
    /// * `path` - The path to the file to write to.
    pub fn new(path: &Path) -> Result<Self> {
        let file = hdf5_metno::File::create_excl(path)
            .map_err(H5Error::from)
            .context(H5FileProblemSnafu {
                operation: format!("opening file {}", path.to_string_lossy()),
            })?;

        write_h5_attr_scalar(&file, MEASUREMENT_STILL_RUNNING, true)?;
        Ok(Self {
            file,
            state: EmptyMeasurement,
        })
    }
    /// Writes the measurement metadata to the file, initializes the audio
    /// dataset for writing measurement data. Initializes the recorded_block_ctr
    /// to 0.
    ///
    /// # Arguments
    ///
    /// * `meta` - The measurement metadata to write.
    pub fn write_meta(
        self,
        meta: MeasurementMetadata,
    ) -> Result<MeasurementWriter<MeasurementWithMeta>> {
        meta.writeToH5(&self.file)?;

        let dset = create_dataset_type(&self.file, &meta, MEASURED_INPUT_DATASET_NAME)?;

        Ok(MeasurementWriter {
            file: self.file,
            state: MeasurementWithMeta {
                dset,
                precapture_dset: None,
                recorded_frames_ctr: 0,
                precapture_frames_ctr: 0,
                meta,
            },
        })
    }
}

impl MeasurementWriter<MeasurementWithMeta> {
    /// Resize a dataset and append one block of data to it.
    /// Updates the frame counter.
    fn write_to_dataset(
        dset: &Dataset,
        frames_ctr: &mut usize,
        meta: &MeasurementMetadata,
        data: &RawStreamData,
    ) -> Result<()> {
        let nchannels = meta.nchannels();
        let new_nframes = data.nsamples() / nchannels;

        dset.resize((*frames_ctr + new_nframes, nchannels))
            .map_err(H5Error::from)
            .context(H5FileProblemSnafu {
                operation: "resizing dataset",
            })?;
        append_to_dset(dset, *frames_ctr, meta, data)?;
        *frames_ctr += new_nframes;
        Ok(())
    }

    /// Write a block of raw stream data to the measurement dataset to record it
    /// in the file. This is the main data writing method.
    pub fn write(&mut self, data: &RawStreamData) -> Result<()> {
        let Self {
            state:
                MeasurementWithMeta {
                    dset,
                    recorded_frames_ctr,
                    meta,
                    ..
                },
            ..
        } = self;
        Self::write_to_dataset(dset, recorded_frames_ctr, meta, data)
    }

    /// Write a block of pre-capture data to the pre-capture dataset.
    /// The dataset is created on first call. All pre-capture data must
    /// be written before any live data is written via `write()`.
    pub fn write_precapture(&mut self, data: &RawStreamData) -> Result<()> {
        // Ensure the pre-capture dataset exists, creating it lazily.
        if self.state.precapture_dset.is_none() {
            let dset = create_dataset_type(
                &self.file,
                &self.state.meta,
                MEASURED_PRECAPTURE_DATASET_NAME,
            )?;
            self.state.precapture_dset = Some(dset);
        }

        Self::write_to_dataset(
            self.state.precapture_dset.as_ref().unwrap(),
            &mut self.state.precapture_frames_ctr,
            &self.state.meta,
            data,
        )
    }
    /// Returns `true` if no blocks have been stored yet.
    pub fn is_empty(&self) -> bool {
        self.frames_recorded() == 0
    }
    #[inline]
    /// Returns the number of blocks recorded so far.
    pub fn frames_recorded(&self) -> usize {
        self.state.recorded_frames_ctr
    }
    /// Finishes the measurement by closing the file and returning the measurement.
    ///
    /// # Arguments
    ///
    /// * `clipped` - Whether the input signal was clipped during recording. This is stored as an
    ///   attribute in the HDF5 file.
    pub fn finish(self, clipped: bool) -> Result<SharedMeasurement> {
        let Self {
            file,
            state:
                MeasurementWithMeta {
                    dset,
                    precapture_dset,
                    ..
                },
        } = self;
        let path = PathBuf::from(file.filename());

        if clipped {
            write_h5_attr_scalar_overwrite(&file, INPUT_SIGNAL_CLIPPED_ATTR, clipped)?;
        }
        write_h5_attr_scalar_overwrite(&file, MEASUREMENT_STILL_RUNNING, false)?;
        // Drop the datasets and file before returning, as the measurement is now
        // complete, and the data is flushed to disk.
        drop(dset);
        drop(precapture_dset);
        drop(file);

        Measurement::from_file(&path)
    }
}

/// Create a resizable HDF5 dataset with the given name, using the
/// metadata's data type, sample rate, and channel count.
fn create_dataset_type(file: &File, meta: &MeasurementMetadata, name: &str) -> Result<Dataset> {
    match meta.dataType() {
        DataType::I8 => create_dataset_type_with_name::<i8>(file, meta, name),
        DataType::I16 => create_dataset_type_with_name::<i16>(file, meta, name),
        DataType::I32 => create_dataset_type_with_name::<i32>(file, meta, name),
        DataType::I24 => create_dataset_type_with_name::<i32>(file, meta, name),
        DataType::F32 => create_dataset_type_with_name::<f32>(file, meta, name),
        DataType::F64 => create_dataset_type_with_name::<f64>(file, meta, name),
    }
}

/// Create a resizable HDF5 dataset with the given name, using the
/// metadata's data type, sample rate, and channel count.
fn create_dataset_type_with_name<T>(
    file: &File,
    meta: &MeasurementMetadata,
    name: &str,
) -> Result<Dataset>
where
    T: H5Type,
{
    let nchannels = meta.nchannels();
    let extents = Extents::new(vec![0, nchannels]).resizable();

    match file
        .new_dataset::<T>()
        .chunk((10.max(*meta.samplerate() as usize / 10), nchannels))
        .shape(extents)
        .create(name)
    {
        Ok(dset) => {
            assert!(dset.is_resizable());
            Ok(dset)
        }
        err @ Err(_) => err.map_err(H5Error::from).context(H5FileProblemSnafu {
            operation: format!("creating dataset '{name}'"),
        }),
    }
}

#[inline]
fn append_to_dset(
    ds: &Dataset,
    start_frame: usize,
    meta: &MeasurementMetadata,
    data: &RawStreamData,
) -> Result<()> {
    // Macro to generate repetitive match arms for different data types
    macro_rules! write_data_slice {
        ($dat:expr, $ds:expr, $ctr:expr, $nchannels:expr) => {{
            let view = ArrayView::from_shape(
                ($dat.len() / $nchannels, $nchannels).strides(($nchannels, 1)),
                &$dat,
            )
            .unwrap();
            $ds.write_slice(&view, (start_frame.., ..))
                .map_err(H5Error::from)
                .context(H5FileProblemSnafu {
                    operation: "writing data to measurement file",
                })?;
        }};
    }

    let nchannels = meta.nchannels();
    match data {
        RawStreamData::Datai8(dat) => write_data_slice!(dat, ds, ctr, nchannels),
        RawStreamData::Datai16(dat) => write_data_slice!(dat, ds, ctr, nchannels),
        RawStreamData::Datai24(dat) => {
            assert!(std::mem::size_of::<dasp_sample::I24>() == std::mem::size_of::<i32>());
            // SAFETY: I24 stores as i32 as inner type
            let dat = unsafe { reinterpret_slice::<_, i32>(dat) };
            write_data_slice!(dat, ds, ctr, nchannels);
        }
        RawStreamData::Datai32(dat) => write_data_slice!(dat, ds, ctr, nchannels),
        RawStreamData::Dataf32(dat) => write_data_slice!(dat, ds, ctr, nchannels),
        RawStreamData::Dataf64(dat) => write_data_slice!(dat, ds, ctr, nchannels),
    };
    Ok(())
}