// Copyright 2018-2026 the Deno authors. MIT license.
use std::io;
use std::num::NonZeroUsize;
use std::result::Result;
use std::sync::Arc;

use anyhow::anyhow;
use futures::FutureExt;
use futures::Stream;
use futures::StreamExt;
use rustls_tokio_stream::TlsStream;
use rustls_tokio_stream::rustls;
use rustls_tokio_stream::rustls::pki_types::CertificateDer;
use rustls_tokio_stream::rustls::pki_types::PrivateKeyDer;
use tokio::net::TcpStream;

use crate::testdata_path;

/// Returns a [`Stream`] of [`TcpStream`]s accepted from the given port.
pub async fn get_tcp_listener_stream(
  name: &'static str,
  port: u16,
) -> impl Stream<Item = Result<TcpStream, std::io::Error>> + Unpin + Send {
  let host_and_port = &format!("localhost:{port}");

  // Listen on ALL addresses that localhost can resolves to.
  let accept = |listener: tokio::net::TcpListener| {
    async {
      let result = listener.accept().await;
      Some((result.map(|r| r.0), listener))
    }
    .boxed()
  };

  let mut addresses = vec![];
  let listeners = tokio::net::lookup_host(host_and_port)
    .await
    .expect(host_and_port)
    .inspect(|address| addresses.push(*address))
    .map(tokio::net::TcpListener::bind)
    .collect::<futures::stream::FuturesUnordered<_>>()
    .collect::<Vec<_>>()
    .await
    .into_iter()
    .map(|s| s.unwrap())
    .map(|listener| futures::stream::unfold(listener, accept))
    .collect::<Vec<_>>();

  // Eye catcher for HttpServerCount
  test_util::println!("ready: {name} on {:?}", addresses);

  futures::stream::select_all(listeners)
}

pub const TLS_BUFFER_SIZE: Option<NonZeroUsize> = NonZeroUsize::new(65536);

#[derive(Default)]
pub enum SupportedHttpVersions {
  #[default]
  All,
  Http1Only,
  Http2Only,
}

pub fn get_tls_listener_stream_from_tcp(
  tls_config: Arc<rustls::ServerConfig>,
  mut tcp: impl Stream<Item = Result<TcpStream, std::io::Error>> + Unpin + 'static,
) -> impl Stream<Item = Result<TlsStream<TcpStream>, std::io::Error>> + Unpin {
  async_stream::stream! {
    while let Some(result) = tcp.next().await {
      match result {
        Ok(tcp) => yield Ok(TlsStream::new_server_side(tcp, tls_config.clone(), TLS_BUFFER_SIZE)),
        Err(e) => yield Err(e),
      };
    }
  }.boxed_local()
}

pub async fn get_tls_listener_stream(
  name: &'static str,
  port: u16,
  http: SupportedHttpVersions,
) -> impl Stream<Item = Result<TlsStream<TcpStream>, std::io::Error>> + Unpin {
  let cert_file = "tls/localhost.crt";
  let key_file = "tls/localhost.key";
  let ca_cert_file = "tls/RootCA.pem";
  let tls_config =
    get_tls_config(cert_file, key_file, ca_cert_file, http).unwrap();

  let tcp = get_tcp_listener_stream(name, port).await;
  get_tls_listener_stream_from_tcp(tls_config, tcp)
}

pub fn get_tls_config(
  cert: &str,
  key: &str,
  ca: &str,
  http_versions: SupportedHttpVersions,
) -> io::Result<Arc<rustls::ServerConfig>> {
  let cert_path = testdata_path().join(cert);
  let key_path = testdata_path().join(key);
  let ca_path = testdata_path().join(ca);

  let cert_file = std::fs::File::open(cert_path)?;
  let key_file = std::fs::File::open(key_path)?;
  let ca_file = std::fs::File::open(ca_path)?;

  let certs_result: Result<Vec<CertificateDer<'static>>, io::Error> = {
    let mut cert_reader = io::BufReader::new(cert_file);
    rustls_pemfile::certs(&mut cert_reader).collect()
  };
  let certs = certs_result?;

  let mut ca_cert_reader = io::BufReader::new(ca_file);
  let ca_cert = rustls_pemfile::certs(&mut ca_cert_reader)
    .collect::<Vec<_>>()
    .remove(0)?;

  let mut key_reader = io::BufReader::new(key_file);
  let key = {
    let pkcs8_keys = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
      .collect::<Result<Vec<_>, _>>()?;
    let rsa_keys = rustls_pemfile::rsa_private_keys(&mut key_reader)
      .collect::<Result<Vec<_>, _>>()?;

    if !pkcs8_keys.is_empty() {
      let key = pkcs8_keys[0].clone_key();
      Some(PrivateKeyDer::from(key))
    } else if !rsa_keys.is_empty() {
      let key = rsa_keys[0].clone_key();
      Some(PrivateKeyDer::from(key))
    } else {
      None
    }
  };

  match key {
    Some(key) => {
      let mut root_cert_store = rustls::RootCertStore::empty();
      root_cert_store.add(ca_cert).unwrap();

      // Allow (but do not require) client authentication.
      let client_verifier = rustls::server::WebPkiClientVerifier::builder(
        Arc::new(root_cert_store),
      )
      .allow_unauthenticated()
      .build()
      .unwrap();

      let mut config = rustls::ServerConfig::builder()
        .with_client_cert_verifier(client_verifier)
        .with_single_cert(certs, key)
        .map_err(|e| anyhow!("Error setting cert: {:?}", e))
        .unwrap();

      match http_versions {
        SupportedHttpVersions::All => {
          config.alpn_protocols = vec!["h2".into(), "http/1.1".into()];
        }
        SupportedHttpVersions::Http1Only => {}
        SupportedHttpVersions::Http2Only => {
          config.alpn_protocols = vec!["h2".into()];
        }
      }

      Ok(Arc::new(config))
    }
    None => Err(io::Error::other("Cannot find key")),
  }
}

/// Like `get_tls_listener_stream`, but the server requires a valid client
/// certificate signed by the test CA
pub async fn get_tls_listener_stream_with_required_client_auth(
  name: &'static str,
  port: u16,
) -> impl Stream<Item = Result<TlsStream<TcpStream>, std::io::Error>> + Unpin {
  let cert_file = "tls/localhost.crt";
  let key_file = "tls/localhost.key";
  let ca_cert_file = "tls/RootCA.pem";
  let tls_config =
    get_tls_config_with_required_client_auth(cert_file, key_file, ca_cert_file)
      .unwrap();

  let tcp = get_tcp_listener_stream(name, port).await;
  get_tls_listener_stream_from_tcp(tls_config, tcp)
}

fn get_tls_config_with_required_client_auth(
  cert: &str,
  key: &str,
  ca: &str,
) -> io::Result<Arc<rustls::ServerConfig>> {
  let cert_path = testdata_path().join(cert);
  let key_path = testdata_path().join(key);
  let ca_path = testdata_path().join(ca);

  let cert_file = std::fs::File::open(cert_path)?;
  let key_file = std::fs::File::open(key_path)?;
  let ca_file = std::fs::File::open(ca_path)?;

  let certs: Vec<CertificateDer<'static>> = {
    let mut cert_reader = io::BufReader::new(cert_file);
    rustls_pemfile::certs(&mut cert_reader).collect::<Result<_, _>>()?
  };

  let mut ca_cert_reader = io::BufReader::new(ca_file);
  let ca_cert = rustls_pemfile::certs(&mut ca_cert_reader)
    .collect::<Vec<_>>()
    .remove(0)?;

  let key = {
    let mut key_reader = io::BufReader::new(key_file);
    let pkcs8_keys = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
      .collect::<Result<Vec<_>, _>>()?;
    if !pkcs8_keys.is_empty() {
      PrivateKeyDer::from(pkcs8_keys[0].clone_key())
    } else {
      return Err(io::Error::other("Cannot find key"));
    }
  };

  let mut root_cert_store = rustls::RootCertStore::empty();
  root_cert_store.add(ca_cert).unwrap();

  let client_verifier =
    rustls::server::WebPkiClientVerifier::builder(Arc::new(root_cert_store))
      .build()
      .unwrap();

  let mut config = rustls::ServerConfig::builder()
    .with_client_cert_verifier(client_verifier)
    .with_single_cert(certs, key)
    .map_err(|e| io::Error::other(format!("Error setting cert: {:?}", e)))?;

  config.alpn_protocols = vec!["h2".into(), "http/1.1".into()];

  Ok(Arc::new(config))
}
