//! Fragmentation and reassembly of `Observation`s over a byte-oriented link.

use crate::consts;
use crate::{Frame, NodeId, Observation, ObservationId, Tick, WireError};
use heapless::Vec;

/// Header allowance per fragment (bytes), pinned.
pub const FRAG_OVERHEAD: usize = 64;

/// The most fragments one observation may be split into: the reassembler
/// tracks received indices in a `u64` bitset.
pub(crate) const MAX_FRAGMENTS: usize = 64;

/// Splits an encoded observation of `len` bytes for a link of `mtu` bytes,
/// returning `(chunk, count)`: the data bytes per fragment and the number of
/// fragments. `None` when `mtu` leaves no room for data or `count` would
/// exceed [`MAX_FRAGMENTS`].
pub(crate) fn plan(len: usize, mtu: usize) -> Option<(usize, usize)> {
    let chunk = mtu.saturating_sub(FRAG_OVERHEAD);
    if chunk == 0 {
        return None;
    }
    let count = len.div_ceil(chunk);
    if count == 0 || count > MAX_FRAGMENTS {
        return None;
    }
    Some((chunk, count))
}

/// Encodes `obs` as one Obs frame if it fits in `mtu`, else as N Frag frames.
/// Returns the number of frames; writes each via `emit(frame_bytes)`.
pub fn fragment(
    from: NodeId,
    obs: &Observation,
    mtu: usize,
    emit: &mut dyn FnMut(&[u8]),
) -> Result<usize, WireError> {
    // 1. Try to send the whole Observation as a single Obs frame.
    let mut buf = [0u8; consts::MAX_FRAME];
    let frame = Frame::Obs {
        from,
        obs: obs.clone(),
    };
    if let Ok(n) = frame.encode(&mut buf) {
        if n <= mtu {
            emit(&buf[..n]);
            return Ok(1);
        }
    }

    // 2. Encode the Observation alone and split it into Frag frames.
    let mut obs_buf = [0u8; consts::MAX_FRAME * 2];
    let obs_bytes = postcard::to_slice(obs, &mut obs_buf).map_err(|_| WireError::TooLarge)?;

    let (chunk, count) = plan(obs_bytes.len(), mtu).ok_or(WireError::TooLarge)?;

    for i in 0..count {
        let start = i * chunk;
        let end = core::cmp::min(start + chunk, obs_bytes.len());
        let data = &obs_bytes[start..end];
        let frag_data = heapless::Vec::from_slice(data).map_err(|_| WireError::TooLarge)?;
        let frame = Frame::Frag {
            from,
            id: obs.id(),
            index: i as u8,
            count: count as u8,
            data: frag_data,
        };
        let mut fb = [0u8; consts::MAX_FRAME];
        let n = frame.encode(&mut fb).map_err(|_| WireError::TooLarge)?;
        emit(&fb[..n]);
    }
    Ok(count)
}

/// A single-slot reassembler for inbound Frag streams.
pub struct Reassembler {
    id: Option<ObservationId>,
    count: u8,
    received: u64,
    started_at: Tick,
    buf: Vec<u8, { consts::MAX_FRAME * 2 }>,
    offs: [u16; MAX_FRAGMENTS],
    lens: [u16; MAX_FRAGMENTS],
}

impl Reassembler {
    pub fn new() -> Self {
        Reassembler {
            id: None,
            count: 0,
            received: 0,
            started_at: 0,
            buf: heapless::Vec::new(),
            offs: [u16::MAX; MAX_FRAGMENTS],
            lens: [0u16; MAX_FRAGMENTS],
        }
    }

    pub fn feed(
        &mut self,
        id: ObservationId,
        index: u8,
        count: u8,
        data: &[u8],
        now: Tick,
    ) -> Option<Observation> {
        // 1. Drop malformed indices: count must be in 1..=64 and index < count.
        //    `index >= 64` is guarded too so the `1 << index` bitmap never
        //    shifts out of bounds on adversarial wire input.
        if count == 0
            || usize::from(count) > MAX_FRAGMENTS
            || usize::from(index) >= MAX_FRAGMENTS
            || index >= count
        {
            return None;
        }

        // 2. Reset the slot if it is empty, belongs to a different id/count, or
        //    has exceeded the reassembly timeout.
        let need_reset = self.id.is_none()
            || self.id != Some(id)
            || self.count != count
            || now.saturating_sub(self.started_at) > consts::REASSEMBLY_TIMEOUT_MS;
        if need_reset {
            self.id = Some(id);
            self.count = count;
            self.received = 0;
            self.started_at = now;
            self.buf.clear();
            self.offs = [u16::MAX; MAX_FRAGMENTS];
            self.lens = [0u16; MAX_FRAGMENTS];
        }

        // 3. Append the fragment data to the flat buffer in arrival order and
        //    record where this index's bytes live.
        let start = self.buf.len();
        if self.buf.len().saturating_add(data.len()) > self.buf.capacity() {
            // No valid observation encodes to more than MAX_FRAME*2 bytes, so a
            // fragment that cannot fit marks the stream as malformed: reset the
            // slot rather than leaving it stuck holding an unreassemblable set.
            self.reset_slot();
            return None;
        }
        self.buf
            .extend_from_slice(data)
            .expect("buf capacity check above");
        self.offs[index as usize] = start as u16;
        self.lens[index as usize] = data.len() as u16;
        self.received |= 1u64 << index;

        // 4. Not all fragments have arrived yet.
        let complete: u64 = if count >= 64 {
            u64::MAX
        } else {
            (1u64 << count) - 1
        };
        if self.received != complete {
            return None;
        }

        // 5. Reassemble in index order and decode.
        let mut recon = [0u8; consts::MAX_FRAME * 2];
        let mut total = 0usize;
        for i in 0..count {
            let o = self.offs[i as usize] as usize;
            let l = self.lens[i as usize] as usize;
            let frag = &self.buf[o..o + l];
            recon[total..total + l].copy_from_slice(frag);
            total += l;
        }

        match postcard::from_bytes::<Observation>(&recon[..total]) {
            Ok(obs) if obs.id() == id => {
                self.reset_slot();
                Some(obs)
            }
            _ => {
                self.reset_slot();
                None
            }
        }
    }
}

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

impl Reassembler {
    /// Clears a completed/aborted slot so it is ready for a fresh stream.
    fn reset_slot(&mut self) {
        self.id = None;
        self.count = 0;
        self.received = 0;
        self.started_at = 0;
        self.buf.clear();
        self.offs = [u16::MAX; MAX_FRAGMENTS];
        self.lens = [0u16; MAX_FRAGMENTS];
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::wire::Frame;
    use crate::{Payload, Topic};
    use std::vec::Vec;

    fn obs_with_payload(observer: [u8; 16], payload_len: usize) -> Observation {
        Observation {
            observer,
            seq: 1,
            age_ms: 0,
            observed_at: None,
            topic: Topic::from_slice(b"door/3").unwrap(),
            payload: Payload::from_slice(&std::vec![0u8; payload_len]).unwrap(),
            to: None,
            lifetime_ms: None,
        }
    }

    #[test]
    fn small_obs_is_one_obs_frame() {
        let obs = obs_with_payload([1; 16], 10);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let n = fragment([1; 16], &obs, 1500, &mut |b| frames.push(b.to_vec()))
            .expect("fragment small obs");
        assert_eq!(n, 1);
        assert_eq!(frames.len(), 1);
        match Frame::decode(&frames[0]).unwrap() {
            Frame::Obs { .. } => {}
            _ => panic!("expected a single Obs frame"),
        }
    }

    #[test]
    fn large_obs_is_fragmented_and_reassembles() {
        let obs = obs_with_payload([2; 16], 1024);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let n = fragment([2; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        assert!(n >= 6, "expected >= 6 frames, got {n}");
        for f in &frames {
            match Frame::decode(f).unwrap() {
                Frame::Frag { count, .. } => {
                    assert_eq!(count, n as u8, "all frags must share the same count");
                }
                other => panic!("expected only Frag frames, got {other:?}"),
            }
        }
        let mut r = Reassembler::new();
        let mut last: Option<Observation> = None;
        for f in frames.iter() {
            if let Frame::Frag {
                id,
                index,
                count,
                data,
                ..
            } = Frame::decode(f).unwrap()
            {
                last = r.feed(id, index, count, &data, 0);
            }
        }
        let reassembled = last.expect("reassembled obs");
        assert_eq!(
            reassembled, obs,
            "reassembled observation must equal original"
        );
    }

    #[test]
    fn out_of_order_reassembles() {
        let obs = obs_with_payload([3; 16], 1024);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let _n = fragment([3; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        let mut r = Reassembler::new();
        let mut last: Option<Observation> = None;
        for f in frames.iter().rev() {
            if let Frame::Frag {
                id,
                index,
                count,
                data,
                ..
            } = Frame::decode(f).unwrap()
            {
                last = r.feed(id, index, count, &data, 0);
            }
        }
        let reassembled = last.expect("reassembled obs");
        assert_eq!(reassembled, obs);
    }

    #[test]
    fn duplicate_fragment_is_harmless() {
        let obs = obs_with_payload([4; 16], 1024);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let _ = fragment([4; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        let frags: Vec<(ObservationId, u8, u8, Vec<u8>)> = frames
            .iter()
            .map(|f| {
                if let Frame::Frag {
                    id,
                    index,
                    count,
                    data,
                    ..
                } = Frame::decode(f).unwrap()
                {
                    (id, index, count, data.to_vec())
                } else {
                    panic!("expected Frag frames");
                }
            })
            .collect();
        let mut r = Reassembler::new();
        let mut last: Option<Observation> = None;
        let (id0, idx0, cnt0, data0) = &frags[0];
        // feed index 0 twice
        r.feed(*id0, *idx0, *cnt0, data0, 0);
        r.feed(*id0, *idx0, *cnt0, data0, 0);
        for item in &frags[1..] {
            let (id, index, count, data) = item;
            last = r.feed(*id, *index, *count, data, 0);
        }
        let reassembled = last.expect("reassembled obs");
        assert_eq!(reassembled, obs);
    }

    #[test]
    fn wrong_count_is_dropped() {
        let obs = obs_with_payload([5; 16], 1024);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let _ = fragment([5; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        let frags: Vec<(ObservationId, u8, u8, Vec<u8>)> = frames
            .iter()
            .map(|f| {
                if let Frame::Frag {
                    id,
                    index,
                    count,
                    data,
                    ..
                } = Frame::decode(f).unwrap()
                {
                    (id, index, count, data.to_vec())
                } else {
                    panic!("expected Frag frames");
                }
            })
            .collect();
        let mut r = Reassembler::new();
        let mut got_some = false;
        for (id, index, _count, data) in &frags {
            // The changing count keeps resetting the slot, so fragments never
            // accumulate and the stream can never complete (no panic).
            let a = r.feed(*id, *index, 2, data, 0);
            let b = r.feed(*id, *index, 3, data, 0);
            if a.is_some() || b.is_some() {
                got_some = true;
            }
        }
        assert!(!got_some, "a wrong-count stream must never complete");
    }

    #[test]
    fn index_ge_count_is_ignored() {
        let mut r = Reassembler::new();
        let id = [6; 16];
        // index 5 with count 3
        assert!(r.feed(id, 5, 3, &[1, 2, 3], 0).is_none());
        // index == count is also invalid
        assert!(r.feed(id, 3, 3, &[1, 2, 3], 0).is_none());
    }

    #[test]
    fn timeout_discards_partial() {
        let obs = obs_with_payload([7; 16], 200);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let n = fragment([7; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        let frags: Vec<(ObservationId, u8, u8, Vec<u8>)> = frames
            .iter()
            .map(|f| {
                if let Frame::Frag {
                    id,
                    index,
                    count,
                    data,
                    ..
                } = Frame::decode(f).unwrap()
                {
                    (id, index, count, data.to_vec())
                } else {
                    panic!("expected Frag frames");
                }
            })
            .collect();
        assert_eq!(frags.len(), n);
        let mut r = Reassembler::new();
        let (id0, idx0, cnt0, data0) = &frags[0];
        // feed index 0 at now=0 (incomplete)
        r.feed(*id0, *idx0, *cnt0, data0, 0);
        // feed index 1 at now=timeout+1 -> slot restarts with only index 1, returns None
        let (id1, idx1, cnt1, data1) = &frags[1];
        assert!(r
            .feed(*id1, *idx1, *cnt1, data1, consts::REASSEMBLY_TIMEOUT_MS + 1)
            .is_none());
        // re-feed index 0 late: proves index 0 from now=0 was discarded, but the
        // restarted slot now holds index 1, so the stream completes to the original
        let (idc, idxc, cntc, datac) = &frags[0];
        let late = r.feed(*idc, *idxc, *cntc, datac, consts::REASSEMBLY_TIMEOUT_MS + 2);
        assert_eq!(
            late,
            Some(obs.clone()),
            "late re-feed must complete the restarted stream"
        );
        // sanity: a fresh feed of the whole set from now=0 also completes
        let mut r2 = Reassembler::new();
        let mut complete: Option<Observation> = None;
        for (id, index, count, data) in &frags {
            complete = r2.feed(*id, *index, *count, data, 0);
        }
        assert_eq!(complete, Some(obs));
    }

    #[test]
    fn reassembled_bytes_must_match_id() {
        let obs = obs_with_payload([8; 16], 1024);
        let mut frames: Vec<Vec<u8>> = Vec::new();
        let _ = fragment([8; 16], &obs, 200, &mut |b| frames.push(b.to_vec()))
            .expect("fragment large obs");
        let mut frags: Vec<(ObservationId, u8, u8, Vec<u8>)> = frames
            .iter()
            .map(|f| {
                if let Frame::Frag {
                    id,
                    index,
                    count,
                    data,
                    ..
                } = Frame::decode(f).unwrap()
                {
                    (id, index, count, data.to_vec())
                } else {
                    panic!("expected Frag frames");
                }
            })
            .collect();
        // corrupt one byte in the payload region of fragment 0
        frags[0].3[0] ^= 0xFF;
        let mut r = Reassembler::new();
        let mut last: Option<Observation> = None;
        for (id, index, count, data) in &frags {
            last = r.feed(*id, *index, *count, data, 0);
        }
        assert!(last.is_none(), "corrupted reassembly must not yield an obs");
    }

    #[test]
    fn fragment_too_large_returns_error() {
        let obs = obs_with_payload([9; 16], consts::MAX_PAYLOAD);
        // mtu smaller than header -> chunk == 0 -> TooLarge
        let got = fragment([9; 16], &obs, FRAG_OVERHEAD, &mut |_| {}).is_err();
        assert!(got, "chunk == 0 must return an error");
    }

    #[test]
    fn buffer_overflow_resets_slot() {
        let mut r = Reassembler::new();
        // Two max-sized fragments fill the 3000-byte buffer.
        assert!(r.feed([0xA; 16], 0, 64, &[0u8; 1500], 0).is_none());
        assert!(r.feed([0xA; 16], 1, 64, &[0u8; 1500], 0).is_none());
        assert_eq!(r.buf.len(), 3000);
        // A third fragment overflows; the slot must reset, not stay stuck.
        assert!(r.feed([0xA; 16], 2, 64, &[0u8; 1], 0).is_none());
        assert!(r.buf.is_empty(), "overflow must reset the slot buffer");
        assert!(r.id.is_none(), "overflow must clear the slot id");
    }

    #[test]
    fn decode_garbage_is_not_a_frame() {
        // smoke check that wire decode of the module's dependency behaves
        assert!(Frame::decode(&[0xFFu8, 0xFF, 0xFF]).is_err());
    }
}