//
//! Raw packet handling, encryption, decryption, parsing

use crate::enc::{
    sym::{HeadLen, TagLen},
    Random,
};

/// Fenrir Connection id
/// 0 is special as it represents the handshake
/// Connection IDs are to be considered u64 little endian
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum ConnectionID {
    /// Connection id 0 represent the handshake
    Handshake,
    /// Non-zero id can represent any connection
    ID(::core::num::NonZeroU64),
}

impl ConnectionID {
    /// Set the conenction id to handshake
    pub fn new_handshake() -> Self {
        Self::Handshake
    }
    /// New id from u64. PLZ NON ZERO
    pub(crate) fn new_u64(raw: u64) -> Self {
        #[allow(unsafe_code)]
        unsafe {
            ConnectionID::ID(::core::num::NonZeroU64::new_unchecked(raw))
        }
    }
    pub(crate) fn as_u64(&self) -> u64 {
        match self {
            ConnectionID::Handshake => 0,
            ConnectionID::ID(id) => id.get(),
        }
    }
    /// New random service ID
    pub fn new_rand(rand: &Random) -> Self {
        let mut raw = [0; 8];
        let mut num = 0;
        while num == 0 {
            rand.fill(&mut raw);
            num = u64::from_le_bytes(raw);
        }
        #[allow(unsafe_code)]
        unsafe {
            ConnectionID::ID(::core::num::NonZeroU64::new_unchecked(num))
        }
    }
    /// Quick check to know if this is an handshake
    pub fn is_handshake(&self) -> bool {
        *self == ConnectionID::Handshake
    }
    /// length if the connection ID in bytes
    pub const fn len() -> usize {
        8
    }
    /// write the ID to a buffer
    pub fn serialize(&self, out: &mut [u8]) {
        assert!(out.len() == 8, "out buffer must be 8 bytes");
        match self {
            ConnectionID::Handshake => out[..].copy_from_slice(&[0; 8]),
            ConnectionID::ID(id) => {
                out[..].copy_from_slice(&id.get().to_le_bytes())
            }
        }
    }
}

impl From<u64> for ConnectionID {
    fn from(raw: u64) -> Self {
        if raw == 0 {
            ConnectionID::Handshake
        } else {
            #[allow(unsafe_code)]
            unsafe {
                ConnectionID::ID(::core::num::NonZeroU64::new_unchecked(raw))
            }
        }
    }
}

impl From<[u8; 8]> for ConnectionID {
    fn from(raw: [u8; 8]) -> Self {
        let raw_u64 = u64::from_le_bytes(raw);
        raw_u64.into()
    }
}

/// Enumerate the possible data in a fenrir packet
#[derive(Debug, Clone)]
pub enum PacketData {
    /// A parsed handshake packet
    Handshake(super::Handshake),
    /// Raw packet. we only have the connection ID and packet length
    Raw(usize),
}

impl PacketData {
    /// total length of the data in bytes
    pub fn len(&self) -> usize {
        match self {
            PacketData::Handshake(h) => h.len(),
            PacketData::Raw(len) => *len,
        }
    }
    /// serialize data into bytes
    /// NOTE: assumes that there is exactly asa much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        tag_len: TagLen,
        out: &mut [u8],
    ) {
        assert!(self.len() == out.len(), "PacketData: wrong buffer length");
        match self {
            PacketData::Handshake(h) => h.serialize(head_len, tag_len, out),
            PacketData::Raw(_) => {
                ::tracing::error!("Tried to serialize a raw PacketData!");
            }
        }
    }
}

const MIN_PACKET_BYTES: usize = 16;

/// Fenrir packet structure
#[derive(Debug, Clone)]
pub struct Packet {
    /// Id of the Fenrir connection.
    pub id: ConnectionID,
    /// actual data inside the packet
    pub data: PacketData,
}

impl Packet {
    /// New recevied packet, yet unparsed
    pub fn deserialize_id(raw: &[u8]) -> Result<Self, ()> {
        // TODO: proper min_packet length. 16 is too conservative.
        if raw.len() < MIN_PACKET_BYTES {
            return Err(());
        }
        let raw_id: [u8; 8] = (raw[..8]).try_into().expect("unreachable");
        Ok(Packet {
            id: raw_id.into(),
            data: PacketData::Raw(raw.len()),
        })
    }
    /// get the total length of the packet
    pub fn len(&self) -> usize {
        ConnectionID::len() + self.data.len()
    }
    /// serialize packet into buffer
    /// NOTE: assumes that there is exactly asa much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        tag_len: TagLen,
        out: &mut [u8],
    ) {
        assert!(
            out.len() > ConnectionID::len(),
            "Packet: not enough buffer to serialize"
        );
        self.id.serialize(&mut out[0..ConnectionID::len()]);
        self.data
            .serialize(head_len, tag_len, &mut out[ConnectionID::len()..]);
    }
}