// Copyright 2025 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only

use std::cmp::min;

use super::*;
use crate::proto::pq_ratchet as pqrpb;
use crate::{Error, SerializedMessage, Version};
use num_enum::IntoPrimitive;

impl States {
    pub fn into_pb(self) -> pqrpb::V1State {
        pqrpb::V1State {
            inner_state: Some(match self {
                // send_ek
                Self::KeysUnsampled(state) => {
                    pqrpb::v1_state::InnerState::KeysUnsampled(state.into_pb())
                }
                Self::KeysSampled(state) => {
                    pqrpb::v1_state::InnerState::KeysSampled(state.into_pb())
                }
                Self::HeaderSent(state) => pqrpb::v1_state::InnerState::HeaderSent(state.into_pb()),
                Self::Ct1Received(state) => {
                    pqrpb::v1_state::InnerState::Ct1Received(state.into_pb())
                }
                Self::EkSentCt1Received(state) => {
                    pqrpb::v1_state::InnerState::EkSentCt1Received(state.into_pb())
                }

                // send_ct
                Self::NoHeaderReceived(state) => {
                    pqrpb::v1_state::InnerState::NoHeaderReceived(state.into_pb())
                }
                Self::HeaderReceived(state) => {
                    pqrpb::v1_state::InnerState::HeaderReceived(state.into_pb())
                }
                Self::Ct1Sampled(state) => pqrpb::v1_state::InnerState::Ct1Sampled(state.into_pb()),
                Self::EkReceivedCt1Sampled(state) => {
                    pqrpb::v1_state::InnerState::EkReceivedCt1Sampled(state.into_pb())
                }
                Self::Ct1Acknowledged(state) => {
                    pqrpb::v1_state::InnerState::Ct1Acknowledged(state.into_pb())
                }
                Self::Ct2Sampled(state) => pqrpb::v1_state::InnerState::Ct2Sampled(state.into_pb()),
            }),
        }
    }

    pub fn from_pb(pb: pqrpb::V1State) -> Result<Self, Error> {
        Ok(match pb.inner_state {
            // send_ek
            Some(pqrpb::v1_state::InnerState::KeysUnsampled(pb)) => {
                Self::KeysUnsampled(send_ek::KeysUnsampled::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::KeysSampled(pb)) => {
                Self::KeysSampled(send_ek::KeysSampled::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::HeaderSent(pb)) => {
                Self::HeaderSent(send_ek::HeaderSent::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::Ct1Received(pb)) => {
                Self::Ct1Received(send_ek::Ct1Received::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::EkSentCt1Received(pb)) => {
                Self::EkSentCt1Received(send_ek::EkSentCt1Received::from_pb(pb)?)
            }

            // send_ct
            Some(pqrpb::v1_state::InnerState::NoHeaderReceived(pb)) => {
                Self::NoHeaderReceived(send_ct::NoHeaderReceived::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::HeaderReceived(pb)) => {
                Self::HeaderReceived(send_ct::HeaderReceived::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::Ct1Sampled(pb)) => {
                Self::Ct1Sampled(send_ct::Ct1Sampled::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::EkReceivedCt1Sampled(pb)) => {
                Self::EkReceivedCt1Sampled(send_ct::EkReceivedCt1Sampled::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::Ct1Acknowledged(pb)) => {
                Self::Ct1Acknowledged(send_ct::Ct1Acknowledged::from_pb(pb)?)
            }
            Some(pqrpb::v1_state::InnerState::Ct2Sampled(pb)) => {
                Self::Ct2Sampled(send_ct::Ct2Sampled::from_pb(pb)?)
            }

            _ => {
                return Err(Error::StateDecode);
            }
        })
    }
}

#[derive(IntoPrimitive)]
#[repr(u8)]
enum MessageType {
    None = 0,
    Hdr = 1,
    Ek = 2,
    EkCt1Ack = 3,
    Ct1Ack = 4,
    Ct1 = 5,
    Ct2 = 6,
}

impl TryFrom<u8> for MessageType {
    type Error = String;
    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(MessageType::None),
            1 => Ok(MessageType::Hdr),
            2 => Ok(MessageType::Ek),
            3 => Ok(MessageType::EkCt1Ack),
            4 => Ok(MessageType::Ct1Ack),
            5 => Ok(MessageType::Ct1),
            6 => Ok(MessageType::Ct2),
            _ => Err("Expected a number between 0 and 6".to_owned()),
        }
    }
}

impl MessageType {
    fn from_payload(mp: &MessagePayload) -> Self {
        match mp {
            MessagePayload::None => Self::None,
            MessagePayload::Hdr(_) => Self::Hdr,
            MessagePayload::Ek(_) => Self::Ek,
            MessagePayload::EkCt1Ack(_) => Self::EkCt1Ack,
            MessagePayload::Ct1Ack(_) => Self::Ct1Ack,
            MessagePayload::Ct1(_) => Self::Ct1,
            MessagePayload::Ct2(_) => Self::Ct2,
        }
    }
}

const MAX_VARINT_BYTES_LEN: usize = 10;

fn encode_varint(mut a: u64, into: &mut SerializedMessage) {
    for _i in 0..MAX_VARINT_BYTES_LEN {
        let byte = (a & 0x7F) as u8;
        if a < 0x80 {
            into.push(byte);
            break;
        } else {
            into.push(0x80 | byte);
            a >>= 7;
        }
    }
}

#[hax_lib::ensures(|res| *at <= *future(at) && if res.is_ok() { *at < from.len() && *future(at) <= from.len() } else { true })]
fn decode_varint(from: &SerializedMessage, at: &mut usize) -> Result<u64, Error> {
    let mut out = 0u64;

    let mut i: usize = 0;
    // Helps prevent return in while loop for Hax
    let mut done = false;
    let start_at: usize = *at;
    if start_at >= from.len() {
        return Err(Error::MsgDecode);
    }

    let max_i = min(MAX_VARINT_BYTES_LEN, from.len() - start_at);

    while i < max_i && !done {
        hax_lib::loop_invariant!(i <= max_i && *at == start_at);
        hax_lib::loop_decreases!(max_i - i);

        let byte = from[start_at + i];
        out |= ((byte as u64) & 0x7f) << (7 * i as i32);

        i += 1;
        done = (byte & 0x80) == 0;
    }

    if done {
        *at += i;
        Ok(out)
    } else {
        Err(Error::MsgDecode)
    }
}

fn encode_chunk(c: &Chunk, into: &mut SerializedMessage) {
    encode_varint(c.index as u64, into);
    hax_lib::assume!(into.len() < usize::MAX - 32);
    into.extend_from_slice(&c.data[..]);
}

fn decode_chunk(from: &SerializedMessage, at: &mut usize) -> Result<Chunk, Error> {
    let index = decode_varint(from, at)?;
    let start = *at;
    hax_lib::assume!(*at < usize::MAX - 32);
    *at += 32;
    if *at > from.len() || index > 65535 {
        return Err(Error::MsgDecode);
    }
    Ok(Chunk {
        index: index as u16,
        data: from[start..*at].try_into().expect("correct size"),
    })
}

#[hax_lib::attributes]
impl Message {
    /// Serialize a message.
    ///
    /// Messages are serialized as:
    ///
    ///   [version]      - 1 byte
    ///   [epoch]        - varint, 1-10 bytes
    ///   [index]        - varint, 1-5 bytes
    ///   [message_type] - 1 byte
    ///
    /// Many of the message types also have a data chunk concatenated to them, of
    /// the form:
    ///
    ///   [index]        - varint, 1-3 bytes
    ///   [chunk_data]   - 32 bytes
    #[hax_lib::ensures(|res| res.len() > 0 && res[0] == Version::V1.into())]
    pub fn serialize(&self, index: u32) -> SerializedMessage {
        let mut into = Vec::with_capacity(40);
        into.push(Version::V1.into());
        encode_varint(self.epoch, &mut into);
        encode_varint(index as u64, &mut into);
        into.push(MessageType::from_payload(&self.payload).into());
        encode_chunk(
            match &self.payload {
                MessagePayload::Hdr(ref chunk) => chunk,
                MessagePayload::Ek(ref chunk) => chunk,
                MessagePayload::EkCt1Ack(ref chunk) => chunk,
                MessagePayload::Ct1(ref chunk) => chunk,
                MessagePayload::Ct2(ref chunk) => chunk,
                _ => {
                    // This assumption could be proven with post-conditions on encode_varint
                    hax_lib::assume!(into.len() > 0 && into[0] == Version::V1.into());
                    return into;
                }
            },
            &mut into,
        );
        // This assumption could be proven with post-conditions on encode_varint and encode_chunk
        hax_lib::assume!(into.len() > 0 && into[0] == Version::V1.into());
        into
    }

    #[hax_lib::ensures(|res| if let Ok((msg, _index, at)) = res { msg.epoch > 0 && at <= from.len() } else { true })]
    pub fn deserialize(from: &SerializedMessage) -> Result<(Self, u32, usize), Error> {
        if from.is_empty() || from[0] != Version::V1.into() {
            return Err(Error::MsgDecode);
        }
        let mut at = 1usize;
        let epoch = decode_varint(from, &mut at)? as Epoch;
        if epoch == 0 {
            return Err(Error::MsgDecode);
        }
        let index: u32 = decode_varint(from, &mut at)?
            .try_into()
            .map_err(|_| Error::MsgDecode)?;
        if at >= from.len() {
            return Err(Error::MsgDecode);
        }
        let msg_type = MessageType::try_from(from[at]).map_err(|_| Error::MsgDecode)?;
        at += 1;
        let payload = match msg_type {
            MessageType::None => MessagePayload::None,
            MessageType::Ct1Ack => MessagePayload::Ct1Ack(true),
            MessageType::Hdr => MessagePayload::Hdr(decode_chunk(from, &mut at)?),
            MessageType::Ek => MessagePayload::Ek(decode_chunk(from, &mut at)?),
            MessageType::EkCt1Ack => MessagePayload::EkCt1Ack(decode_chunk(from, &mut at)?),
            MessageType::Ct1 => MessagePayload::Ct1(decode_chunk(from, &mut at)?),
            MessageType::Ct2 => MessagePayload::Ct2(decode_chunk(from, &mut at)?),
        };
        // We allow for there to be additional trailing data in the message, so it's
        // possible that `at < from.len()`.  This allows for us to potentially
        // upgrade sessions in future versions of the protocol.
        Ok((Self { epoch, payload }, index, at))
    }
}

#[cfg(test)]
mod test {
    use super::{decode_varint, encode_varint};
    use rand::RngCore;
    use rand::TryRngCore;
    use rand_core::OsRng;

    #[test]
    fn encoding_varint() {
        let mut v = vec![];
        encode_varint(0x012C, &mut v);
        assert_eq!(&v, &[0xAC, 0x02][..]);
    }

    #[test]
    fn decoding_varint() {
        let v = vec![0xFF, 0xAC, 0x02, 0xFF];
        let mut at = 1usize;
        assert_eq!(0x012C, decode_varint(&v, &mut at).unwrap());
        assert_eq!(at, 3, "at <= v.len()");
    }

    #[test]
    fn decoding_varint_zero() {
        let v = vec![0x00];
        let mut at = 0usize;
        assert_eq!(0x0, decode_varint(&v, &mut at).unwrap());
        assert_eq!(at, 1, "at <= v.len()");
    }

    #[test]
    fn roundtrip_varint() {
        let mut rng = OsRng.unwrap_err();
        for _i in 0..10000 {
            let u = rng.next_u64();
            let mut v = vec![];
            encode_varint(u, &mut v);
            let mut at = 0usize;
            assert_eq!(u, decode_varint(&v, &mut at).unwrap());
            assert_eq!(at, v.len());
        }
    }
}
