//! Handhsake handling

pub mod dirsync;
#[cfg(test)]
mod tests;

use crate::{
    auth::ServiceID,
    connection::{self, Connection, IDRecv, ProtocolVersion},
    enc::{
        asym::{KeyID, PrivKey, PubKey},
        sym::{HeadLen, TagLen},
    },
};
use ::num_traits::FromPrimitive;

/// Handshake errors
#[derive(::thiserror::Error, Debug, Copy, Clone)]
#[non_exhaustive]
pub enum Error {
    /// Error while parsing the handshake packet
    /// TODO: more detailed parsing errors
    #[error("not an handshake packet")]
    Parsing,
    /// No such Key ID
    #[error("unknown key id")]
    UnknownKeyID,
    /// Key error
    #[error("key: {0:?}")]
    Key(#[from] crate::enc::Error),
    /// Not enough data
    #[error("not enough data")]
    NotEnoughData,
    /// Could not find common cryptography
    #[error("Negotiation of keys/hkdfs/ciphers failed")]
    Negotiation,
    /// Could not generate Keys
    #[error("Key generation failed")]
    KeyGeneration,
    /// Too many client handshakes currently running
    #[error("Too many client handshakes")]
    TooManyClientHandshakes,
}

/// List of possible handshakes
#[derive(
    ::num_derive::FromPrimitive,
    Debug,
    Clone,
    Copy,
    PartialEq,
    ::strum_macros::EnumString,
    ::strum_macros::IntoStaticStr,
)]
#[repr(u8)]
pub enum HandshakeID {
    /// 1-RTT Directory synchronized handshake. Fast, no forward secrecy
    #[strum(serialize = "directory_synchronized")]
    DirectorySynchronized = 0,
    /// 2-RTT Stateful exchange. Little DDos protection
    #[strum(serialize = "stateful")]
    Stateful,
    /// 3-RTT stateless exchange. Forward secrecy and ddos protection
    #[strum(serialize = "stateless")]
    Stateless,
}
impl HandshakeID {
    /// The length of the serialized field
    pub const fn len() -> usize {
        1
    }
}

pub(crate) struct HandshakeServer {
    pub id: KeyID,
    pub key: PrivKey,
}

pub(crate) struct HandshakeClient {
    pub service_id: ServiceID,
    pub service_conn_id: IDRecv,
    pub connection: Connection,
    pub timeout: Option<::tokio::task::JoinHandle<()>>,
}

/// Tracks the keys used by the client and the handshake
/// they are associated with
pub(crate) struct HandshakeClientList {
    used: Vec<::bitmaps::Bitmap<1024>>, // index = KeyID
    keys: Vec<Option<(PrivKey, PubKey)>>,
    list: Vec<Option<HandshakeClient>>,
}

impl HandshakeClientList {
    pub(crate) fn new() -> Self {
        Self {
            used: [::bitmaps::Bitmap::<1024>::new()].to_vec(),
            keys: Vec::with_capacity(16),
            list: Vec::with_capacity(16),
        }
    }
    pub(crate) fn get(&self, id: KeyID) -> Option<&HandshakeClient> {
        if id.0 as usize >= self.list.len() {
            return None;
        }
        self.list[id.0 as usize].as_ref()
    }
    pub(crate) fn remove(&mut self, id: KeyID) -> Option<HandshakeClient> {
        if id.0 as usize >= self.list.len() {
            return None;
        }
        let used_vec_idx = id.0 as usize / 1024;
        let used_bitmap_idx = id.0 as usize % 1024;
        let used_iter = match self.used.get_mut(used_vec_idx) {
            Some(used_iter) => used_iter,
            None => return None,
        };
        used_iter.set(used_bitmap_idx, false);
        self.keys[id.0 as usize] = None;
        let mut owned = None;
        ::core::mem::swap(&mut self.list[id.0 as usize], &mut owned);
        owned
    }
    pub(crate) fn add(
        &mut self,
        priv_key: PrivKey,
        pub_key: PubKey,
        service_id: ServiceID,
        service_conn_id: IDRecv,
        connection: Connection,
    ) -> Result<(KeyID, &mut HandshakeClient), ()> {
        let maybe_free_key_idx =
            self.used.iter().enumerate().find_map(|(idx, bmap)| {
                match bmap.first_false_index() {
                    Some(false_idx) => Some(((idx * 1024), false_idx)),
                    None => None,
                }
            });
        let free_key_idx = match maybe_free_key_idx {
            Some((idx, false_idx)) => {
                let free_key_idx = idx * 1024 + false_idx;
                if free_key_idx > KeyID::MAX as usize {
                    return Err(());
                }
                self.used[idx].set(false_idx, true);
                free_key_idx
            }
            None => {
                let mut bmap = ::bitmaps::Bitmap::<1024>::new();
                bmap.set(0, true);
                self.used.push(bmap);
                self.used.len() * 1024
            }
        };
        if self.keys.len() >= free_key_idx {
            self.keys.push(None);
            self.list.push(None);
        }
        self.keys[free_key_idx] = Some((priv_key, pub_key));
        self.list[free_key_idx] = Some(HandshakeClient {
            service_id,
            service_conn_id,
            connection,
            timeout: None,
        });
        Ok((
            KeyID(free_key_idx as u16),
            self.list[free_key_idx].as_mut().unwrap(),
        ))
    }
}

/// Parsed handshake
#[derive(Debug, Clone, PartialEq)]
pub enum HandshakeData {
    /// Directory synchronized handhsake
    DirSync(dirsync::DirSync),
}

impl HandshakeData {
    /// actual length of the handshake data
    pub fn len(&self) -> usize {
        match self {
            HandshakeData::DirSync(d) => d.len(),
        }
    }
    /// Serialize into raw 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],
    ) {
        match self {
            HandshakeData::DirSync(d) => d.serialize(head_len, tag_len, out),
        }
    }
}

/// Kind of handshake
#[derive(::num_derive::FromPrimitive, Debug, Clone, Copy)]
#[repr(u8)]
pub enum HandshakeKind {
    /// 1-RTT, Directory synchronized handshake
    /// Request
    DirSyncReq = 0,
    /// 1-RTT, Directory synchronized handshake
    /// Response
    DirSyncResp,
    /*
    /// 2-RTT, Stateful connection
    Stateful,
    ...
    /// 3 RTT, Stateless connection
    Stateless,
    ....
    */
}
impl HandshakeKind {
    /// Length of the serialized field
    pub const fn len() -> usize {
        1
    }
}

/// Parsed handshake
#[derive(Debug, Clone, PartialEq)]
pub struct Handshake {
    /// Fenrir Protocol version
    pub fenrir_version: ProtocolVersion,
    /// enum for the parsed data
    pub data: HandshakeData,
}

impl Handshake {
    /// Build new handshake from the data
    pub fn new(data: HandshakeData) -> Self {
        Handshake {
            fenrir_version: ProtocolVersion::V0,
            data,
        }
    }
    /// return the total length of the handshake
    pub fn len(&self) -> usize {
        ProtocolVersion::len() + HandshakeKind::len() + self.data.len()
    }
    const MIN_PKT_LEN: usize = 8;
    /// Parse the packet and return the parsed handshake
    pub fn deserialize(raw: &[u8]) -> Result<Self, Error> {
        if raw.len() < Self::MIN_PKT_LEN {
            return Err(Error::NotEnoughData);
        }
        let fenrir_version = match ProtocolVersion::from_u8(raw[0]) {
            Some(fenrir_version) => fenrir_version,
            None => return Err(Error::Parsing),
        };
        let handshake_kind = match HandshakeKind::from_u8(raw[1]) {
            Some(handshake_kind) => handshake_kind,
            None => return Err(Error::Parsing),
        };
        let data = match handshake_kind {
            HandshakeKind::DirSyncReq => dirsync::Req::deserialize(&raw[2..])?,
            HandshakeKind::DirSyncResp => {
                dirsync::Resp::deserialize(&raw[2..])?
            }
        };
        Ok(Self {
            fenrir_version,
            data,
        })
    }
    /// serialize the handshake into bytes
    /// NOTE: assumes that there is exactly as much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        tag_len: TagLen,
        out: &mut [u8],
    ) {
        out[0] = self.fenrir_version as u8;
        out[1] = match &self.data {
            HandshakeData::DirSync(d) => match d {
                dirsync::DirSync::Req(_) => HandshakeKind::DirSyncReq,
                dirsync::DirSync::Resp(_) => HandshakeKind::DirSyncResp,
            },
        } as u8;
        self.data.serialize(head_len, tag_len, &mut out[2..]);
    }
}

trait HandshakeParsing {
    fn deserialize(raw: &[u8]) -> Result<HandshakeData, Error>;
}