quicproquo/crates/quicprochat-p2p/src/transport_lora.rs

//! LoRa-style constrained transport with mock RF medium, fragmentation, and EU868 duty-cycle budgeting.
//!
//! Real hardware typically uses a UART-attached module; this crate ships a [`LoRaMockMedium`] that
//! delivers frames between registered node addresses for tests and the integration example.
//!
//! # Wire format (mock / modem-passthrough oriented)
//!
//! - **Whole datagram** (`0x01`): `LR` magic, type, 4-byte source, 4-byte destination, `u16` BE length, payload.
//! - **Fragment** (`0x02`): same header prefix + `frag_id` (u32 BE), `idx`, `total`, `u16` BE chunk length, chunk.

use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{bail, Result};
use tokio::sync::Mutex;
use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};

use crate::transport::{MeshTransport, TransportAddr, TransportInfo, TransportPacket};

const FRAME_MAGIC: [u8; 2] = *b"LR";
const TYPE_WHOLE: u8 = 0x01;
const TYPE_FRAG: u8 = 0x02;

const WHOLE_HEADER: usize = 2 + 1 + 4 + 4 + 2;
const FRAG_HEADER: usize = 2 + 1 + 4 + 4 + 4 + 1 + 1 + 2;

/// LoRa radio and serial link parameters (modem AT layer is out of scope; UART path is optional extension).
#[derive(Clone, Debug)]
pub struct LoRaConfig {
    /// Serial device path when using hardware (informational / future UART backend).
    pub port: String,
    pub baud_rate: u32,
    pub frequency: u64,
    pub spreading_factor: u8,
    pub bandwidth: u32,
    /// LoRa coding rate denominator n in 4/n (5..=8 → 4/5 .. 4/8).
    pub coding_rate: u8,
    pub tx_power: i8,
    /// Max frame size including headers (modem MTU). If `None`, derived from spreading factor.
    pub max_frame_len: Option<usize>,
}

impl Default for LoRaConfig {
    fn default() -> Self {
        Self {
            port: String::new(),
            baud_rate: 115_200,
            frequency: 868_100_000,
            spreading_factor: 7,
            bandwidth: 125_000,
            coding_rate: 5,
            tx_power: 14,
            max_frame_len: None,
        }
    }
}

impl LoRaConfig {
    /// Typical max MAC payload for EU868 / 125 kHz (order-of-magnitude; modem-specific).
    pub fn default_max_frame_len(&self) -> usize {
        if let Some(m) = self.max_frame_len {
            return m.clamp(WHOLE_HEADER + 1, 256);
        }
        let mtu = match self.spreading_factor {
            7 | 8 => 222,
            9 => 115,
            10 | 11 | 12 => 51,
            _ => 128,
        };
        mtu.clamp(WHOLE_HEADER + 1, 256)
    }

    fn cr_index(&self) -> u64 {
        match self.coding_rate.clamp(5, 8) {
            5 => 1,
            6 => 2,
            7 => 3,
            8 => 4,
            _ => 1,
        }
    }
}

/// Approximate LoRa time-on-air in milliseconds for a given PHY payload length (including our framing).
pub fn lora_airtime_ms(payload_len: usize, cfg: &LoRaConfig) -> u64 {
    let sf = cfg.spreading_factor.clamp(7, 12) as u64;
    let bw = cfg.bandwidth.max(7_800) as u64;
    let t_sym_us = ((1u64 << sf) * 1_000_000u64) / bw;
    let preamble_syms = 12u64 + 4;
    let preamble_us = preamble_syms * t_sym_us;

    let de = 0i64;
    let pl = payload_len as i64;
    let sf_i = sf as i64;
    let numerator = 8 * pl - 4 * sf_i + 28 + 16 - 20;
    let denom = 4 * (sf_i - 2 * de);
    let payload_symb = if denom > 0 && numerator > 0 {
        let ceiled = (numerator + denom - 1) / denom;
        let cr = cfg.cr_index() as i64;
        8 + ceiled * (cr + 4)
    } else {
        8i64
    };
    let payload_us = (payload_symb as u64).saturating_mul(t_sym_us);
    (preamble_us + payload_us) / 1000
}

/// Rough PHY bitrate estimate (bits/s) for routing metrics — not precise at low SNR.
pub fn lora_nominal_bitrate_bps(cfg: &LoRaConfig) -> u64 {
    let sf = cfg.spreading_factor.clamp(7, 12) as u32;
    let bw = cfg.bandwidth.max(7_800);
    // bits per symbol ≈ SF; symbol rate ≈ BW / 2^SF
    let sym_rate = (bw as u64) / (1u64 << sf);
    sym_rate.saturating_mul(sf as u64)
}

/// EU868-style 1% duty cycle: at most 36_000 ms airtime per rolling hour.
#[derive(Debug)]
pub struct DutyCycleTracker {
    max_ms_per_hour: u64,
    window: Mutex<VecDeque<(Instant, u64)>>,
}

impl DutyCycleTracker {
    pub fn new(max_ms_airtime_per_hour: u64) -> Self {
        Self {
            max_ms_per_hour: max_ms_airtime_per_hour,
            window: Mutex::new(VecDeque::new()),
        }
    }

    /// 1% of one hour = 36 seconds of transmission time.
    pub fn eu868_one_percent() -> Self {
        Self::new(36_000)
    }

    fn prune_old( deque: &mut VecDeque<(Instant, u64)>) {
        let cutoff = Instant::now() - Duration::from_secs(3600);
        while let Some(&(t, _)) = deque.front() {
            if t < cutoff {
                deque.pop_front();
            } else {
                break;
            }
        }
    }

    fn sum_ms(deque: &VecDeque<(Instant, u64)>) -> u64 {
        deque.iter().map(|(_, m)| m).sum()
    }

    /// Wait until `airtime_ms` fits in the budget, then record it.
    pub async fn acquire(&self, airtime_ms: u64) {
        loop {
            let sleep_for = {
                let mut deque = self.window.lock().await;
                Self::prune_old(&mut deque);
                let used = Self::sum_ms(&deque);
                if used + airtime_ms <= self.max_ms_per_hour {
                    deque.push_back((Instant::now(), airtime_ms));
                    return;
                }
                if let Some(&(oldest, _)) = deque.front() {
                    let elapsed = oldest.elapsed();
                    let until_refresh = Duration::from_secs(3600).saturating_sub(elapsed);
                    until_refresh.max(Duration::from_millis(1))
                } else {
                    Duration::from_millis(1)
                }
            };
            tokio::time::sleep(sleep_for).await;
        }
    }

    /// Total recorded airtime in the current window (for tests / diagnostics).
    pub async fn used_ms_in_window(&self) -> u64 {
        let mut deque = self.window.lock().await;
        Self::prune_old(&mut deque);
        Self::sum_ms(&deque)
    }
}

/// In-process mock RF cloud: addressed delivery between registered 4-byte LoRa addresses.
#[derive(Debug)]
pub struct LoRaMockMedium {
    nodes: Mutex<HashMap<[u8; 4], UnboundedSender<Vec<u8>>>>,
}

impl LoRaMockMedium {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            nodes: Mutex::new(HashMap::new()),
        })
    }

    /// Register a node; returns a transport bound to `my_addr`.
    pub async fn connect(
        self: &Arc<Self>,
        my_addr: [u8; 4],
        config: LoRaConfig,
        duty: Arc<DutyCycleTracker>,
    ) -> Result<LoRaTransport> {
        let (tx, rx) = unbounded_channel();
        let mut map = self.nodes.lock().await;
        if map.contains_key(&my_addr) {
            bail!("LoRa address already registered on this medium");
        }
        map.insert(my_addr, tx);
        drop(map);

        Ok(LoRaTransport {
            medium: Arc::clone(self),
            my_addr,
            inbox: Mutex::new(rx),
            config,
            duty,
            assembler: Mutex::new(FragmentAssembler::default()),
        })
    }

    async fn deliver(self: &Arc<Self>, dest: [u8; 4], frame: Vec<u8>) -> Result<()> {
        let sender = {
            let map = self.nodes.lock().await;
            map.get(&dest)
                .cloned()
                .ok_or_else(|| anyhow::anyhow!("unknown LoRa destination {dest:02x?}"))?
        };
        sender
            .send(frame)
            .map_err(|_| anyhow::anyhow!("LoRa peer inbox closed"))?;
        Ok(())
    }

    async fn unregister(self: &Arc<Self>, addr: [u8; 4]) {
        let mut map = self.nodes.lock().await;
        map.remove(&addr);
    }
}

/// LoRa [`MeshTransport`] using [`LoRaMockMedium`].
pub struct LoRaTransport {
    medium: Arc<LoRaMockMedium>,
    my_addr: [u8; 4],
    inbox: Mutex<UnboundedReceiver<Vec<u8>>>,
    config: LoRaConfig,
    duty: Arc<DutyCycleTracker>,
    assembler: Mutex<FragmentAssembler>,
}

impl LoRaTransport {
    pub fn local_address(&self) -> [u8; 4] {
        self.my_addr
    }

    pub fn transport_addr(&self) -> TransportAddr {
        TransportAddr::LoRa(self.my_addr)
    }

    fn max_frame_len(&self) -> usize {
        self.config.default_max_frame_len()
    }

    fn whole_payload_cap(&self) -> usize {
        self.max_frame_len().saturating_sub(WHOLE_HEADER)
    }

    fn frag_payload_cap(&self) -> usize {
        self.max_frame_len().saturating_sub(FRAG_HEADER)
    }

    fn build_whole(src: [u8; 4], dst: [u8; 4], payload: &[u8]) -> Result<Vec<u8>> {
        let len = payload.len();
        if len > u16::MAX as usize {
            bail!("LoRa payload too large");
        }
        let mut v = Vec::with_capacity(WHOLE_HEADER + len);
        v.extend_from_slice(&FRAME_MAGIC);
        v.push(TYPE_WHOLE);
        v.extend_from_slice(&src);
        v.extend_from_slice(&dst);
        v.extend_from_slice(&(len as u16).to_be_bytes());
        v.extend_from_slice(payload);
        Ok(v)
    }

    fn build_frag(
        src: [u8; 4],
        dst: [u8; 4],
        frag_id: u32,
        idx: u8,
        total: u8,
        chunk: &[u8],
    ) -> Result<Vec<u8>> {
        let len = chunk.len();
        if len > u16::MAX as usize {
            bail!("fragment chunk too large");
        }
        let mut v = Vec::with_capacity(FRAG_HEADER + len);
        v.extend_from_slice(&FRAME_MAGIC);
        v.push(TYPE_FRAG);
        v.extend_from_slice(&src);
        v.extend_from_slice(&dst);
        v.extend_from_slice(&frag_id.to_be_bytes());
        v.push(idx);
        v.push(total);
        v.extend_from_slice(&(len as u16).to_be_bytes());
        v.extend_from_slice(chunk);
        Ok(v)
    }

    fn parse_frame(buf: &[u8]) -> Result<ParsedFrame> {
        if buf.len() < 2 || buf[0] != FRAME_MAGIC[0] || buf[1] != FRAME_MAGIC[1] {
            bail!("invalid LoRa frame magic");
        }
        if buf.len() < 3 {
            bail!("truncated LoRa frame");
        }
        match buf[2] {
            TYPE_WHOLE => {
                if buf.len() < WHOLE_HEADER {
                    bail!("truncated whole frame");
                }
                let mut src = [0u8; 4];
                src.copy_from_slice(&buf[3..7]);
                let mut dst = [0u8; 4];
                dst.copy_from_slice(&buf[7..11]);
                let plen = u16::from_be_bytes([buf[11], buf[12]]) as usize;
                if buf.len() != WHOLE_HEADER + plen {
                    bail!("whole frame length mismatch");
                }
                Ok(ParsedFrame::Whole {
                    src,
                    dst,
                    payload: buf[WHOLE_HEADER..].to_vec(),
                })
            }
            TYPE_FRAG => {
                if buf.len() < FRAG_HEADER {
                    bail!("truncated fragment frame");
                }
                let mut src = [0u8; 4];
                src.copy_from_slice(&buf[3..7]);
                let mut dst = [0u8; 4];
                dst.copy_from_slice(&buf[7..11]);
                let frag_id = u32::from_be_bytes([buf[11], buf[12], buf[13], buf[14]]);
                let idx = buf[15];
                let total = buf[16];
                let clen = u16::from_be_bytes([buf[17], buf[18]]) as usize;
                if buf.len() != FRAG_HEADER + clen {
                    bail!("fragment length mismatch");
                }
                Ok(ParsedFrame::Frag {
                    src,
                    dst,
                    frag_id,
                    idx,
                    total,
                    chunk: buf[FRAG_HEADER..].to_vec(),
                })
            }
            t => bail!("unknown LoRa frame type {t}"),
        }
    }
}

enum ParsedFrame {
    Whole {
        src: [u8; 4],
        dst: [u8; 4],
        payload: Vec<u8>,
    },
    Frag {
        src: [u8; 4],
        dst: [u8; 4],
        frag_id: u32,
        idx: u8,
        total: u8,
        chunk: Vec<u8>,
    },
}

#[derive(Default)]
struct FragmentAssembler {
    partials: HashMap<(u32, [u8; 4]), PartialFrag>,
}

struct PartialFrag {
    total: u8,
    pieces: HashMap<u8, Vec<u8>>,
    started: Instant,
}

impl FragmentAssembler {
    const TIMEOUT: Duration = Duration::from_secs(120);

    fn push(
        &mut self,
        src: [u8; 4],
        frag_id: u32,
        idx: u8,
        total: u8,
        chunk: Vec<u8>,
    ) -> Result<Option<Vec<u8>>> {
        self.gc();
        let key = (frag_id, src);
        let entry = self
            .partials
            .entry(key)
            .or_insert_with(|| PartialFrag {
                total,
                pieces: HashMap::new(),
                started: Instant::now(),
            });
        if entry.total != total {
            bail!("fragment total mismatch");
        }
        entry.pieces.insert(idx, chunk);
        if entry.pieces.len() == total as usize {
            let mut out = Vec::new();
            for i in 0..total {
                let piece = entry
                    .pieces
                    .get(&i)
                    .ok_or_else(|| anyhow::anyhow!("missing fragment index {i}"))?;
                out.extend_from_slice(piece);
            }
            self.partials.remove(&key);
            return Ok(Some(out));
        }
        Ok(None)
    }

    fn gc(&mut self) {
        let now = Instant::now();
        self.partials
            .retain(|_, p| now.duration_since(p.started) < Self::TIMEOUT);
    }
}

#[async_trait::async_trait]
impl MeshTransport for LoRaTransport {
    fn info(&self) -> TransportInfo {
        TransportInfo {
            name: "lora".to_string(),
            mtu: self.whole_payload_cap(),
            bitrate: lora_nominal_bitrate_bps(&self.config),
            bidirectional: true,
        }
    }

    async fn send(&self, dest: &TransportAddr, data: &[u8]) -> Result<()> {
        let dst = match dest {
            TransportAddr::LoRa(a) => *a,
            other => bail!("LoRaTransport cannot send to {other}"),
        };

        let max_frame = self.max_frame_len();
        let cap_whole = self.whole_payload_cap();
        let cap_frag = self.frag_payload_cap().max(1);

        let frames: Vec<Vec<u8>> = if data.len() <= cap_whole {
            vec![Self::build_whole(self.my_addr, dst, data)?]
        } else {
            let frag_id = random_frag_id();
            let chunk_sz = cap_frag;
            let total = data.chunks(chunk_sz).count();
            if total > u8::MAX as usize {
                bail!("payload needs more than 255 fragments");
            }
            let total_u8 = total as u8;
            let mut out = Vec::with_capacity(total);
            for (idx, chunk) in data.chunks(chunk_sz).enumerate() {
                out.push(Self::build_frag(
                    self.my_addr,
                    dst,
                    frag_id,
                    idx as u8,
                    total_u8,
                    chunk,
                )?);
            }
            out
        };

        for frame in frames {
            let air = lora_airtime_ms(frame.len(), &self.config);
            self.duty.acquire(air).await;
            if frame.len() > max_frame {
                bail!("LoRa frame exceeds configured MTU");
            }
            self.medium.deliver(dst, frame).await?;
        }

        Ok(())
    }

    async fn recv(&self) -> Result<TransportPacket> {
        loop {
            let raw = {
                let mut inbox = self.inbox.lock().await;
                inbox
                    .recv()
                    .await
                    .ok_or_else(|| anyhow::anyhow!("LoRa inbox closed"))?
            };

            match Self::parse_frame(&raw)? {
                ParsedFrame::Whole { src, dst, payload } => {
                    if dst != self.my_addr {
                        continue;
                    }
                    return Ok(TransportPacket {
                        from: TransportAddr::LoRa(src),
                        data: payload,
                    });
                }
                ParsedFrame::Frag {
                    src,
                    dst,
                    frag_id,
                    idx,
                    total,
                    chunk,
                } => {
                    if dst != self.my_addr {
                        continue;
                    }
                    let mut asm = self.assembler.lock().await;
                    if let Some(complete) = asm.push(src, frag_id, idx, total, chunk)? {
                        return Ok(TransportPacket {
                            from: TransportAddr::LoRa(src),
                            data: complete,
                        });
                    }
                }
            }
        }
    }

    async fn close(&self) -> Result<()> {
        self.medium.unregister(self.my_addr).await;
        Ok(())
    }
}

fn random_frag_id() -> u32 {
    use rand::Rng;
    rand::thread_rng().gen::<u32>()
}

/// Split `data` into chunks suitable for a transport with `max_payload` bytes per frame (application layer).
pub fn split_for_mtu(data: &[u8], max_payload: usize) -> Vec<&[u8]> {
    if max_payload == 0 {
        return vec![data];
    }
    data.chunks(max_payload).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn airtime_increases_with_sf() {
        let mut low = LoRaConfig::default();
        low.spreading_factor = 7;
        let mut high = LoRaConfig::default();
        high.spreading_factor = 12;
        let n = 64;
        assert!(lora_airtime_ms(n, &high) >= lora_airtime_ms(n, &low));
    }

    #[tokio::test]
    async fn mock_roundtrip() {
        let medium = LoRaMockMedium::new();
        let duty = Arc::new(DutyCycleTracker::new(3600 * 1000));
        let a = medium
            .connect([1, 0, 0, 0], LoRaConfig::default(), Arc::clone(&duty))
            .await
            .expect("connect a");
        let b = medium
            .connect([2, 0, 0, 0], LoRaConfig::default(), Arc::clone(&duty))
            .await
            .expect("connect b");

        let dest = TransportAddr::LoRa([2, 0, 0, 0]);
        let payload = b"mesh-over-lora";

        let recv_h = tokio::spawn(async move {
            let pkt = b.recv().await.expect("recv");
            assert_eq!(pkt.data, payload.to_vec());
            match pkt.from {
                TransportAddr::LoRa(addr) => assert_eq!(addr, [1, 0, 0, 0]),
                _ => panic!("expected LoRa from-address"),
            }
            b.close().await.expect("close b");
        });

        tokio::time::sleep(Duration::from_millis(20)).await;
        a.send(&dest, payload).await.expect("send");

        recv_h.await.expect("join");
        a.close().await.expect("close a");
    }

    #[tokio::test]
    async fn fragmentation_roundtrip() {
        let medium = LoRaMockMedium::new();
        let duty = Arc::new(DutyCycleTracker::new(3600 * 1000));
        let mut cfg = LoRaConfig::default();
        cfg.max_frame_len = Some(48);
        let a = medium
            .connect([0x10, 0, 0, 0], cfg.clone(), Arc::clone(&duty))
            .await
            .expect("a");
        let b = medium
            .connect([0x20, 0, 0, 0], cfg, Arc::clone(&duty))
            .await
            .expect("b");

        let dest = TransportAddr::LoRa([0x20, 0, 0, 0]);
        let payload: Vec<u8> = (0u8..200).collect();
        let expected = payload.clone();

        let recv_h = tokio::spawn(async move {
            let pkt = b.recv().await.expect("recv");
            assert_eq!(pkt.data, expected);
            b.close().await.ok();
        });

        tokio::time::sleep(Duration::from_millis(20)).await;
        a.send(&dest, &payload).await.expect("send frag");

        recv_h.await.expect("join");
        a.close().await.ok();
    }

    #[tokio::test]
    async fn duty_cycle_records_airtime() {
        let duty = Arc::new(DutyCycleTracker::new(100_000));
        duty.acquire(55).await;
        let used = duty.used_ms_in_window().await;
        assert!(used >= 55, "expected recorded airtime, got {used}");
    }

    #[test]
    fn split_for_mtu_chunks() {
        let data = [1u8, 2, 3, 4, 5];
        let parts = split_for_mtu(&data, 2);
        assert_eq!(parts.len(), 3);
        assert_eq!(parts[0], &[1, 2][..]);
        assert_eq!(parts[1], &[3, 4][..]);
        assert_eq!(parts[2], &[5][..]);
    }
}