//! Lock-free single-producer, single-consumer latest-value mailboxes.

use core::sync::atomic::{AtomicU32, Ordering};

use kaptan_sensor::ImuSnapshot;
use kaptan_time::Timestamp;

/// A latest-value IMU mailbox.
///
/// One IMU worker publishes and one telemetry worker consumes. Publishing a
/// newer sample intentionally replaces an unread older sample: telemetry is a
/// best-effort observer and must never delay acquisition. The sequence counter
/// is a seqlock; odd values indicate a writer and even values identify a
/// stable snapshot.
pub struct LatestImuMailbox {
    sequence: AtomicU32,
    timestamp_low: AtomicU32,
    timestamp_high: AtomicU32,
    acceleration: [AtomicU32; 3],
    angular_rate: [AtomicU32; 3],
    temperature: AtomicU32,
}

impl LatestImuMailbox {
    /// Creates an empty mailbox.
    pub const fn new() -> Self {
        Self {
            sequence: AtomicU32::new(0),
            timestamp_low: AtomicU32::new(0),
            timestamp_high: AtomicU32::new(0),
            acceleration: [const { AtomicU32::new(0) }; 3],
            angular_rate: [const { AtomicU32::new(0) }; 3],
            temperature: AtomicU32::new(0),
        }
    }

    /// Publishes a sample, replacing any unread older value.
    pub fn publish(&self, sample: ImuSnapshot) {
        let sequence = self.sequence.fetch_add(1, Ordering::Relaxed);
        debug_assert_eq!(sequence & 1, 0, "SPSC producer must not reenter");

        let timestamp = sample.timestamp.as_microseconds();
        self.timestamp_low
            .store(timestamp as u32, Ordering::Relaxed);
        self.timestamp_high
            .store((timestamp >> 32) as u32, Ordering::Relaxed);
        store_vector(&self.acceleration, sample.acceleration_milligravity);
        store_vector(
            &self.angular_rate,
            sample.angular_rate_milliradians_per_second,
        );
        self.temperature.store(
            sample.temperature_centidegrees_celsius as u16 as u32,
            Ordering::Relaxed,
        );

        self.sequence
            .store(sequence.wrapping_add(2), Ordering::Release);
    }

    /// Returns the newest stable sample different from `observed_sequence`.
    pub fn take_newer(&self, observed_sequence: &mut u32) -> Option<ImuSnapshot> {
        let before = self.sequence.load(Ordering::Acquire);
        if before == *observed_sequence || before & 1 != 0 {
            return None;
        }

        let timestamp = u64::from(self.timestamp_low.load(Ordering::Relaxed))
            | (u64::from(self.timestamp_high.load(Ordering::Relaxed)) << 32);
        let acceleration = load_vector(&self.acceleration);
        let angular_rate = load_vector(&self.angular_rate);
        let temperature = self.temperature.load(Ordering::Relaxed) as u16 as i16;

        if self.sequence.load(Ordering::Acquire) != before {
            return None;
        }

        *observed_sequence = before;
        Some(ImuSnapshot::new(
            Timestamp::from_microseconds(timestamp),
            acceleration,
            angular_rate,
            temperature,
        ))
    }
}

impl Default for LatestImuMailbox {
    fn default() -> Self {
        Self::new()
    }
}

fn store_vector(destination: &[AtomicU32; 3], source: [i16; 3]) {
    for (destination, source) in destination.iter().zip(source) {
        destination.store(source as u16 as u32, Ordering::Relaxed);
    }
}

fn load_vector(source: &[AtomicU32; 3]) -> [i16; 3] {
    [
        source[0].load(Ordering::Relaxed) as u16 as i16,
        source[1].load(Ordering::Relaxed) as u16 as i16,
        source[2].load(Ordering::Relaxed) as u16 as i16,
    ]
}

#[cfg(test)]
mod tests {
    use super::LatestImuMailbox;
    use kaptan_sensor::ImuSnapshot;
    use kaptan_time::Timestamp;

    #[test]
    fn publishes_only_the_newest_unread_sample() {
        let mailbox = LatestImuMailbox::new();
        let mut observed = 0;
        mailbox.publish(ImuSnapshot::new(
            Timestamp::from_microseconds(1),
            [1; 3],
            [2; 3],
            3,
        ));
        mailbox.publish(ImuSnapshot::new(
            Timestamp::from_microseconds(2),
            [4; 3],
            [5; 3],
            6,
        ));

        let sample = mailbox.take_newer(&mut observed).unwrap();
        assert_eq!(sample.timestamp, Timestamp::from_microseconds(2));
        assert_eq!(sample.acceleration_milligravity, [4; 3]);
        assert!(mailbox.take_newer(&mut observed).is_none());
    }
}