//! The wire format of a message: one byte for the format's version (`VERSION`), one
//! saying what the message is (`MessageKind`), then the message itself in MessagePack, as derived from its type in `transport` (structs as
//! arrays, fields in declaration order). The codecs below make hashes, IDs and bytes
//! compact (MessagePack binaries, not arrays of numbers) and check their lengths.
//!
//! The wire format may change freely until the first release: there is no compatibility
//! with older nodes to keep. But every change bumps `VERSION`, so nodes on different
//! builds tell each other apart instead of misreading each other's messages. The
//! transport's topic says which application a message is for, not which version.

use thiserror::Error;

use super::{MessageKind, ParsedMessage};

/// The wire format's version, the first byte of every payload. Bump it with every change
/// to a message type or a codec (`golden_bytes` fails then).
pub(crate) const VERSION: u8 = 2;

/// What the second byte of a payload says.
fn tag(kind: MessageKind) -> u8 {
    match kind {
        MessageKind::Hello => 0,
        MessageKind::HelloReply => 1,
        MessageKind::Sync => 2,
        MessageKind::Ping => 3,
        MessageKind::Pong => 4,
        MessageKind::Blob => 5,
    }
}

#[derive(Error, Debug)]
pub enum WireError {
    #[error("empty payload")]
    Empty,

    #[error("wire format version {0} (we speak version {VERSION})")]
    Version(u8),

    #[error("unknown message kind {0}")]
    UnknownKind(u8),

    #[error("MessagePack serialization failed: {0}")]
    Encode(#[from] rmp_serde::encode::Error),

    #[error("MessagePack deserialization failed: {0}")]
    Decode(#[from] rmp_serde::decode::Error),
}

/// The payload of `message`.
pub(crate) fn encode(message: &ParsedMessage) -> Result<Vec<u8>, WireError> {
    let mut payload = vec![VERSION, tag(message.kind())];
    match message {
        ParsedMessage::Hello(hello) => rmp_serde::encode::write(&mut payload, hello)?,
        ParsedMessage::HelloReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
        ParsedMessage::Sync(envelope) => rmp_serde::encode::write(&mut payload, envelope)?,
        ParsedMessage::Ping(keepalive) | ParsedMessage::Pong(keepalive) => {
            rmp_serde::encode::write(&mut payload, keepalive)?
        }
        ParsedMessage::Blob(blob) => rmp_serde::encode::write(&mut payload, blob)?,
    }
    Ok(payload)
}

/// The message in `payload` (as made by [`encode`]).
pub(crate) fn decode(payload: &[u8]) -> Result<ParsedMessage, WireError> {
    let (&version, payload) = payload.split_first().ok_or(WireError::Empty)?;
    if version != VERSION {
        return Err(WireError::Version(version));
    }
    let (&kind, body) = payload.split_first().ok_or(WireError::Empty)?;
    let is = |wanted: MessageKind| kind == tag(wanted);
    Ok(if is(MessageKind::Hello) {
        ParsedMessage::Hello(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::HelloReply) {
        ParsedMessage::HelloReply(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Sync) {
        ParsedMessage::Sync(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Ping) {
        ParsedMessage::Ping(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Pong) {
        ParsedMessage::Pong(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Blob) {
        ParsedMessage::Blob(rmp_serde::from_slice(body)?)
    } else {
        return Err(WireError::UnknownKind(kind));
    })
}

/// Field codecs (`#[serde(with = "...")]`) for the message types.
pub(crate) mod codec {
    use serde::de::Error as _;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use serde_bytes::{ByteBuf, Bytes};

    /// Exactly `N` bytes, or a decoding error naming `what`.
    fn fixed<const N: usize, E: serde::de::Error>(bytes: &[u8], what: &str) -> Result<[u8; N], E> {
        bytes
            .try_into()
            .map_err(|_| E::custom(format!("{what} must be {N} bytes")))
    }

    /// A run ID (`RunId`).
    pub(crate) mod run_id {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(value: &[u8; 16], s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 16], D::Error> {
            fixed(&ByteBuf::deserialize(d)?, "a run id")
        }
    }

    /// A 32-byte hash.
    pub(crate) mod bytes32 {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(value: &[u8; 32], s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 32], D::Error> {
            fixed(&ByteBuf::deserialize(d)?, "a hash")
        }
    }

    /// 32-byte hashes.
    pub(crate) mod hashes32 {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(
            value: &[[u8; 32]],
            s: S,
        ) -> Result<S::Ok, S::Error> {
            s.collect_seq(value.iter().map(|hash| Bytes::new(hash)))
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<Vec<[u8; 32]>, D::Error> {
            Vec::<ByteBuf>::deserialize(d)?
                .iter()
                .map(|hash| fixed(hash, "a hash"))
                .collect()
        }
    }

    /// An Automerge change hash.
    pub(crate) mod change_hash {
        use super::*;
        use automerge::ChangeHash;

        pub(crate) fn serialize<S: Serializer>(
            value: &ChangeHash,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            Bytes::new(&value.0).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<ChangeHash, D::Error> {
            Ok(ChangeHash(fixed(
                &ByteBuf::deserialize(d)?,
                "a change hash",
            )?))
        }
    }

    /// An epoch: generation and ID.
    pub(crate) mod epoch {
        use super::*;
        use crate::document::replica::Epoch;

        pub(crate) fn serialize<S: Serializer>(value: &Epoch, s: S) -> Result<S::Ok, S::Error> {
            (value.generation, Bytes::new(&value.id)).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Epoch, D::Error> {
            let (generation, id): (u64, ByteBuf) = Deserialize::deserialize(d)?;
            Ok(Epoch {
                generation,
                id: fixed(&id, "an epoch id")?,
            })
        }
    }

    /// Heads as a keepalive carries them: listed (at most `MAX_LISTED_HEADS`, digest
    /// empty), or as a digest (none listed).
    pub(crate) mod heads {
        use super::*;
        use crate::transport::{HeadsSummary, MAX_LISTED_HEADS};
        use automerge::ChangeHash;

        pub(crate) fn serialize<S: Serializer>(
            value: &HeadsSummary,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            let (listed, digest): (Vec<&Bytes>, &Bytes) = match value {
                HeadsSummary::Heads(heads) => (
                    heads.iter().map(|head| Bytes::new(&head.0)).collect(),
                    Bytes::new(&[]),
                ),
                HeadsSummary::Digest { digest, sample } => (
                    sample.iter().map(|head| Bytes::new(&head.0)).collect(),
                    Bytes::new(digest),
                ),
            };
            (listed, digest).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<HeadsSummary, D::Error> {
            let (listed, digest): (Vec<ByteBuf>, ByteBuf) = Deserialize::deserialize(d)?;
            if listed.len() > MAX_LISTED_HEADS {
                return Err(D::Error::custom(format!(
                    "more than {MAX_LISTED_HEADS} heads listed"
                )));
            }
            let listed = listed
                .iter()
                .map(|head| -> Result<ChangeHash, D::Error> {
                    Ok(ChangeHash(fixed(head, "a head")?))
                })
                .collect::<Result<Vec<_>, _>>()?;
            if digest.is_empty() {
                return Ok(HeadsSummary::Heads(listed));
            }
            Ok(HeadsSummary::Digest {
                digest: fixed(&digest, "a heads digest")?,
                sample: listed,
            })
        }
    }

    /// An Automerge sync message, in Automerge's own encoding.
    pub(crate) mod sync_message {
        use super::*;
        use automerge::sync::Message;

        pub(crate) fn serialize<S: Serializer>(value: &Message, s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(&value.clone().encode()).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Message, D::Error> {
            Message::decode(&ByteBuf::deserialize(d)?)
                .map_err(|error| D::Error::custom(format!("Automerge decode failed: {error:?}")))
        }
    }

    /// Shared bytes (`bytes::Bytes`).
    pub(crate) mod shared_bytes {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(
            value: &bytes::Bytes,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<bytes::Bytes, D::Error> {
            Ok(ByteBuf::deserialize(d)?.into_vec().into())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::replica::Epoch;
    use crate::transport::{
        BlobBody, BlobId, BlobMessage, HeadsSummary, HelloReply, HelloRequest, Keepalive,
        MAX_LISTED_HEADS,
    };
    use automerge::ChangeHash;

    #[test]
    fn handshakes_round_trip() {
        let epoch = Epoch {
            generation: 2,
            id: [3; 32],
        };
        let hello = ParsedMessage::Hello(HelloRequest {
            swarm_id: Some("swarm".into()),
            epoch,
            epoch_seq: 1,
            auth: vec![1, 2, 3],
            ready: true,
            run_id: [4; 16],
            handshake: 7,
        });
        let ParsedMessage::Hello(decoded) = decode(&encode(&hello).unwrap()).unwrap() else {
            panic!("not a HELLO");
        };
        assert_eq!(decoded.swarm_id.as_deref(), Some("swarm"));
        assert_eq!((decoded.epoch, decoded.epoch_seq), (epoch, 1));
        assert_eq!(decoded.auth, vec![1, 2, 3]);
        assert_eq!(decoded.run_id, [4; 16]);
        assert_eq!(decoded.handshake, 7);

        let reply = ParsedMessage::HelloReply(HelloReply {
            swarm_id: "swarm".into(),
            peers: vec!["p".into()],
            epoch,
            epoch_seq: 1,
            auth: Vec::new(),
            ready: false,
            run_id: [4; 16],
            away_ms: Some(5),
            sender_run_id: [6; 16],
        });
        let ParsedMessage::HelloReply(decoded) = decode(&encode(&reply).unwrap()).unwrap() else {
            panic!("not a HELLO_REPLY");
        };
        assert_eq!(decoded.peers, vec!["p".to_owned()]);
        assert_eq!(decoded.away_ms, Some(5));
        assert_eq!(decoded.sender_run_id, [6; 16]);
    }

    #[test]
    fn keepalives_keep_their_kind_and_heads() {
        let keepalive = Keepalive {
            heads: HeadsSummary::Digest {
                digest: [8; 32],
                sample: vec![ChangeHash([9; 32])],
            },
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: false,
        };
        let decoded = decode(&encode(&ParsedMessage::Pong(keepalive.clone())).unwrap()).unwrap();
        let ParsedMessage::Pong(decoded) = decoded else {
            panic!("not a PONG");
        };
        assert_eq!(decoded.heads, keepalive.heads);
        assert_eq!(decoded.heads.listed(), [ChangeHash([9; 32])]);
    }

    #[test]
    fn too_many_listed_heads_are_refused() {
        let keepalive = Keepalive {
            heads: HeadsSummary::Heads(
                (0..=MAX_LISTED_HEADS as u8)
                    .map(|n| ChangeHash([n; 32]))
                    .collect(),
            ),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: true,
        };
        assert!(decode(&encode(&ParsedMessage::Ping(keepalive)).unwrap()).is_err());
    }

    #[test]
    fn unknown_kinds_and_bad_payloads_are_refused() {
        assert!(matches!(decode(&[]), Err(WireError::Empty)));
        assert!(matches!(decode(&[VERSION]), Err(WireError::Empty)));
        assert!(matches!(
            decode(&[VERSION, 99]),
            Err(WireError::UnknownKind(99))
        ));
        assert!(matches!(
            decode(&[VERSION + 1, 0]),
            Err(WireError::Version(version)) if version == VERSION + 1
        ));

        let blob = ParsedMessage::Blob(BlobMessage {
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            id: BlobId::Content([1; 32]),
            body: BlobBody::Missing,
        });
        let payload = encode(&blob).unwrap();
        assert!(matches!(decode(&payload).unwrap(), ParsedMessage::Blob(_)));
        // A payload cut short doesn't decode.
        assert!(decode(&payload[..payload.len() - 1]).is_err());
        // Nor one of another kind.
        let mut relabelled = payload.clone();
        relabelled[1] = tag(MessageKind::Hello);
        assert!(decode(&relabelled).is_err());
    }

    fn hex(bytes: &[u8]) -> String {
        bytes.iter().map(|byte| format!("{byte:02x}")).collect()
    }

    /// The exact bytes of a few messages. If this fails after a change to a message type
    /// or a codec, the change is a new wire format: bump `VERSION` and take the new bytes.
    #[test]
    fn golden_bytes() {
        let epoch = Epoch {
            generation: 2,
            id: [3; 32],
        };
        let epoch_hex = format!("9202c420{}", "03".repeat(32));
        let run = |byte: &str| format!("c410{}", byte.repeat(16));

        let hello = ParsedMessage::Hello(HelloRequest {
            swarm_id: Some("swarm".into()),
            epoch,
            epoch_seq: 1,
            auth: vec![1, 2, 3],
            ready: true,
            run_id: [4; 16],
            handshake: 7,
        });
        assert_eq!(
            hex(&encode(&hello).unwrap()),
            format!("020097a5737761726d{epoch_hex}01c403010203c3{}07", run("04"))
        );

        let reply = ParsedMessage::HelloReply(HelloReply {
            swarm_id: "swarm".into(),
            peers: vec!["p".into()],
            epoch,
            epoch_seq: 1,
            auth: Vec::new(),
            ready: false,
            run_id: [4; 16],
            away_ms: Some(5),
            sender_run_id: [6; 16],
        });
        assert_eq!(
            hex(&encode(&reply).unwrap()),
            format!(
                "020199a5737761726d91a170{epoch_hex}01c400c2{}05{}",
                run("04"),
                run("06")
            )
        );

        let ping = ParsedMessage::Ping(Keepalive {
            heads: HeadsSummary::Heads(Vec::new()),
            epoch,
            epoch_seq: 1,
            ready: true,
            syncs_with_you: false,
        });
        assert_eq!(
            hex(&encode(&ping).unwrap()),
            format!("0203959290c400{epoch_hex}01c3c2")
        );
    }
}