path: root/crates/windlass/src/messages.rs

use std::collections::BTreeMap;

use crate::dictionary::Dictionary;
use crate::encoding::{FieldType, FieldValue, MessageDecodeError};

/// A parser for a single message type
pub struct MessageParser {
    /// The name of the message
    pub name: String,

    /// The fields of the message, and their types.
    pub fields: Vec<(String, FieldType)>,

    /// How the message should be debug printed
    pub output: Option<OutputFormat>,
}

impl MessageParser {
    /// Create a parser for a message with the given name and printf declaration parts
    pub(crate) fn new<'a>(
        name: &str,
        parts: impl Iterator<Item = &'a str>,
    ) -> Result<Self, MessageSkipperError> {
        let mut fields = vec![];
        for part in parts {
            let (arg, ty) = part
                .split_once('=')
                .ok_or_else(|| MessageSkipperError::InvalidArgumentFormat(part.into()))?;

            let field_type = FieldType::from_msg(ty)?;
            fields.push((arg.to_string(), field_type));
        }
        Ok(Self {
            name: name.to_string(),
            fields,
            output: None,
        })
    }

    /// Create a parser for a message type with the given printf-style specifier
    pub(crate) fn new_output(msg: &str) -> Result<Self, MessageSkipperError> {
        let mut fields = vec![];
        let mut parts = vec![];

        let mut work = msg;
        while let Some(pos) = work.find('%') {
            let (pre, rest) = work.split_at(pos);
            if !pre.is_empty() {
                parts.push(FormatBlock::Static(pre.to_string()));
            }
            if let Some(rest) = rest.strip_prefix("%%") {
                parts.push(FormatBlock::Static("%".to_string()));
                work = rest;
                break;
            }
            let (format, rest) = FieldType::from_format(rest)?;
            parts.push(FormatBlock::Field);
            fields.push((format!("field_{}", fields.len()), format));
            work = rest;
        }
        if !work.is_empty() {
            parts.push(FormatBlock::Static(work.to_string()));
        }

        Ok(Self {
            name: msg.to_string(),
            fields,
            output: Some(OutputFormat { parts }),
        })
    }

    /// Skip over this message at the top of `input`.
    #[allow(dead_code)]
    pub fn skip(&self, input: &mut &[u8]) -> Result<(), MessageDecodeError> {
        for (_, field) in &self.fields {
            field.skip(input)?;
        }
        Ok(())
    }

    /// Skip over this message at the top of `input`, but try to read the `oid` field if it is part of this message.
    pub fn skip_with_oid(&self, input: &mut &[u8]) -> Result<Option<u8>, MessageDecodeError> {
        let mut oid = None;
        for (name, field) in &self.fields {
            if name == "oid" {
                if let FieldValue::U8(read_oid) = field.read(input)? {
                    oid = Some(read_oid);
                }
            } else {
                field.skip(input)?;
            }
        }
        Ok(oid)
    }

    /// Parse a message of this type from the top of `input`.
    pub fn parse(
        &self,
        input: &mut &[u8],
    ) -> Result<BTreeMap<String, FieldValue>, MessageDecodeError> {
        let mut output = BTreeMap::new();
        for (name, field) in &self.fields {
            output.insert(name.to_string(), field.read(input)?);
        }
        Ok(output)
    }
}

/// A message type
pub trait Message: 'static {
    type Pod<'a>: Into<Self::PodOwned> + std::fmt::Debug;
    type PodOwned: Clone + Send + std::fmt::Debug + 'static;

    /// Get the message ID from the given data dictionary
    fn get_id(dict: Option<&Dictionary>) -> Option<u16>;

    /// Get the message name
    // TODO: this could be an associated constant?
    fn get_name() -> &'static str;

    /// Decode this message type from the top of `input`.
    fn decode<'a>(input: &mut &'a [u8]) -> Result<Self::Pod<'a>, MessageDecodeError>;

    /// Get a list of field names and types
    fn fields() -> Vec<(&'static str, FieldType)>;
}

/// Marker trait for messages with an oid field
pub trait WithOid: 'static {}

/// Marker trait for messages without an oid field
pub trait WithoutOid: 'static {}

/// Represents an encoded message, with a type-level link to the message kind
pub struct EncodedMessage<M> {
    pub payload: FrontTrimmableBuffer,
    pub _message_kind: std::marker::PhantomData<M>,
}

/// Wraps a `Vec<u8>` allowing removal of front bytes in a zero-copy way
pub struct FrontTrimmableBuffer {
    pub content: Vec<u8>,
    pub offset: usize,
}

impl FrontTrimmableBuffer {
    /// Get the rest of the buffer as a slice
    pub fn as_slice(&self) -> &[u8] {
        &self.content[self.offset..]
    }
}

/// Holds the format of a `output()` style debug message
#[derive(Debug)]
pub struct OutputFormat {
    parts: Vec<FormatBlock>,
}

impl OutputFormat {
    /// Format the given fields according to this output format.
    pub fn format<'a>(&self, mut fields: impl Iterator<Item = &'a FieldValue>) -> String {
        let mut buf = String::new();
        for part in &self.parts {
            match part {
                FormatBlock::Static(s) => buf.push_str(s),
                FormatBlock::Field => {
                    if let Some(v) = fields.next() {
                        std::fmt::write(&mut buf, format_args!("{v}")).ok();
                    }
                }
            }
        }
        buf
    }
}

/// Part of an [`OutputFormat`].
#[derive(Debug)]
enum FormatBlock {
    Static(String),
    Field,
}

/// Format the given name and type pairs as a printf style declaration.
pub(crate) fn format_command_args<'a>(
    fields: impl Iterator<Item = (&'a str, FieldType)>,
) -> String {
    let mut buf = String::new();
    for (idx, (name, ty)) in fields.enumerate() {
        if idx != 0 {
            buf.push(' ');
        }
        buf.push_str(name);
        buf.push('=');
        buf.push_str(match ty {
            FieldType::U32 => "%u",
            FieldType::I32 => "%i",
            FieldType::U16 => "%hu",
            FieldType::I16 => "%hi",
            FieldType::U8 => "%c",
            FieldType::String => "%s",
            FieldType::ByteArray => "%*s",
        });
    }
    buf
}

/// An error enountered when parsing a message format string
#[derive(thiserror::Error, Debug)]
pub enum MessageSkipperError {
    #[error("invalid argument format: {0}")]
    InvalidArgumentFormat(String),

    #[error("unknown type '{1}' for argument '{0}'")]
    UnknownType(String, String),

    #[error("invalid format field type '%{0}'")]
    InvalidFormatFieldType(String),
}

impl std::fmt::Debug for FrontTrimmableBuffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Debug::fmt(self.as_slice(), f)
    }
}

impl std::fmt::Debug for MessageParser {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_map()
            .entry(&"name", &self.name)
            .entry(&"fields", &self.fields)
            .finish()
    }
}

impl<R: Message> std::fmt::Debug for EncodedMessage<R> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EncodedMessage")
            .field("kind", &R::get_name())
            .field("payload", &self.payload)
            .finish()
    }
}