lasprs 0.6.7

use super::*;
use crate::config::Flt;
use crate::rt::SimpleClipDetector;
use anyhow::{anyhow, bail, Error, Result};
use clap::builder::OsStr;
use crossbeam::atomic::AtomicCell;
use hdf5::types::{VarLenArray, VarLenUnicode};
use hdf5::{dataset, datatype, Dataset, File, H5Type};
use parking_lot::Mutex;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
use std::sync::Arc;
use std::thread::{spawn, JoinHandle};
use std::time::Duration;
use streamdata::*;
use streammgr::*;
use streammsg::InStreamMsg;
use strum::EnumMessage;

/// Status of a recording
#[cfg_attr(feature = "python-bindings", pyclass(get_all, eq))]
#[derive(Clone, Debug, PartialEq)]
pub enum RecordStatus {
    /// Nothing to update
    NoUpdate {},

    /// Not yet started, waiting for first msg
    Idle {},

    /// Waiting for start delay to be processed.
    Waiting {},

    /// Recording in progress
    Recording {
        /// Whether a clip has already happened during recording
        clipped: bool,
        /// The amount of time that is currently recorded
        recorded: Flt,
        /// The percentage done [0 to 100%]. Stays at 0% for infinite duration
        /// recordings
        pct_done: Flt,
    },

    /// Recording finished
    Finished {
        /// Whether there occured a signal clip during the recording. Can only
        /// be detected when the flag 'record_clip`
        clipped: bool,
        /// Possibly, the recording finished, but an error happened, due to
        /// which the data is invalid.
        error: Option<String>,
    },
}

/// Settings used to start a recording.
#[derive(Clone)]
#[cfg_attr(feature = "python-bindings", pyclass(get_all))]
pub struct RecordSettings {
    /// File name to record to.
    pub filename: PathBuf,

    /// The recording time. Set to 0 to perform indefinite recording
    pub duration: Duration,

    /// The delay to wait before adding data
    pub startDelay: Duration,
}
impl RecordSettings {
    /// Create new record settings. Convenience wrapper to fill in fields in
    /// right form. Start delay is optional
    ///
    /// * args:
    ///     filename: Name of file to record to
    ///     duration: How long recording should be. Zero means record indefinitely.
    ///     startDelay: Optional start delay.
    pub fn new<T, U>(filename: T, duration: U, startDelay: Option<U>) -> RecordSettings
    where
        T: Into<PathBuf>,
        U: Into<Duration> + Default,
    {
        RecordSettings {
            filename: filename.into(),
            duration: duration.into(),
            startDelay: startDelay
                .map(|s| s.into())
                .unwrap_or_else(|| Duration::ZERO),
        }
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", pymethods)]
impl RecordSettings {
    #[new]
    #[pyo3(signature=(filename, duration=None, startDelay=None))]
    fn new_py(filename: String, duration: Option<Flt>, startDelay: Option<Flt>) -> PyResult<Self> {
        let startDelay = if let Some(delay) = startDelay {
            let delay = Duration::try_from_secs_f64(delay)
                .map_err(|e| anyhow! {"Invalid start delay specified: {e}"})?;
            Some(delay)
        } else {
            None
        };
        let duration = if let Some(duration) = duration {
            Duration::try_from_secs_f64(duration)
                .map_err(|e| anyhow! {"Invalid duration specified: {e}"})?
        } else {
            Duration::ZERO
        };
        Ok(RecordSettings::new(filename, duration, startDelay))
    }
}

/// This struct lets a recording run on a stream, waits till the first data arrives and records for a given period of time. Usage:
///
/// ```
/// use lasprs::daq::{RecordSettings, StreamMgr, Recording};
/// use std::time::Duration;
///
/// fn main() -> anyhow::Result<()> {
/// let mut smgr = StreamMgr::new();
/// smgr.startDefaultInputStream()?;
///
/// // Create record settings
/// let settings = RecordSettings::new(
///  "test.h5",
///  Duration::from_millis(100),
///  None,
/// );
/// let rec = Recording::new(settings, &mut smgr)?;
/// Ok(())
/// }
/// ```
#[cfg_attr(feature = "python-bindings", pyclass)]
pub struct Recording {
    // Recording settings
    settings: RecordSettings,

    // Stop the recording. This stops the thread
    stopThread: Arc<AtomicBool>,

    // Obtain status from thread.
    status_from_thread: Arc<AtomicCell<RecordStatus>>,

    // Stores latest status from thread, if no update comes from status_from_thread
    last_status: RecordStatus,
}

impl Recording {
    fn create_dataset_type<T>(file: &File, meta: &StreamMetaData) -> Result<Dataset>
    where
        T: H5Type,
    {
        let bs = meta.framesPerBlock;
        let nch = meta.nchannels();
        match file
            .new_dataset::<T>()
            .chunk((1, bs, nch))
            .shape((1.., bs, nch))
            // .deflate(3)
            .create("audio")
        {
            Ok(f) => Ok(f),
            Err(e) => bail!("{}", e),
        }
    }

    fn create_dataset(file: &File, meta: &StreamMetaData) -> Result<Dataset> {
        match meta.rawDatatype {
            DataType::I8 => Recording::create_dataset_type::<i8>(file, meta),
            DataType::I16 => Recording::create_dataset_type::<i16>(file, meta),
            DataType::I32 => Recording::create_dataset_type::<i32>(file, meta),
            DataType::F32 => Recording::create_dataset_type::<f32>(file, meta),
            DataType::F64 => Recording::create_dataset_type::<f64>(file, meta),
        }
    }

    fn write_hdf5_attr_scalar<T>(file: &File, name: &str, val: T) -> Result<()>
    where
        T: H5Type,
    {
        let attr = file.new_attr::<T>().create(name)?;
        attr.write_scalar(&val)?;
        Ok(())
    }
    fn write_hdf5_attr_list<T>(file: &File, name: &str, val: &[T]) -> Result<()>
    where
        T: H5Type,
    {
        let attr = file.new_attr::<T>().shape([val.len()]).create(name)?;
        attr.write(&val)?;
        Ok(())
    }

    #[inline]
    fn append_to_dset(
        ds: &Dataset,
        ctr: usize,
        data: &InStreamData,
        framesPerBlock: usize,
        nchannels: usize,
    ) -> Result<()> {
        match data.getRaw() {
            // The code below uses ndarray 0.15.6, which is the version required
            // to communicate with rust-hdf5. It requires input to be C-ordered,
            // or interleaved. This happens to be the default for ndarray as
            // well.
            RawStreamData::Datai8(dat) => {
                let arr =
                    ndarray15p6::ArrayView2::<i8>::from_shape((framesPerBlock, nchannels), dat)?;
                ds.write_slice(arr, (ctr, .., ..))?;
            }
            RawStreamData::Datai16(dat) => {
                let arr =
                    ndarray15p6::ArrayView2::<i16>::from_shape((framesPerBlock, nchannels), dat)?;
                ds.write_slice(arr, (ctr, .., ..))?;
            }
            RawStreamData::Datai32(dat) => {
                let arr =
                    ndarray15p6::ArrayView2::<i32>::from_shape((framesPerBlock, nchannels), dat)?;
                ds.write_slice(arr, (ctr, .., ..))?;
            }
            RawStreamData::Dataf32(dat) => {
                let arr =
                    ndarray15p6::ArrayView2::<f32>::from_shape((framesPerBlock, nchannels), dat)?;
                ds.write_slice(arr, (ctr, .., ..))?;
            }
            RawStreamData::Dataf64(dat) => {
                let arr =
                    ndarray15p6::ArrayView2::<f64>::from_shape((framesPerBlock, nchannels), dat)?;
                ds.write_slice(arr, (ctr, .., ..))?;
            }
        }
        Ok(())
    }

    /// Start a new recording
    ///
    /// # Arguments
    ///
    /// * setttings: The settings to use for the recording
    /// * smgr: Stream manager to use to start the recording
    ///
    pub fn new(mut settings: RecordSettings, mgr: &mut StreamMgr) -> Result<Recording> {
        // Append extension if not yet there
        match settings.filename.extension() {
            Some(a) if a == "h5" => {}
            None | Some(_) => {
                settings.filename =
                    (settings.filename.to_string_lossy().to_string() + ".h5").into();
            }
        };

        // Detector for signal clipping
        let clipdetector = SimpleClipDetector::new(mgr);

        let stopThread = Arc::new(AtomicBool::new(false));
        let stopThread_clone = stopThread.clone();

        // Fail if filename already exists
        if settings.filename.exists() {
            bail!(
                "Filename '{}' already exists in filesystem",
                settings.filename.to_string_lossy()
            );
        }
        let settings_clone = settings.clone();

        let status = Arc::new(AtomicCell::new(RecordStatus::Idle {}));
        let status_clone = status.clone();

        let (tx, rx) = crossbeam::channel::unbounded();
        mgr.addInQueue(tx.clone());

        // The thread doing the actual work
        rayon::spawn(move || {
            let recording_closure = || {
                let file = File::create(settings.filename)?;

                let firstmsg = match rx.recv() {
                    Ok(msg) => msg,
                    Err(_) => bail!("Queue handle error"),
                };

                let meta = match firstmsg {
                    InStreamMsg::StreamStarted(meta) => meta,
                    _ => bail!("Recording failed. Missed stream metadata message."),
                };

                // Samplerate, block size, number of channels
                Recording::write_hdf5_attr_scalar(&file, "samplerate", meta.samplerate)?;
                Recording::write_hdf5_attr_scalar(&file, "nchannels", meta.nchannels())?;
                Recording::write_hdf5_attr_scalar(&file, "blocksize", meta.framesPerBlock)?;

                // Store sensitivity
                let sens: Vec<Flt> = meta.channelInfo.iter().map(|ch| ch.sensitivity).collect();
                Recording::write_hdf5_attr_list(&file, "sensitivity", &sens)?;

                // Timestamp
                use chrono::DateTime;
                let now_utc = chrono::Utc::now();
                let timestamp = now_utc.timestamp();
                Recording::write_hdf5_attr_scalar(&file, "time", timestamp)?;

                // Create UUID for measurement
                use hdf5::types::VarLenUnicode;
                let uuid = uuid::Uuid::new_v4();
                let uuid_unicode: VarLenUnicode =
                    VarLenUnicode::from_str(&uuid.to_string()).unwrap();
                Recording::write_hdf5_attr_scalar(&file, "UUID", uuid_unicode)?;

                // Channel names
                let chnames: Vec<VarLenUnicode> = meta
                    .channelInfo
                    .iter()
                    .map(|ch| VarLenUnicode::from_str(&ch.name).unwrap())
                    .collect();
                let chname_attr = file
                    .new_attr::<VarLenUnicode>()
                    .shape([chnames.len()])
                    .create("channelNames")?;
                chname_attr.write(&chnames)?;

                // Create the dataset
                let ds = Recording::create_dataset(&file, &meta)?;

                let framesPerBlock = meta.framesPerBlock as usize;
                let mut wait_block_ctr = 0;
                // Indicate we are ready to rec!
                if settings.startDelay > Duration::ZERO {
                    status.store(RecordStatus::Waiting {});
                    let startdelay_s = settings.startDelay.as_micros() as Flt / 1e6;
                    wait_block_ctr =
                        (meta.samplerate as Flt * startdelay_s / framesPerBlock as Flt) as u32;
                } else {
                    status.store(RecordStatus::Recording {
                        clipped: clipdetector.hasClipped(),
                        recorded: 0.,
                        pct_done: 0.,
                    });
                }

                // Counter of stored blocks
                let mut recorded_ctr = 0;

                // Offset in stream
                let mut incoming_ctr = 0;

                // Flag indicating that the first RawStreamData package still has 
                // arrived
                let mut firstblockhasarrived = false;

                // Indicating the file is still empty (does not contain recorded data)
                let mut empty_file = true;

                let nchannels = meta.nchannels() as usize;
                'recloop: loop {
                    if stopThread.load(Relaxed) {
                        break 'recloop;
                    }
                    match rx.recv().unwrap() {
                        InStreamMsg::StreamError(e) => {
                            bail!("Recording failed due to stream error: {}.", e)
                        }
                        InStreamMsg::StreamStarted(_) => {
                            bail!("Stream started again?")
                        }
                        InStreamMsg::StreamStopped => {
                            // Early stop. User stopped it.
                            break 'recloop;
                        }
                        InStreamMsg::InStreamData(instreamdata) => {
                            if ! firstblockhasarrived {
                                // Toggle flag.
                                firstblockhasarrived = true;
                                // Initialize counter offset
                                incoming_ctr = instreamdata.ctr;
                            } else {
                                incoming_ctr += 1;
                            }
                            if instreamdata.ctr != incoming_ctr {
                                eprintln!("********** PACKAGES MISSED ***********");
                                bail!("Stream data blocks missed. Recording is invalid.")
                            }

                            if wait_block_ctr > 0 {
                                // We are still waiting
                                wait_block_ctr -= 1;
                                if wait_block_ctr == 0 {
                                    status.store(RecordStatus::Recording {
                                        clipped: clipdetector.hasClipped(),
                                        recorded: 0.,
                                        pct_done: 0.,
                                    });
                                }
                                continue 'recloop;
                            }

                            ds.resize((recorded_ctr + 1, framesPerBlock, nchannels))?;
                            Recording::append_to_dset(
                                &ds,
                                recorded_ctr,
                                &instreamdata,
                                framesPerBlock,
                                nchannels,
                            )?;
                            // Once we have added to the file, this flag is swapped
                            // and a file should only be deleted in case of an error.
                            empty_file = false;

                            // Recorded time rounded of to milliseconds.
                            let recorded_time = Duration::from_millis(
                                ((1000 * (recorded_ctr + 1) * framesPerBlock) as Flt
                                    / meta.samplerate) as u64,
                            );

                            if !settings.duration.is_zero() {
                                // Duration not equal to zero, meaning we record up to a
                                // certain duration.
                                if recorded_time >= settings.duration {
                                    break 'recloop;
                                }
                            }
                            // println!("\n... {} {} {}", recorded_time.as_millis(), meta.samplerate, framesPerBlock);
                            recorded_ctr += 1;
                            status.store(RecordStatus::Recording {
                                clipped: clipdetector.hasClipped(),
                                recorded: recorded_time.as_secs_f64(),
                                pct_done: if settings.duration > Duration::ZERO {
                                    100. * recorded_time.as_millis() as Flt
                                        / settings.duration.as_millis() as Flt
                                } else {
                                    0.
                                },
                            });
                        }
                    }
                } // end of 'recloop

                if empty_file {
                    bail!("Recording stopped before any data is stored.");
                }
                Ok(())
            };
            if let Err(e) = recording_closure() {
                status.store(RecordStatus::Finished {
                    clipped: clipdetector.hasClipped(),
                    error: Some(format!("{e}")),
                });
            } else {
                status.store(RecordStatus::Finished {
                    clipped: clipdetector.hasClipped(),
                    error: None,
                });
                // End of thread
            }
        });

        Ok(Recording {
            settings: settings_clone,
            stopThread: stopThread_clone,
            last_status: RecordStatus::NoUpdate {},
            status_from_thread: status_clone,
        })
    }

    // Delete recording file, should be done when something went wrong (an error
    // occured), or when cancel() is called, or when recording object is dropped
    // while thread is still running.
    fn deleteFile(&self) {
        // File should not be un use anymore, as thread is joined.
        // In case of error, we try to delete the file
        if let Err(e) = std::fs::remove_file(&self.settings.filename) {
            eprintln!("Recording failed, but file removal failed as well: {}", e);
        }
    }
}
#[cfg_attr(feature = "python-bindings", pymethods)]
impl Recording {
    #[cfg(feature = "python-bindings")]
    #[new]
    fn new_py(settings: RecordSettings, mgr: &mut StreamMgr) -> PyResult<Recording> {
        Ok(Recording::new(settings, mgr)?)
    }

    /// Get current record status
    pub fn status(&mut self) -> RecordStatus {
        // Update status due to normal messaging
        let status_from_thread = self.status_from_thread.swap(RecordStatus::NoUpdate {});
        match status_from_thread {
            RecordStatus::NoUpdate {} => {}
            _ => {
                self.last_status = status_from_thread;
            }
        }

        // Return latest status
        self.last_status.clone()
    }

    /// Stop existing recording early. At the current time, or st
    pub fn stop(&mut self) {
        // Stop thread and wait for message of finished
        self.stopThread.store(true, Relaxed);
        self.waitForThread();
        match self.status() {
            RecordStatus::Finished { error, .. } => {
                if error.is_some() {
                    // Delete file if there was an error
                    self.deleteFile();
                }
            }
            _ => {
                panic!("RecordStatus should be finished!");
            }
        }
    }

    /// Cancel recording. Deletes the recording file
    pub fn cancel(&mut self) {
        self.stopThread.store(true, Relaxed);
        self.waitForThread();
        self.deleteFile();
    }
    fn waitForThread(&mut self) {
        while !matches!(self.status(), RecordStatus::Finished { .. }) {
            std::thread::sleep(Duration::from_millis(10));
        }
    }
}

impl Drop for Recording {
    fn drop(&mut self) {
        // If we enter here, stop() or cancel() has not been called. In that
        // case, we cleanup here by cancelling the recording and deleting the file.
        if !matches!(self.status(), RecordStatus::Finished { .. }) {
            self.stopThread.store(true, Relaxed);
            self.waitForThread();
            self.deleteFile();
        }
    }
}