use crate::proto::{Decode, DecodeError, Decoder, Encode, EncodeError, Encoder};

pub mod client;
pub mod config;
pub mod server;
pub mod signature;
pub mod stream;
pub mod symmetric;
pub mod util;

// v1 supports only one DH algo
pub const V1_KEX_X25519_DALEK: u8 = 0x10;

pub const V1_SIG_ED25519: u8 = 0x10;

pub const V1_CLIENT_RANDOM_LEN: usize = 16;

#[derive(Debug)]
pub enum ServerReject {
    NoReason,
    UnsupportedKexAlgorithm,
    UnsupportedCiphersuite,
    UnsupportedSigAlgorithm,
    InvalidParameter,
}

// v1 protocol
#[derive(Debug)]
pub struct ServerHello<'a> {
    // TODO server will send its public fingerprint and advertise
    //      supported asymmetric key formats
    pub symmetric_ciphersuites: &'a [u8],
    pub kex_algos: &'a [u8],
    pub sig_algos: &'a [u8],
}

#[derive(Debug)]
pub enum ClientNegotiationMessage<'a> {
    Hello {
        protocol: u8,
    },
    StartKex {
        kex_algorithm: u8,
        sig_algorithm: u8,
        ciphersuite: u8,
        client_random: &'a [u8; V1_CLIENT_RANDOM_LEN],
        key_data: &'a [u8],
    },
    DHPublicKey(&'a [u8]),
    Agreed,
}

#[derive(Debug)]
pub enum ServerNegotiationMessage<'a> {
    Hello(ServerHello<'a>),
    Reject(ServerReject),
    StartKex {
        sig_algorithm: u8,
        key_data: &'a [u8],
    },
    DHPublicKey(&'a [u8]),
    Agreed,
}

impl Encode for ServerHello<'_> {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        buffer.write_variable_bytes(self.kex_algos)?;
        buffer.write_variable_bytes(self.symmetric_ciphersuites)?;
        buffer.write_variable_bytes(self.sig_algos)?;
        Ok(())
    }
}

impl<'de> Decode<'de> for ServerHello<'de> {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let kex_algos = buffer.read_variable_bytes()?;
        let symmetric_ciphersuites = buffer.read_variable_bytes()?;
        let sig_algos = buffer.read_variable_bytes()?;
        Ok(Self {
            kex_algos,
            symmetric_ciphersuites,
            sig_algos,
        })
    }
}

impl ClientNegotiationMessage<'_> {
    const TAG_HELLO: u8 = 0x11;
    const TAG_START_KEX: u8 = 0x12;
    const TAG_DH_PUBLIC_KEY: u8 = 0x13;
    const TAG_AGREED: u8 = 0x14;
}

impl Encode for ClientNegotiationMessage<'_> {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        match *self {
            Self::Hello { protocol } => buffer.write(&[Self::TAG_HELLO, protocol]),
            Self::StartKex {
                kex_algorithm,
                sig_algorithm,
                ciphersuite,
                client_random,
                key_data,
            } => {
                buffer.write(&[
                    Self::TAG_START_KEX,
                    kex_algorithm,
                    sig_algorithm,
                    ciphersuite,
                ])?;
                buffer.write(client_random)?;
                buffer.write_variable_bytes(key_data)
            }
            Self::DHPublicKey(key_data) => {
                buffer.write(&[Self::TAG_DH_PUBLIC_KEY])?;
                buffer.write_variable_bytes(key_data)
            }
            Self::Agreed => buffer.write(&[Self::TAG_AGREED]),
        }
    }
}

impl<'de> Decode<'de> for ClientNegotiationMessage<'de> {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let tag = buffer.read_u8()?;
        match tag {
            Self::TAG_HELLO => {
                let protocol = buffer.read_u8()?;
                Ok(Self::Hello { protocol })
            }
            Self::TAG_START_KEX => {
                let kex_algorithm = buffer.read_u8()?;
                let sig_algorithm = buffer.read_u8()?;
                let ciphersuite = buffer.read_u8()?;
                let client_random = buffer.read_bytes(V1_CLIENT_RANDOM_LEN)?.try_into().unwrap();
                let key_data = buffer.read_variable_bytes()?;
                Ok(Self::StartKex {
                    kex_algorithm,
                    sig_algorithm,
                    ciphersuite,
                    client_random,
                    key_data,
                })
            }
            Self::TAG_DH_PUBLIC_KEY => buffer.read_variable_bytes().map(Self::DHPublicKey),
            Self::TAG_AGREED => Ok(Self::Agreed),
            _ => Err(DecodeError::InvalidMessage(
                "Invalid ClientNegotiationMessage tag",
            )),
        }
    }
}

impl ServerNegotiationMessage<'_> {
    const TAG_HELLO: u8 = 0x21;
    const TAG_REJECT: u8 = 0x22;
    const TAG_START_KEX: u8 = 0x23;
    const TAG_DH_PUBLIC_KEY: u8 = 0x24;
    const TAG_AGREED: u8 = 0x25;
}

impl Encode for ServerNegotiationMessage<'_> {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        match self {
            Self::Hello(hello) => {
                buffer.write(&[Self::TAG_HELLO])?;
                hello.encode(buffer)
            }
            Self::Reject(_reject) => todo!(),
            &Self::StartKex {
                sig_algorithm,
                key_data,
            } => {
                buffer.write(&[Self::TAG_START_KEX, sig_algorithm])?;
                buffer.write_variable_bytes(key_data)
            }
            Self::DHPublicKey(key_data) => {
                buffer.write(&[Self::TAG_DH_PUBLIC_KEY])?;
                buffer.write_variable_bytes(key_data)
            }
            Self::Agreed => buffer.write(&[Self::TAG_AGREED]),
        }
    }
}

impl<'de> Decode<'de> for ServerNegotiationMessage<'de> {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let tag = buffer.read_u8()?;
        match tag {
            Self::TAG_HELLO => ServerHello::decode(buffer).map(Self::Hello),
            Self::TAG_REJECT => todo!(),
            Self::TAG_START_KEX => {
                let sig_algorithm = buffer.read_u8()?;
                let key_data = buffer.read_variable_bytes()?;
                Ok(Self::StartKex {
                    sig_algorithm,
                    key_data,
                })
            }
            Self::TAG_DH_PUBLIC_KEY => buffer.read_variable_bytes().map(Self::DHPublicKey),
            Self::TAG_AGREED => Ok(Self::Agreed),

            _ => Err(DecodeError::InvalidMessage(
                "Invalid ServerNegotiationMessage tag",
            )),
        }
    }
}

pub fn sig_algo_name(sig: u8) -> Option<&'static str> {
    match sig {
        V1_SIG_ED25519 => Some("ed25519"),
        _ => None,
    }
}
