//! This module provides [HeapMessage], an implementation of all the message traits backed by heap
//! memory.
//!
//! It serves both to easily get an owned copy of a message (using [`HeapMessage::new()`] and
//! [`MinimalWritableMessage::set_from_message()`]), especially if that needs frequent manipulation
//! of options, and as an example of how to the message traits can be implemented.
//!
//! This module is available only if the `alloc` feature is enabled on this crate.
//!
//! Cross references:
//!
//! * The [coap-lite] crate's Packet struct is very similar to this, to the point where their
//!   [MutableWritableMessage] implementations are identical. They differ by whether or not they
//!   carry CoAP-over-UDP header information.
//!
//! * This implementation used to be part of the [coap_message] crate before its 0.3 version.
//!
//! [coap-lite]: https://crates.io/crates/coap-lite

extern crate alloc;
use alloc::{vec, vec::Vec};
use core::convert::Infallible;

use coap_message::{
    MinimalWritableMessage, MutableWritableMessage, ReadableMessage, SeekWritableMessage,
    WithSortedOptions,
};

/// A heap CoAP message backed by allocated memory
///
/// It stores its payload in a [`Vec`], and uses a [`BTreeMap<_, Vec<_>>`](alloc::collections::BTreeMap) to
/// store all the individual option values.
///
/// It offers a few methods for direct manipulation of options, even out of sequence, that can not
/// be expected from a general message buffer and are thus not captured in traits.
///
/// ```
/// # use coap_message_implementations::heap::HeapMessage;
/// use coap_message::{MinimalWritableMessage, ReadableMessage};
/// let mut m = HeapMessage::new();
/// m.set_code(1);
/// m.add_option(11, b".well-known");
/// m.add_option(11, b"core");
///
/// let mut m2 = HeapMessage::new();
/// m2.set_from_message(&m);
/// assert!(m.code() == 1);
/// ```
#[derive(Debug)]
pub struct HeapMessage {
    code: u8,
    options: alloc::collections::BTreeMap<u16, Vec<Vec<u8>>>,
    payload: Vec<u8>,
}

impl HeapMessage {
    pub fn new() -> Self {
        Self {
            code: 0,
            options: alloc::collections::BTreeMap::new(),
            payload: vec![],
        }
    }

    /// Replace the occurrence'th value of the optnum option in the message
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn change_option(&mut self, optnum: u16, occurrence: usize, value: impl Into<Vec<u8>>) {
        self.options.get_mut(&optnum).unwrap()[occurrence] = value.into();
    }

    /// Remove the occurrence'th option of option number optnum
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn remove_option(&mut self, optnum: u16, occurrence: usize) {
        let vec = self.options.get_mut(&optnum).unwrap();
        vec.remove(occurrence);
        if vec.len() == 0 {
            self.options.remove(&optnum);
        }
    }
}

impl MinimalWritableMessage for HeapMessage {
    type AddOptionError = Infallible;
    type SetPayloadError = Infallible;
    type UnionError = Infallible;

    type Code = u8;
    type OptionNumber = u16;

    fn set_code(&mut self, code: u8) {
        self.code = code;
    }

    fn add_option(&mut self, optnum: u16, data: &[u8]) -> Result<(), Infallible> {
        self.options
            .entry(optnum)
            .or_insert(vec![])
            .push(data.to_vec());
        Ok(())
    }

    fn set_payload(&mut self, payload: &[u8]) -> Result<(), Infallible> {
        self.payload = payload.to_vec();
        Ok(())
    }
}

impl MutableWritableMessage for HeapMessage {
    fn available_space(&self) -> usize {
        core::usize::MAX
    }

    fn payload_mut_with_len(&mut self, len: usize) -> Result<&mut [u8], Infallible> {
        self.payload.resize(len, 0);
        Ok(&mut self.payload)
    }

    fn truncate(&mut self, len: usize) -> Result<(), Infallible> {
        self.payload.truncate(len);
        Ok(())
    }

    fn mutate_options<F>(&mut self, mut callback: F)
    where
        F: FnMut(Self::OptionNumber, &mut [u8]),
    {
        for (&number, ref mut values) in self.options.iter_mut() {
            for v in values.iter_mut() {
                callback(number.into(), v);
            }
        }
    }
}

impl SeekWritableMessage for HeapMessage {}

pub struct MessageOption<'a> {
    number: u16,
    value: &'a [u8],
}

impl<'a> coap_message::MessageOption for MessageOption<'a> {
    fn number(&self) -> u16 {
        self.number
    }
    fn value(&self) -> &[u8] {
        self.value
    }
}

pub struct ReadCursor<'a> {
    iter: alloc::collections::btree_map::Iter<'a, u16, Vec<Vec<u8>>>,
    popped: (u16, &'a [Vec<u8>]),
}

impl<'a> Iterator for ReadCursor<'a> {
    type Item = MessageOption<'a>;
    fn next(&mut self) -> Option<MessageOption<'a>> {
        if self.popped.1.len() == 0 {
            self.popped = self.iter.next().map(|(k, v)| (*k, v.as_slice()))?;

            // Avoided by construction
            debug_assert!(
                self.popped.1.len() > 0,
                "HeapMessage had present but empty option"
            );
        }

        let (first, rest) = self.popped.1.split_at(1);
        self.popped.1 = rest;
        Some(MessageOption {
            number: self.popped.0,
            value: first[0].as_slice(),
        })
    }
}

impl WithSortedOptions for HeapMessage {
    // By virtue of the order of options used in BTreeMap
}

impl ReadableMessage for HeapMessage {
    type Code = u8;
    type MessageOption<'a> = MessageOption<'a>;
    type OptionsIter<'a> = ReadCursor<'a>;

    fn code(&self) -> u8 {
        self.code
    }
    fn payload(&self) -> &[u8] {
        &self.payload
    }
    fn options<'m>(&'m self) -> ReadCursor<'m> {
        ReadCursor {
            iter: self.options.iter(),
            popped: (0, &[]),
        }
    }
}