//! Provides a TLS 1.3 stream, accepting only LE for root cert.
//! SNI is disabled.
use std::{
    io::{self, ErrorKind},
    pin::Pin,
    sync::Arc,
    task::{self, Poll},
};

use hyper_util::client::legacy::connect::{Connected, Connection};
use rustls_pki_types::{CertificateDer, pem::PemObject};
use std::sync::LazyLock;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio_rustls::{
    TlsConnector,
    rustls::{self, ClientConfig, pki_types::ServerName},
};

const LE_ROOT_CERT: &[u8] = include_bytes!("../le_root_cert.pem");

pub struct TlsStream<S: AsyncRead + AsyncWrite + Unpin> {
    stream: tokio_rustls::client::TlsStream<S>,
}

static TLS_CONFIG: LazyLock<Arc<ClientConfig>> = LazyLock::new(|| {
    let config = {
        let mut config =
            ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
                .with_protocol_versions(&[&rustls::version::TLS13])
                .expect("ring crypt-prover should support TLS 1.3")
                .with_root_certificates(read_cert_store().expect("Failed to parse pem file"))
                .with_no_client_auth();
        // This assumes that the server hello/certificates will include certificate for the domain.
        config.enable_sni = false;
        config
    };
    Arc::new(config)
});

impl<S> TlsStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    pub async fn connect_https(stream: S, domain: &str) -> io::Result<TlsStream<S>> {
        Self::connect_https_with_client_config(stream, domain, Arc::clone(&TLS_CONFIG)).await
    }

    pub async fn connect_https_with_client_config(
        stream: S,
        domain: &str,
        client_config: Arc<ClientConfig>,
    ) -> io::Result<TlsStream<S>> {
        let connector = TlsConnector::from(client_config);

        let host = match ServerName::try_from(domain.to_owned()) {
            Ok(n) => n,
            Err(_) => {
                return Err(io::Error::new(
                    ErrorKind::InvalidInput,
                    format!("invalid hostname \"{domain}\""),
                ));
            }
        };

        let stream = connector.connect(host, stream).await?;

        Ok(TlsStream { stream })
    }
}

fn read_cert_store() -> Result<rustls::RootCertStore, rustls_pki_types::pem::Error> {
    let mut cert_store = rustls::RootCertStore::empty();

    let certs = CertificateDer::pem_reader_iter(&mut std::io::BufReader::new(LE_ROOT_CERT))
        .collect::<Result<Vec<_>, _>>()?;
    let (num_certs_added, num_failures) = cert_store.add_parsable_certificates(certs);
    if num_failures > 0 || num_certs_added != 1 {
        panic!("Failed to add root cert");
    }

    Ok(cert_store)
}

impl<S> AsyncRead for TlsStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.stream).poll_read(cx, buf)
    }
}

impl<S> AsyncWrite for TlsStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.stream).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.stream).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.stream).poll_shutdown(cx)
    }
}

impl<S> Connection for TlsStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn connected(&self) -> Connected {
        Connected::new()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_cert_loading() {
        let _certs = read_cert_store();
    }
}
