use crate::authenticated::data::Data;
use commonware_codec::{Buf, EncodeSize, Error, Read, ReadExt, Write};
use commonware_runtime::BufMut;

/// Prefix that identifies the message as a Data message.
pub const DATA_PREFIX: u8 = crate::authenticated::data::DATA_PREFIX; // 0
/// Prefix that identifies the message as a Ping message.
pub const PING_PREFIX: u8 = 1;

/// The messages that can be sent between peers.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
pub enum Message {
    Data(Data),
    Ping,
}

impl From<Data> for Message {
    fn from(data: Data) -> Self {
        Self::Data(data)
    }
}

impl EncodeSize for Message {
    fn encode_size(&self) -> usize {
        (match self {
            Self::Data(data) => data.encode_size(),
            Self::Ping => 0, // Ping has no payload
        }) + 1 // 1 bytes for Message discriminant
    }
}

impl Write for Message {
    fn write(&self, buf: &mut impl BufMut) {
        match self {
            Self::Data(data) => {
                DATA_PREFIX.write(buf); // Discriminant for Data
                data.write(buf);
            }
            Self::Ping => {
                PING_PREFIX.write(buf); // Discriminant for Ping
            }
        }
    }
}

impl Read for Message {
    type Cfg = usize; // Maximum amount of data to read

    fn read_cfg(buf: &mut impl Buf, max_data_length: &Self::Cfg) -> Result<Self, Error> {
        let message_type = <u8>::read(buf)?;
        match message_type {
            DATA_PREFIX => {
                let data = Data::read_cfg(buf, &(..=*max_data_length).into())?;
                Ok(Self::Data(data))
            }
            PING_PREFIX => Ok(Self::Ping),
            other => Err(Error::InvalidEnum(other)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::authenticated::{
        MAX_PAYLOAD_OVERHEAD,
        data::allocations::{allocation_calls_during, frame},
    };
    use commonware_codec::{Decode as _, Encode as _, Error};
    use commonware_runtime::IoBuf;
    use commonware_utils::Widen;

    #[test]
    fn test_data_prefix_value() {
        assert_eq!(DATA_PREFIX, 0);
    }

    #[test]
    fn test_max_payload_overhead() {
        let message = IoBuf::from(vec![0; 1 << 29]);
        let message_len = message.len();
        let payload = Message::Data(Data {
            channel: u64::MAX,
            message,
        });
        assert_eq!(
            payload.encode_size(),
            message_len + MAX_PAYLOAD_OVERHEAD as usize
        );
    }

    #[test]
    fn test_ping_fits_framing_overhead() {
        let overhead: usize = Widen::widen(MAX_PAYLOAD_OVERHEAD);
        assert!(Message::Ping.encode_size() <= overhead);
    }

    #[test]
    fn test_decode_data_within_limit() {
        let payload = Message::Data(Data {
            channel: 7,
            message: IoBuf::from(b"ping"),
        });
        let encoded = payload.encode();

        let decoded = Message::decode_cfg(encoded, &4).expect("within limit");
        match decoded {
            Message::Data(data) => {
                assert_eq!(data.channel, 7);
                assert_eq!(data.message, IoBuf::from(b"ping"));
            }
            other => panic!("unexpected message variant: {other:?}"),
        }
    }

    #[test]
    fn test_decode_data_exceeding_limit() {
        let payload = Message::Data(Data {
            channel: 9,
            message: IoBuf::from(b"hello"),
        });
        let encoded = payload.encode();

        let result = Message::decode_cfg(encoded, &4);
        assert!(matches!(result, Err(Error::InvalidLength(5))));
    }

    #[cfg(feature = "arbitrary")]
    mod conformance {
        use super::*;
        use commonware_codec::conformance::CodecConformance;

        commonware_conformance::conformance_tests! {
            CodecConformance<Message>,
        }
    }

    /// Decoding a data frame through the message type the peer actor uses neither copies the
    /// message nor allocates owner records for it.
    #[commonware_macros::test_group("miri")]
    #[test]
    fn test_decode_does_not_allocate() {
        let original = Data {
            channel: 12345,
            message: IoBuf::from(b"Hello, world!"),
        };
        let (frame, bytes) = frame(&Message::Data(original.clone()).encode());
        let (decoded, calls) =
            allocation_calls_during(|| Message::decode_cfg(frame, &1024).unwrap());
        assert_eq!(calls, 0);
        let Message::Data(decoded) = decoded else {
            panic!("expected a data frame");
        };
        assert_eq!(decoded, original);
        assert!(bytes.contains(&decoded.message.as_ptr()));
    }
}