use std::{
    io::{self, Read, Write},
    net::{SocketAddr, TcpStream},
};

use rand::RngCore;
use x25519_dalek::{EphemeralSecret, PublicKey};

use crate::{
    crypt::{
        self,
        signature::VerificationMethod,
        stream::ClientSocket,
        symmetric::{data::Iv, SymmetricCipher},
        V1_CLIENT_RANDOM_LEN,
    },
    proto::{Decode, DecodeError, Decoder, Encode, EncodeError, Encoder},
};

use super::{
    config::ClientConfig,
    signature::{self, SignatureMethod},
    symmetric, ClientNegotiationMessage, ServerNegotiationMessage,
};

pub(crate) struct Negotiation {
    stream: TcpStream,
    remote: SocketAddr,
    buffer: [u8; 512],
    config: ClientConfig,
    client_random: [u8; V1_CLIENT_RANDOM_LEN],
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),
    #[error("Message encode error: {0}")]
    Encode(#[from] EncodeError),
    #[error("Message decode error: {0}")]
    Decode(#[from] DecodeError),

    #[error("Signature/verification error: {0}")]
    Signature(#[from] signature::Error),
    #[error("Symmetric cipher error: {0}")]
    Symmetric(#[from] symmetric::Error),

    #[error("Unexpected server reply")]
    UnexpectedServerReply,
    #[error("Could not select a signature key")]
    NoSignatureMethods,
    #[error("Could not select a ciphersuite")]
    NoCiphersuites,
    #[error("Could not select a key exchange algorithm")]
    NoKexAlgorithms,
    #[error("Server public key rejected")]
    ServerKeyRejected,
}

impl Negotiation {
    pub fn new(stream: TcpStream, remote: SocketAddr, config: ClientConfig) -> Self {
        Self {
            stream,
            remote,
            buffer: [0; 512],
            config,
            client_random: [0; V1_CLIENT_RANDOM_LEN],
        }
    }

    fn fill_client_random(&mut self) {
        let mut rng = rand::rngs::OsRng;
        rng.fill_bytes(&mut self.client_random);
    }

    fn hello(&mut self, recv_buf: &mut [u8]) -> Result<(SignatureMethod, u8, u8), Error> {
        log::info!("Send ClientHello v1");
        self.send(None, &ClientNegotiationMessage::Hello { protocol: 1 })?;

        let hello = match self.recv(None, recv_buf)? {
            ServerNegotiationMessage::Hello(hello) => hello,
            ServerNegotiationMessage::Reject(_) => todo!(),
            _ => return Err(Error::UnexpectedServerReply),
        };

        log::info!("Server ciphersuites:");
        // for &cipher in hello.symmetric_ciphersuites {
        //     if let Some(name) = crypt::ciphersuite_name(cipher) {
        //         log::info!(" * {name:?} ({cipher:#x})");
        //     } else {
        //         log::info!(" * {cipher:#x}");
        //     }
        // }
        log::info!("Server signature algorithms:");
        for &sig in hello.sig_algos {
            if let Some(name) = crypt::sig_algo_name(sig) {
                log::info!(" * {name:?} ({sig:#x})");
            } else {
                log::info!(" * {sig:#x}");
            }
        }

        let signer = self
            .config
            .signature_keystore
            .signer(hello.sig_algos)
            .ok_or(Error::NoSignatureMethods)?;
        let ciphersuite = (self.config.select_ciphersuite)(hello.symmetric_ciphersuites)
            .ok_or(Error::NoCiphersuites)?;
        let kex_algorithm =
            (self.config.select_kex_algorithm)(hello.kex_algos).ok_or(Error::NoKexAlgorithms)?;

        Ok((signer, ciphersuite, kex_algorithm))
    }

    fn start_kex(
        &mut self,
        recv_buf: &mut [u8],
        signer: &SignatureMethod,
        ciphersuite: u8,
        kex_algorithm: u8,
    ) -> Result<VerificationMethod, Error> {
        // Fill the client random buffer
        self.fill_client_random();

        let sig_algorithm = signer.algorithm();
        let key_data = signer.verifying_key_bytes();

        let sig_algorithm_name = crypt::sig_algo_name(sig_algorithm).unwrap_or("unknown");
        let offered_fingerprint =
            crypt::signature::fingerprint_sha256(sig_algorithm_name, &key_data);

        log::info!("Offer {offered_fingerprint}");
        let ciphersuite_name = symmetric::ciphersuite(ciphersuite)
            .map(|t| t.1)
            .unwrap_or("???");
        log::info!("With ciphersuite {ciphersuite_name:?} ({ciphersuite:#x})");
        let client_random = self.client_random;

        self.send(
            None,
            &ClientNegotiationMessage::StartKex {
                kex_algorithm,
                sig_algorithm,
                ciphersuite,
                client_random: &client_random,
                key_data: &key_data,
            },
        )?;

        let (server_sig, server_verifying_key) = match self.recv(None, recv_buf)? {
            ServerNegotiationMessage::StartKex {
                sig_algorithm,
                key_data,
            } => (sig_algorithm, key_data),
            ServerNegotiationMessage::Reject(_) => todo!(),
            _ => return Err(Error::UnexpectedServerReply),
        };

        let server_sig_algorithm_name = crypt::sig_algo_name(server_sig).unwrap_or("unknown");
        let server_fingerprint =
            crypt::signature::fingerprint_sha256(server_sig_algorithm_name, server_verifying_key);
        log::info!("Server fingerprint {server_fingerprint}");

        let verifier = self
            .config
            .signature_keystore
            .accept_public_key(server_sig, server_verifying_key)
            .ok_or(Error::ServerKeyRejected)?;

        Ok(verifier)
    }

    fn key_exchange(
        &mut self,
        recv_buf: &mut [u8],
        signer: &mut SignatureMethod,
        verifier: &mut VerificationMethod,
        ciphersuite: u8,
    ) -> Result<SymmetricCipher, Error> {
        let mut rng = rand::thread_rng();
        let ephemeral = EphemeralSecret::random_from_rng(&mut rng);
        let public = PublicKey::from(&ephemeral);

        self.send(
            Some(signer),
            &ClientNegotiationMessage::DHPublicKey(public.as_bytes()),
        )?;
        let server_public_key = match self.recv(Some(verifier), recv_buf)? {
            ServerNegotiationMessage::DHPublicKey(key_data) => key_data,
            ServerNegotiationMessage::Reject(_reject) => todo!(),
            _ => return Err(Error::UnexpectedServerReply),
        };
        if server_public_key.len() != 32 {
            todo!()
        }

        let (ciphersuite_algo, ciphersuite_name) = symmetric::ciphersuite(ciphersuite).ok_or(
            Error::Symmetric(symmetric::Error::InvalidCiphersuite(ciphersuite)),
        )?;
        let key_size = ciphersuite_algo.key_size();

        log::info!("Using ciphersuite: {ciphersuite_name}");

        let mut server_key = [0; 32];
        server_key.copy_from_slice(server_public_key);
        let server_key = PublicKey::from(server_key);

        let shared = ephemeral.diffie_hellman(&server_key);
        let cipher = SymmetricCipher::from_ciphersuite(
            ciphersuite_algo,
            &shared.as_bytes()[..key_size],
            Iv::from(self.client_random),
        )?;
        Ok(cipher)
    }

    fn finish(
        &mut self,
        recv_buf: &mut [u8],
        signer: &mut SignatureMethod,
        verifier: &mut VerificationMethod,
    ) -> Result<(), Error> {
        self.send(Some(signer), &ClientNegotiationMessage::Agreed)?;
        match self.recv(Some(verifier), recv_buf)? {
            ServerNegotiationMessage::Agreed => Ok(()),
            ServerNegotiationMessage::Reject(_reject) => todo!(),
            _ => Err(Error::UnexpectedServerReply),
        }
    }

    pub fn perform(mut self) -> Result<ClientSocket, Error> {
        let mut recv_buf = [0; 512];

        let (mut signer, ciphersuite, kex_algorithm) = self.hello(&mut recv_buf)?;
        let mut verifier = self.start_kex(&mut recv_buf, &signer, ciphersuite, kex_algorithm)?;
        let symmetric =
            self.key_exchange(&mut recv_buf, &mut signer, &mut verifier, ciphersuite)?;
        self.finish(&mut recv_buf, &mut signer, &mut verifier)?;

        log::info!("Established");
        Ok(ClientSocket::from_parts(
            self.stream,
            self.remote,
            signer,
            verifier,
            symmetric,
        ))
    }

    fn recv<'de>(
        &mut self,
        _verifier: Option<&mut VerificationMethod>,
        recv_buf: &'de mut [u8],
    ) -> Result<ServerNegotiationMessage<'de>, Error> {
        self.stream.read_exact(&mut recv_buf[..size_of::<u16>()])?;
        let len = u16::from_le_bytes([recv_buf[0], recv_buf[1]]) as usize;
        if len >= recv_buf.len() {
            todo!()
        }
        self.stream.read_exact(&mut recv_buf[..len])?;
        let mut decoder = Decoder::new(&recv_buf[..len]);
        let message = ServerNegotiationMessage::decode(&mut decoder)?;
        Ok(message)
    }

    fn send(
        &mut self,
        _signer: Option<&mut SignatureMethod>,
        message: &ClientNegotiationMessage,
    ) -> Result<(), Error> {
        let mut encoder = Encoder::new(&mut self.buffer[2..]);
        message.encode(&mut encoder)?;
        let payload = encoder.get();
        let len: u16 = payload.len().try_into().unwrap();
        self.buffer[0..2].copy_from_slice(&len.to_le_bytes());

        self.stream.write_all(&self.buffer[..len as usize + 2])?;

        Ok(())
    }
}
