//! Complete outbound MAVLink data flow.
//!
//! This crate owns fixed-rate scheduling, mailbox consumption, vertical-data
//! freshness, MAVLink encoding, and bounded queue classification. A platform
//! supplies only mailbox references, a declared MAVLink state, arming state,
//! and a transport sink.

#![no_std]

use kaptan_dma::TxCompletion;
use kaptan_err_propagate::KaptanError;
use kaptan_mavlink::message::MessageData;
use kaptan_mavlink::{
    MavlinkFrame,
    encoder::{DistanceSensorConfig, MavlinkEncoder},
    message::MavState,
    transmitter::{EnqueueResult, FrameClass},
};
use kaptan_sensor::{optical_flow::OpticalFlowSnapshot, range::RangeSnapshot};

mod mailbox;

pub use mailbox::{LatestImuMailbox, LatestOpticalFlowMailbox, LatestRangeMailbox};

/// Validated publication schedule consumed by the outbound data channel.
///
/// This is a derived runtime input, not a persisted configuration root. Board
/// packages own the factory values and the platform constructs this view only
/// after validating the active user configuration against immutable hardware.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct DataChannelRates {
    /// Frequency of the platform service tick.
    pub service_hz: usize,
    /// Raw-IMU publication rate.
    pub raw_imu_hz: usize,
    /// Range publication rate; `None` disables the stream.
    pub range_hz: Option<usize>,
    /// Optical-flow publication rate; `None` disables the stream.
    pub optical_flow_hz: Option<usize>,
    /// Heartbeat publication rate.
    pub heartbeat_hz: usize,
}

/// Platform transport used by a data channel.
///
/// The channel decides exactly which encoded frame is admitted and with which
/// delivery class. The platform retains only DMA endpoint ownership.
pub trait FrameSink {
    /// Admits one fully encoded frame to the platform's bounded transport.
    fn send(&mut self, class: FrameClass, frame: MavlinkFrame) -> EnqueueResult;

    /// Applies one terminal DMA interrupt and starts the next retained frame
    /// when possible.
    fn on_dma_complete(&mut self, completion: TxCompletion);
}

/// Mailboxes read by one data-channel worker.
#[derive(Clone, Copy)]
pub struct DataSources<'a> {
    /// Latest IMU observation from the acquisition worker.
    pub imu: &'a LatestImuMailbox,
    /// Latest range observation from either supported vertical route.
    pub range: &'a LatestRangeMailbox,
    /// Latest optical-flow observation from MTF-01.
    pub optical_flow: &'a LatestOpticalFlowMailbox,
}

/// Vertical-sensor metadata and freshness policy for an outbound data channel.
#[derive(Clone, Copy)]
pub struct VerticalDataConfig {
    /// Metadata required to encode `DISTANCE_SENSOR`; `None` disables it.
    pub distance: Option<DistanceSensorConfig>,
    /// MTF-01 sensor ID for `OPTICAL_FLOW`; `None` disables it.
    pub optical_flow_sensor_id: Option<u8>,
    /// Maximum age, in service ticks, for cached vertical observations.
    pub maximum_age_ticks: u8,
}

/// Data-channel construction failure.
#[derive(Clone, Copy, Debug, Eq, PartialEq, KaptanError)]
pub enum DataChannelError {
    /// A configured channel must be due at least once per second.
    #[kaptan_error(id = 4001)]
    ZeroRate,
    /// A configured channel exceeds the service frequency.
    #[kaptan_error(id = 4002)]
    RateExceedsServiceFrequency,
    /// Two configured channels cannot share the one-frame-per-tick budget.
    #[kaptan_error(id = 4003)]
    ScheduleFull,
}

/// Complete outbound data flow for one MAVLink endpoint.
pub struct DataChannel<'a, Sink, const SERVICE_HZ: usize> {
    sink: Sink,
    encoder: MavlinkEncoder,
    schedule: Schedule<SERVICE_HZ>,
    sources: DataSources<'a>,
    vertical: VerticalDataConfig,
    observed_imu: u32,
    observed_range: u32,
    last_range: Option<RangeSnapshot>,
    range_age_ticks: u8,
    observed_optical_flow: u32,
    last_optical_flow: Option<OpticalFlowSnapshot>,
    optical_flow_age_ticks: u8,
}

impl<'a, Sink, const SERVICE_HZ: usize> DataChannel<'a, Sink, SERVICE_HZ>
where
    Sink: FrameSink,
{
    /// Creates an outbound data flow from already-validated runtime inputs.
    pub fn new(
        sink: Sink,
        system_id: u8,
        component_id: u8,
        rates: DataChannelRates,
        vertical: VerticalDataConfig,
        sources: DataSources<'a>,
    ) -> Result<Self, DataChannelError> {
        let schedule = Schedule::from_rates(rates)?;
        Ok(Self {
            sink,
            encoder: MavlinkEncoder::new(system_id, component_id),
            schedule,
            sources,
            vertical,
            observed_imu: 0,
            observed_range: 0,
            last_range: None,
            range_age_ticks: u8::MAX,
            observed_optical_flow: 0,
            last_optical_flow: None,
            optical_flow_age_ticks: u8::MAX,
        })
    }

    /// Runs all outbound work due for one platform service tick.
    ///
    /// No schedule token escapes this method. It refreshes cached vertical
    /// observations, selects due data, encodes the MAVLink message, and admits
    /// it with the correct bounded-queue class.
    pub fn on_service_tick(&mut self, system_status: MavState, motors_armed: bool) {
        self.refresh_vertical_data();

        match self.schedule.next_due() {
            Some(DueData::RawImu) => {
                if let Some(sample) = self.sources.imu.take_newer(&mut self.observed_imu) {
                    let _ = self.sink.send(
                        FrameClass::BestEffort,
                        self.encoder.raw_imu_snapshot(&sample),
                    );
                }
            }
            Some(DueData::Distance) if self.range_age_ticks <= self.vertical.maximum_age_ticks => {
                if let (Some(range), Some(config)) = (self.last_range, self.vertical.distance) {
                    let _ = self.sink.send(
                        FrameClass::BestEffort,
                        self.encoder.distance_sensor_snapshot(&range, config),
                    );
                }
            }
            Some(DueData::OpticalFlow)
                if self.optical_flow_age_ticks <= self.vertical.maximum_age_ticks =>
            {
                if let (Some(flow), Some(sensor_id)) =
                    (self.last_optical_flow, self.vertical.optical_flow_sensor_id)
                {
                    let _ = self.sink.send(
                        FrameClass::BestEffort,
                        self.encoder.optical_flow_snapshot(&flow, sensor_id),
                    );
                }
            }
            Some(DueData::Heartbeat) => {
                let _ = self.sink.send(
                    FrameClass::Vital,
                    self.encoder.heartbeat(system_status, motors_armed),
                );
            }
            Some(DueData::Distance | DueData::OpticalFlow) | None => {}
        }
    }

    /// Completes the active DMA lifecycle for this channel's transport.
    pub fn on_dma_complete(&mut self, completion: TxCompletion) {
        self.sink.on_dma_complete(completion)
    }

    /// Encodes an adapter-owned generated MAVLink data record through the
    /// channel's sole encoder and bounded transport sink.
    pub fn send_custom<M: MessageData>(&mut self, class: FrameClass, message: &M) -> EnqueueResult {
        self.sink.send(class, self.encoder.custom(message))
    }

    fn refresh_vertical_data(&mut self) {
        if let Some(range) = self.sources.range.take_newer(&mut self.observed_range) {
            self.last_range = Some(range);
            self.range_age_ticks = 0;
        } else {
            self.range_age_ticks = self.range_age_ticks.saturating_add(1);
        }
        if let Some(flow) = self
            .sources
            .optical_flow
            .take_newer(&mut self.observed_optical_flow)
        {
            self.last_optical_flow = Some(flow);
            self.optical_flow_age_ticks = 0;
        } else {
            self.optical_flow_age_ticks = self.optical_flow_age_ticks.saturating_add(1);
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum DueData {
    Heartbeat,
    RawImu,
    Distance,
    OpticalFlow,
}

struct Schedule<const SERVICE_HZ: usize> {
    phase: usize,
    due: [Option<DueData>; SERVICE_HZ],
}

impl<const SERVICE_HZ: usize> Schedule<SERVICE_HZ> {
    fn from_rates(rates: DataChannelRates) -> Result<Self, DataChannelError> {
        assert!(
            SERVICE_HZ != 0,
            "data-channel service frequency must be nonzero"
        );
        if rates.service_hz != SERVICE_HZ {
            return Err(DataChannelError::RateExceedsServiceFrequency);
        }

        let schedule = Self {
            phase: 0,
            due: [None; SERVICE_HZ],
        }
        .with(DueData::RawImu, rates.raw_imu_hz)?
        .with(DueData::Heartbeat, rates.heartbeat_hz)?;
        let schedule = match rates.range_hz {
            Some(hertz) => schedule.with(DueData::Distance, hertz)?,
            None => schedule,
        };
        match rates.optical_flow_hz {
            Some(hertz) => schedule.with(DueData::OpticalFlow, hertz),
            None => Ok(schedule),
        }
    }

    fn next_due(&mut self) -> Option<DueData> {
        let due = self.due[self.phase];
        self.phase = if self.phase + 1 == SERVICE_HZ {
            0
        } else {
            self.phase + 1
        };
        due
    }

    fn with(mut self, data: DueData, hertz: usize) -> Result<Self, DataChannelError> {
        if hertz == 0 {
            return Err(DataChannelError::ZeroRate);
        }
        if hertz > SERVICE_HZ {
            return Err(DataChannelError::RateExceedsServiceFrequency);
        }

        let mut count = 0;
        while count < hertz {
            let mut index = (count * SERVICE_HZ) / hertz;
            let mut probes = 0;
            while self.due[index].is_some() {
                index = if index + 1 == SERVICE_HZ {
                    0
                } else {
                    index + 1
                };
                probes += 1;
                if probes == SERVICE_HZ {
                    return Err(DataChannelError::ScheduleFull);
                }
            }
            self.due[index] = Some(data);
            count += 1;
        }
        Ok(self)
    }
}

#[cfg(test)]
mod tests {
    use super::{
        DataChannel, DataChannelRates, DataSources, FrameSink, LatestImuMailbox,
        LatestOpticalFlowMailbox, LatestRangeMailbox, VerticalDataConfig,
    };
    use kaptan_dma::TxCompletion;
    use kaptan_mavlink::{
        MavlinkFrame,
        message::MavState,
        transmitter::{EnqueueResult, FrameClass},
    };
    use kaptan_sensor::{
        imu::ImuSnapshot, optical_flow::OpticalFlowSnapshot, range::RangeSnapshot,
    };
    use kaptan_time::Timestamp;

    struct Sink {
        best_effort: u8,
        critical: u8,
        completions: u8,
    }

    impl FrameSink for Sink {
        fn send(&mut self, class: FrameClass, _: MavlinkFrame) -> EnqueueResult {
            match class {
                FrameClass::Vital => self.critical += 1,
                FrameClass::BestEffort => self.best_effort += 1,
                FrameClass::Critical => self.critical += 1,
            }
            EnqueueResult::Enqueued
        }

        fn on_dma_complete(&mut self, _: TxCompletion) {
            self.completions += 1;
        }
    }

    const RATES: DataChannelRates = DataChannelRates {
        service_hz: 100,
        raw_imu_hz: 20,
        range_hz: Some(20),
        optical_flow_hz: Some(20),
        heartbeat_hz: 2,
    };

    #[test]
    fn owns_rate_selection_encoding_admission_and_completion() {
        let imu = LatestImuMailbox::new();
        let range = LatestRangeMailbox::new();
        let flow = LatestOpticalFlowMailbox::new();
        let mut channel: DataChannel<'_, Sink, 100> = DataChannel::new(
            Sink {
                best_effort: 0,
                critical: 0,
                completions: 0,
            },
            1,
            1,
            RATES,
            VerticalDataConfig {
                distance: Some(kaptan_mavlink::encoder::DistanceSensorConfig {
                    sensor_id: 0,
                    orientation:
                        kaptan_mavlink::message::MavSensorOrientation::MAV_SENSOR_ROTATION_NONE,
                    minimum_distance_centimetres: 1,
                    maximum_distance_centimetres: 1_000,
                }),
                optical_flow_sensor_id: Some(0),
                maximum_age_ticks: 10,
            },
            DataSources {
                imu: &imu,
                range: &range,
                optical_flow: &flow,
            },
        )
        .unwrap();

        for tick in 0..100 {
            imu.publish(ImuSnapshot::new(
                Timestamp::from_microseconds(tick),
                [0; 3],
                [0; 3],
                0,
            ));
            range.publish(RangeSnapshot::new(Timestamp::from_microseconds(tick), 100));
            flow.publish(OpticalFlowSnapshot::new(
                Timestamp::from_microseconds(tick),
                [0; 2],
                100,
                100,
            ));
            channel.on_service_tick(MavState::MAV_STATE_STANDBY, false);
        }
        channel.on_dma_complete(TxCompletion::Complete);
        let sink = channel.sink;
        assert_eq!(sink.best_effort, 60);
        assert_eq!(sink.critical, 2);
        assert_eq!(sink.completions, 1);
    }
}