//! MPU6000 register protocol and sample driver.

mod config;
mod error;
mod registers;
mod sample;
pub mod service;

use kaptan_sensor::imu::ImuSnapshot;
use kaptan_spi::{SpiCompletion, SpiDevice};
use kaptan_time::Timestamp;

pub use config::{AccelerometerRange, DigitalLowPassFilter, GyroscopeRange, Mpu6000Config};
pub use error::Error;
pub use service::{Mpu6000Service, Mpu6000ServiceEvent};

use registers::{
    ACCEL_CONFIG, ACCEL_XOUT_H, CLOCK_SOURCE_PLL_X_GYRO, CONFIG, DATA_READY_INTERRUPT,
    DEVICE_RESET, ENABLE_ALL_SENSOR_AXES, EXPECTED_WHO_AM_I, GYRO_CONFIG, I2C_INTERFACE_DISABLE,
    INT_ENABLE, INT_PIN_CFG, INTERRUPT_CLEAR_ON_ANY_READ, PWR_MGMT_1, PWR_MGMT_2, SAMPLE_BYTES,
    SMPLRT_DIV, USER_CTRL, WHO_AM_I,
};

const TRANSFER_CAPACITY: usize = SAMPLE_BYTES + 1;
const CONFIGURATION: [(u8, u8); 9] = [
    (PWR_MGMT_1, CLOCK_SOURCE_PLL_X_GYRO),
    (USER_CTRL, I2C_INTERFACE_DISABLE),
    (PWR_MGMT_2, ENABLE_ALL_SENSOR_AXES),
    (CONFIG, 0),
    (SMPLRT_DIV, 0),
    (ACCEL_CONFIG, 0),
    (GYRO_CONFIG, 0),
    (INT_PIN_CFG, INTERRUPT_CLEAR_ON_ANY_READ),
    (INT_ENABLE, DATA_READY_INTERRUPT),
];

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum State {
    Dormant,
    Resetting,
    WaitingForReset,
    ReadingIdentity,
    WritingConfiguration(usize),
    Ready,
    ReadingSample { timestamp: Timestamp },
}

/// Observable outcome of one SPI terminal event.
#[derive(Debug, Eq, PartialEq)]
pub enum Mpu6000Event<BusError> {
    /// The selected SPI interrupt had no terminal transfer event.
    NotPending,
    /// One initialization transfer completed; initialization continues.
    Initializing,
    /// The MPU6000 identity and configuration sequence completed.
    Initialized,
    /// One coherent IMU sample completed.
    Sample(ImuSnapshot),
    /// Initialization or acquisition failed.
    Fault(Error<BusError>),
}

/// Configured MPU6000 driver.
///
/// The driver owns fixed transfer buffers and advances only from external
/// interrupt events. It never waits for SPI status flags.
pub struct Mpu6000<Bus> {
    bus: Bus,
    config: Mpu6000Config,
    state: State,
    transmit: [u8; TRANSFER_CAPACITY],
    receive: [u8; TRANSFER_CAPACITY],
}

impl<Bus> Mpu6000<Bus> {
    /// Creates an inactive MPU6000 state machine.
    pub const fn new(bus: Bus, config: Mpu6000Config) -> Self {
        Self {
            bus,
            config,
            state: State::Dormant,
            transmit: [0; TRANSFER_CAPACITY],
            receive: [0; TRANSFER_CAPACITY],
        }
    }
}

impl<Bus: SpiDevice> Mpu6000<Bus> {
    /// Starts a device reset without waiting for SPI.
    ///
    /// The caller must subsequently call [`Self::advance_initialization`] only
    /// after the platform-selected reset-settling interval has elapsed.
    pub fn begin_initialization(&mut self) -> Result<(), Error<Bus::StartError>> {
        if self.state != State::Dormant {
            return Err(Error::NotReady);
        }

        self.prepare_write(PWR_MGMT_1, DEVICE_RESET);
        self.submit(2)?;
        self.state = State::Resetting;
        Ok(())
    }

    /// Starts identity verification after the required reset-settling interval.
    ///
    /// This operation only submits the next asynchronous SPI transfer; it
    /// never waits for a peripheral or sensor status flag.
    pub fn advance_initialization(&mut self) -> Result<(), Error<Bus::StartError>> {
        if self.state != State::WaitingForReset {
            return Ok(());
        }

        self.prepare_read(WHO_AM_I, 2);
        if let Err(error) = self.submit(2) {
            self.state = State::Dormant;
            return Err(error);
        }
        self.state = State::ReadingIdentity;
        Ok(())
    }

    /// Starts one fourteen-byte sensor read when initialization is complete.
    pub fn request_sample(&mut self, timestamp: Timestamp) -> Result<(), Error<Bus::StartError>> {
        if self.state != State::Ready {
            return Err(Error::NotReady);
        }

        self.prepare_read(ACCEL_XOUT_H, TRANSFER_CAPACITY);
        self.submit(TRANSFER_CAPACITY)?;
        self.state = State::ReadingSample { timestamp };
        Ok(())
    }

    /// Advances the state machine from an already-acknowledged DMA terminal event.
    pub fn on_transfer_complete(
        &mut self,
        completion: SpiCompletion,
    ) -> Mpu6000Event<Bus::StartError> {
        self.bus.complete_transfer(completion);
        match completion {
            SpiCompletion::NotPending => Mpu6000Event::NotPending,
            SpiCompletion::TransferFault => {
                self.state = State::Dormant;
                Mpu6000Event::Fault(Error::TransferFault)
            }
            SpiCompletion::Complete => self.advance_after_completion(),
        }
    }

    /// Returns whether an IMU data-ready event can start a sensor transfer.
    pub fn is_ready(&self) -> bool {
        self.state == State::Ready
    }

    fn advance_after_completion(&mut self) -> Mpu6000Event<Bus::StartError> {
        match self.state {
            State::Dormant => Mpu6000Event::NotPending,
            State::Resetting => {
                self.state = State::WaitingForReset;
                Mpu6000Event::Initializing
            }
            State::WaitingForReset => Mpu6000Event::NotPending,
            State::ReadingIdentity => {
                let identity = self.receive[1];
                if identity != EXPECTED_WHO_AM_I {
                    self.state = State::Dormant;
                    return Mpu6000Event::Fault(Error::UnexpectedIdentity(identity));
                }

                self.start_configuration(0)
            }
            State::WritingConfiguration(index) => {
                let next = index + 1;
                if next == CONFIGURATION.len() {
                    self.state = State::Ready;
                    Mpu6000Event::Initialized
                } else {
                    self.start_configuration(next)
                }
            }
            State::Ready => Mpu6000Event::NotPending,
            State::ReadingSample { timestamp } => {
                let mut bytes = [0; SAMPLE_BYTES];
                bytes.copy_from_slice(&self.receive[1..]);
                self.state = State::Ready;
                Mpu6000Event::Sample(sample::decode(&bytes, timestamp, self.config))
            }
        }
    }

    fn start_configuration(&mut self, index: usize) -> Mpu6000Event<Bus::StartError> {
        let (register, value) = CONFIGURATION[index];
        let value = match register {
            CONFIG => self.config.digital_low_pass_filter.register_value(),
            SMPLRT_DIV => self.config.sample_rate_divider,
            ACCEL_CONFIG => self.config.accelerometer.register_value(),
            GYRO_CONFIG => self.config.gyroscope.register_value(),
            _ => value,
        };
        self.prepare_write(register, value);
        match self.submit(2) {
            Ok(()) => {
                self.state = State::WritingConfiguration(index);
                Mpu6000Event::Initializing
            }
            Err(error) => {
                self.state = State::Dormant;
                Mpu6000Event::Fault(error)
            }
        }
    }

    fn prepare_read(&mut self, register: u8, length: usize) {
        self.transmit = [0; TRANSFER_CAPACITY];
        self.receive = [0; TRANSFER_CAPACITY];
        self.transmit[0] = register | 0x80;
        debug_assert!(length <= TRANSFER_CAPACITY);
    }

    fn prepare_write(&mut self, register: u8, value: u8) {
        self.transmit = [0; TRANSFER_CAPACITY];
        self.receive = [0; TRANSFER_CAPACITY];
        self.transmit[..2].copy_from_slice(&[register, value]);
    }

    fn submit(&mut self, length: usize) -> Result<(), Error<Bus::StartError>> {
        let (bus, transmit, receive) = (&mut self.bus, &self.transmit, &mut self.receive);
        bus.start_transfer(&transmit[..length], &mut receive[..length])
            .map_err(Error::Start)
    }
}

#[cfg(test)]
mod tests {
    use super::{Error, Mpu6000, Mpu6000Config, Mpu6000Event, TRANSFER_CAPACITY};
    use kaptan_spi::{SpiCompletion, SpiDevice};
    use kaptan_time::Timestamp;

    #[derive(Debug, Eq, PartialEq)]
    enum FakeError {
        Start,
    }

    struct FakeBus {
        receive: [u8; TRANSFER_CAPACITY],
        length: usize,
        fail_start: bool,
    }

    impl FakeBus {
        fn ready() -> Self {
            Self {
                receive: [0; TRANSFER_CAPACITY],
                length: 0,
                fail_start: false,
            }
        }
    }

    impl SpiDevice for FakeBus {
        type StartError = FakeError;

        fn start_transfer(&mut self, _: &[u8], receive: &mut [u8]) -> Result<(), Self::StartError> {
            if self.fail_start {
                return Err(FakeError::Start);
            }

            self.length = receive.len();
            receive.copy_from_slice(&self.receive[..self.length]);
            Ok(())
        }

        fn complete_transfer(&mut self, _: SpiCompletion) {}
    }

    #[test]
    fn configures_a_verified_device_without_waiting() {
        let mut bus = FakeBus::ready();
        bus.receive[1] = 0x68;
        let mut driver = Mpu6000::new(bus, Mpu6000Config::default());

        driver.begin_initialization().unwrap();
        assert_eq!(
            driver.on_transfer_complete(SpiCompletion::Complete),
            Mpu6000Event::Initializing
        );
        driver.advance_initialization().unwrap();
        for _ in 0..9 {
            assert_eq!(
                driver.on_transfer_complete(SpiCompletion::Complete),
                Mpu6000Event::Initializing
            );
        }
        assert_eq!(
            driver.on_transfer_complete(SpiCompletion::Complete),
            Mpu6000Event::Initialized
        );
        assert!(driver.is_ready());
    }

    #[test]
    fn rejects_an_unexpected_identity() {
        let mut bus = FakeBus::ready();
        bus.receive[1] = 0x70;
        let mut driver = Mpu6000::new(bus, Mpu6000Config::default());

        driver.begin_initialization().unwrap();
        let _ = driver.on_transfer_complete(SpiCompletion::Complete);
        driver.advance_initialization().unwrap();
        assert!(matches!(
            driver.on_transfer_complete(SpiCompletion::Complete),
            Mpu6000Event::Fault(Error::UnexpectedIdentity(0x70))
        ));
    }

    #[test]
    fn converts_a_complete_burst_to_core_units() {
        let burst = [
            0x20, 0x00, 0xE0, 0x00, 0x00, 0x00, // accel: 1g, -1g, 0g
            0x00, 0x00, // 36.53 C
            0x40, 0x10, 0xBF, 0xF0, 0x00, 0x00, // gyro: 500, -500, 0 dps
        ];
        let mut bus = FakeBus::ready();
        bus.receive[1] = 0x68;
        let mut driver = Mpu6000::new(bus, Mpu6000Config::default());

        driver.begin_initialization().unwrap();
        let _ = driver.on_transfer_complete(SpiCompletion::Complete);
        driver.advance_initialization().unwrap();
        for _ in 0..9 {
            let _ = driver.on_transfer_complete(SpiCompletion::Complete);
        }
        let _ = driver.on_transfer_complete(SpiCompletion::Complete);
        driver.bus.receive[1..].copy_from_slice(&burst);
        driver
            .request_sample(Timestamp::from_microseconds(42))
            .unwrap();
        let Mpu6000Event::Sample(sample) = driver.on_transfer_complete(SpiCompletion::Complete)
        else {
            panic!("sample transfer must complete");
        };

        assert_eq!(sample.acceleration_milligravity, [1_000, -1_000, 0]);
        assert_eq!(
            sample.angular_rate_milliradians_per_second,
            [8_726, -8_726, 0]
        );
        assert_eq!(sample.temperature_centidegrees_celsius, 3_653);
    }

    #[test]
    fn reports_a_route_transfer_fault() {
        let mut bus = FakeBus::ready();
        bus.receive[1] = 0x68;
        let mut driver = Mpu6000::new(bus, Mpu6000Config::default());

        driver.begin_initialization().unwrap();
        assert!(matches!(
            driver.on_transfer_complete(SpiCompletion::TransferFault),
            Mpu6000Event::Fault(Error::TransferFault)
        ));
    }
}