use crate::proxy::DomainFrontingConfig;
use crate::{
    DnsResolver,
    abortable_stream::{AbortableStream, AbortableStreamHandle},
    proxy::{ApiConnection, ApiConnectionMode, ProxyConfig},
    tls_stream::TlsStream,
};
use domain_fronting::ProxyConnection;
use futures::{StreamExt, channel::mpsc, future, pin_mut};
#[cfg(target_os = "android")]
use futures::{channel::oneshot, sink::SinkExt};
use http::uri::Scheme;
use hyper::Uri;
use hyper_util::rt::TokioIo;
use mullvad_encrypted_dns_proxy::{
    Forwarder as EncryptedDNSForwarder, config::ProxyConfig as EncryptedDNSConfig,
};
use shadowsocks::{
    ServerConfig,
    config::{ServerConfigError, ServerType},
    context::{Context as SsContext, SharedContext},
    crypto::CipherKind,
    relay::tcprelay::ProxyClientStream,
};
#[cfg(target_os = "android")]
use std::os::unix::io::{AsRawFd, RawFd};
use std::{
    fmt,
    future::Future,
    io,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    str,
    sync::{Arc, Mutex},
    task::{Context, Poll},
    time::Duration,
};
use talpid_types::net::proxy;
use tokio::{
    io::{AsyncRead, AsyncWrite},
    net::{TcpSocket, TcpStream},
    sync::Notify,
    time::timeout,
};
use tower::Service;

#[cfg(any(feature = "api-override", test))]
use crate::proxy::ConnectionDecorator;

const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);

/// Communicate asynchronously with an [HttpsConnector] actor.
/// Created via [HttpsConnector::spawn].
#[derive(Clone)]
pub struct HttpsConnectorHandle {
    tx: mpsc::UnboundedSender<HttpsConnectorRequest>,
}

impl HttpsConnectorHandle {
    /// Stop all streams produced by this connector
    pub fn reset(&self) {
        let _ = self.tx.unbounded_send(HttpsConnectorRequest::Reset);
    }

    /// Change the proxy settings for the connector
    pub fn set_connection_mode(&self, proxy: ApiConnectionMode) {
        let _ = self
            .tx
            .unbounded_send(HttpsConnectorRequest::SetConnectionMode(proxy));
    }
}

/// Requests for the HttpsConnector actor.
enum HttpsConnectorRequest {
    Reset,
    SetConnectionMode(ApiConnectionMode),
}

#[derive(Clone)]
pub(crate) enum InnerConnectionMode {
    /// Connect directly to the target.
    Direct,
    /// Connect to the destination via a Shadowsocks proxy.
    Shadowsocks(ShadowsocksConfig),
    /// Connect to the destination via a Socks proxy.
    Socks5(SocksConfig),
    /// Connect to the destination via Mullvad Encrypted DNS proxy.
    /// See [`mullvad_encrypted_dns_proxy`] for how the proxy works.
    EncryptedDnsProxy(EncryptedDNSConfig),
    /// Connect to the destination via domain fronting.
    DomainFronting(DomainFrontingConfig),
}

impl InnerConnectionMode {
    async fn connect(
        self,
        hostname: &str,
        addr: &SocketAddr,
        #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
        #[cfg(any(feature = "api-override", test))] disable_tls: bool,
    ) -> Result<ApiConnection, std::io::Error> {
        match self {
            // Set up a TCP-socket connection.
            InnerConnectionMode::Direct => {
                let first_hop = *addr;
                let make_proxy_stream = |tcp_stream| async { Ok(tcp_stream) };
                Self::connect_proxied(
                    first_hop,
                    hostname,
                    make_proxy_stream,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx,
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                )
                .await
            }
            // Set up a Shadowsocks-connection.
            InnerConnectionMode::Shadowsocks(shadowsocks) => {
                let first_hop = shadowsocks.params.peer;
                let make_proxy_stream = |tcp_stream| async {
                    Ok(ProxyClientStream::from_stream(
                        shadowsocks.proxy_context,
                        tcp_stream,
                        &ServerConfig::try_from(shadowsocks.params)
                            .map_err(|_| std::io::Error::other("Invalid shadowsocks config"))?,
                        *addr,
                    ))
                };
                Self::connect_proxied(
                    first_hop,
                    hostname,
                    make_proxy_stream,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx,
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                )
                .await
            }
            // Set up a SOCKS5-connection.
            InnerConnectionMode::Socks5(socks) => {
                let first_hop = socks.peer;
                let make_proxy_stream = |tcp_stream| async {
                    match socks.authentication {
                        None => {
                            tokio_socks::tcp::Socks5Stream::connect_with_socket(tcp_stream, addr)
                                .await
                        }
                        Some(credentials) => {
                            tokio_socks::tcp::Socks5Stream::connect_with_password_and_socket(
                                tcp_stream,
                                addr,
                                credentials.username(),
                                credentials.password(),
                            )
                            .await
                        }
                    }
                    .map_err(|error| io::Error::other(format!("SOCKS error: {error}")))
                };
                Self::connect_proxied(
                    first_hop,
                    hostname,
                    make_proxy_stream,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx,
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                )
                .await
            }
            InnerConnectionMode::EncryptedDnsProxy(proxy_config) => {
                let first_hop = SocketAddr::V4(proxy_config.addr);
                let make_proxy_stream = |tcp_stream| async {
                    let forwarder = EncryptedDNSForwarder::from_stream(&proxy_config, tcp_stream);
                    Ok(forwarder)
                };
                Self::connect_proxied(
                    first_hop,
                    hostname,
                    make_proxy_stream,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx,
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                )
                .await
            }
            InnerConnectionMode::DomainFronting(config) => {
                let front = config.domain_fronting.front().to_owned();
                let proxy_host = config.domain_fronting.proxy_host().to_owned();
                let session_header_key = config.domain_fronting.session_header_key().to_owned();
                let make_proxy_stream = |tcp_stream| async {
                    let tls_stream = cdn_tls_connect(tcp_stream, &front).await?;
                    let forwarder =
                        ProxyConnection::from_stream(tls_stream, proxy_host, session_header_key)
                            .await
                            .map_err(std::io::Error::other)?;
                    Ok(forwarder)
                };
                Self::connect_proxied(
                    config.addr,
                    hostname,
                    make_proxy_stream,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx,
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                )
                .await
            }
        }
    }

    /// Create an [`ApiConnection`] from a [`TcpStream`].
    ///
    /// The `make_proxy_stream` closure receives a [`TcpStream`] and produces a
    /// stream which can send to and receive data from some server using any
    /// proxy protocol. The only restriction is that this stream must implement
    /// [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`], as well as
    /// [`Unpin`] and [`Send`].
    ///
    /// If a direct connection is to be established (i.e. the stream will not be
    /// using any proxy protocol) `make_proxy_stream` may return the
    /// [`TcpStream`] itself. See for example how a connection is established
    /// from connection mode [`InnerConnectionMode::Direct`].
    async fn connect_proxied<ProxyFactory, ProxyFuture, Proxy>(
        first_hop: SocketAddr,
        hostname: &str,
        make_proxy_stream: ProxyFactory,
        #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
        #[cfg(any(feature = "api-override", test))] disable_tls: bool,
    ) -> Result<ApiConnection, io::Error>
    where
        ProxyFactory: FnOnce(TcpStream) -> ProxyFuture,
        ProxyFuture: Future<Output = io::Result<Proxy>>,
        Proxy: AsyncRead + AsyncWrite + Unpin + Send + 'static,
    {
        let socket = HttpsConnector::open_socket(
            first_hop,
            #[cfg(target_os = "android")]
            socket_bypass_tx,
        )
        .await?;

        let proxy = make_proxy_stream(socket).await?;

        #[cfg(any(feature = "api-override", test))]
        if disable_tls {
            return Ok(ApiConnection::new(Box::new(ConnectionDecorator(proxy))));
        }

        let tls_stream = TlsStream::connect_https(proxy, hostname).await?;
        Ok(ApiConnection::new(Box::new(tls_stream)))
    }
}

/// Establish a TLS connection to the CDN without certificate verification.
///
/// Certificate verification is intentionally skipped because the security of domain fronting
/// does not depend on the CDN TLS — the actual API connection is secured by the inner TLS
/// layer to `api.mullvad.net`.
async fn cdn_tls_connect(
    stream: TcpStream,
    front_domain: &str,
) -> io::Result<tokio_rustls::client::TlsStream<TcpStream>> {
    use std::sync::LazyLock;
    use tokio_rustls::rustls::{self, client::danger, pki_types};

    #[derive(Debug)]
    struct NoCertVerification;

    impl danger::ServerCertVerifier for NoCertVerification {
        fn verify_server_cert(
            &self,
            _: &pki_types::CertificateDer<'_>,
            _: &[pki_types::CertificateDer<'_>],
            _: &pki_types::ServerName<'_>,
            _: &[u8],
            _: pki_types::UnixTime,
        ) -> Result<danger::ServerCertVerified, rustls::Error> {
            Ok(danger::ServerCertVerified::assertion())
        }

        fn verify_tls12_signature(
            &self,
            _: &[u8],
            _: &pki_types::CertificateDer<'_>,
            _: &rustls::DigitallySignedStruct,
        ) -> Result<danger::HandshakeSignatureValid, rustls::Error> {
            Ok(danger::HandshakeSignatureValid::assertion())
        }

        fn verify_tls13_signature(
            &self,
            _: &[u8],
            _: &pki_types::CertificateDer<'_>,
            _: &rustls::DigitallySignedStruct,
        ) -> Result<danger::HandshakeSignatureValid, rustls::Error> {
            Ok(danger::HandshakeSignatureValid::assertion())
        }

        fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
            rustls::crypto::ring::default_provider()
                .signature_verification_algorithms
                .supported_schemes()
        }
    }

    static CDN_TLS_CONFIG: LazyLock<Arc<rustls::ClientConfig>> = LazyLock::new(|| {
        Arc::new(
            rustls::ClientConfig::builder()
                .dangerous()
                .with_custom_certificate_verifier(Arc::new(NoCertVerification))
                .with_no_client_auth(),
        )
    });

    let connector = tokio_rustls::TlsConnector::from(Arc::clone(&CDN_TLS_CONFIG));
    let server_name = pki_types::ServerName::try_from(front_domain.to_owned())
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid front domain"))?;
    connector.connect(server_name, stream).await
}

#[derive(Clone)]
pub(crate) struct ShadowsocksConfig {
    proxy_context: SharedContext,
    params: ParsedShadowsocksConfig,
}

#[derive(Clone)]
struct ParsedShadowsocksConfig {
    peer: SocketAddr,
    password: String,
    cipher: CipherKind,
}

impl TryFrom<ParsedShadowsocksConfig> for ServerConfig {
    type Error = ServerConfigError;

    fn try_from(config: ParsedShadowsocksConfig) -> Result<Self, Self::Error> {
        ServerConfig::new(config.peer, config.password, config.cipher)
    }
}

#[derive(Clone)]
pub(crate) struct SocksConfig {
    peer: SocketAddr,
    authentication: Option<proxy::SocksAuth>,
}

impl From<ApiConnectionMode> for InnerConnectionMode {
    fn from(config: ApiConnectionMode) -> Self {
        use std::net::Ipv4Addr;
        match config {
            ApiConnectionMode::Direct => InnerConnectionMode::Direct,
            ApiConnectionMode::Proxied(proxy_settings) => match proxy_settings {
                ProxyConfig::Shadowsocks(config) => {
                    InnerConnectionMode::Shadowsocks(ShadowsocksConfig {
                        params: ParsedShadowsocksConfig {
                            peer: config.endpoint,
                            password: config.plaintext_password().to_string(),
                            cipher: config.cipher.kind(),
                        },
                        proxy_context: SsContext::new_shared(ServerType::Local),
                    })
                }
                ProxyConfig::Socks5Local(config) => InnerConnectionMode::Socks5(SocksConfig {
                    peer: SocketAddr::new(IpAddr::from(Ipv4Addr::LOCALHOST), config.local_port),
                    authentication: None,
                }),
                ProxyConfig::Socks5Remote(config) => InnerConnectionMode::Socks5(SocksConfig {
                    peer: config.endpoint,
                    authentication: config.auth,
                }),
                ProxyConfig::EncryptedDnsProxy(config) => {
                    InnerConnectionMode::EncryptedDnsProxy(config)
                }
                ProxyConfig::DomainFronting(config) => InnerConnectionMode::DomainFronting(config),
            },
        }
    }
}

/// A Connector for the `https` scheme.
#[derive(Clone)]
pub struct HttpsConnector {
    inner: Arc<Mutex<HttpsConnectorInner>>,
    abort_notify: Arc<Notify>,
    dns_resolver: Arc<dyn DnsResolver>,
    #[cfg(target_os = "android")]
    socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
    #[cfg(any(feature = "api-override", test))]
    disable_tls: bool,
}

struct HttpsConnectorInner {
    stream_handles: Vec<AbortableStreamHandle>,
    proxy_config: InnerConnectionMode,
}

#[cfg(target_os = "android")]
pub type SocketBypassRequest = (RawFd, oneshot::Sender<()>);

impl HttpsConnector {
    /// Creates a new [HttpsConnector] actor instance, which notably implements [`Service<Uri>`].
    ///
    /// Call [HttpsConnector::spawn] to start listening for [events](HttpsConnectorRequest).
    pub fn new(
        dns_resolver: Arc<dyn DnsResolver>,
        proxy_config: InnerConnectionMode,
        #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
        #[cfg(any(feature = "api-override", test))] disable_tls: bool,
    ) -> Self {
        let abort_notify = Arc::new(tokio::sync::Notify::new());
        let connector = Arc::new(Mutex::new(HttpsConnectorInner {
            proxy_config,
            stream_handles: Default::default(),
        }));

        HttpsConnector {
            inner: connector,
            abort_notify,
            dns_resolver,
            #[cfg(target_os = "android")]
            socket_bypass_tx,
            #[cfg(any(feature = "api-override", test))]
            disable_tls,
        }
    }

    /// Start listening for incoming [HttpsConnectorRequest]s from the created [HttpsConnectorHandle].
    pub fn spawn(&self) -> HttpsConnectorHandle {
        let (tx, requests) = mpsc::unbounded();
        let connector = self.inner.clone();
        let notify = self.abort_notify.clone();
        let request_handler = RequestHandler {
            connector,
            notify,
            requests,
        };
        tokio::spawn(request_handler.run());
        HttpsConnectorHandle { tx }
    }

    /// Establishes a TCP connection with a peer at the specified socket address.
    ///
    /// Will timeout after [`CONNECT_TIMEOUT`] seconds.
    async fn open_socket(
        addr: SocketAddr,
        #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
    ) -> std::io::Result<TcpStream> {
        let socket = match addr {
            SocketAddr::V4(_) => TcpSocket::new_v4()?,
            SocketAddr::V6(_) => TcpSocket::new_v6()?,
        };

        #[cfg(target_os = "android")]
        if let Some(mut tx) = socket_bypass_tx {
            let (done_tx, done_rx) = oneshot::channel();
            let _ = tx.send((socket.as_raw_fd(), done_tx)).await;
            if done_rx.await.is_err() {
                log::error!("Failed to bypass socket, connection might fail");
            }
        }

        timeout(CONNECT_TIMEOUT, socket.connect(addr))
            .await
            .map_err(|err| io::Error::new(io::ErrorKind::TimedOut, err))?
    }

    /// Resolve the provided `uri` to an IP and port. If the URI contains an IP, that IP will be
    /// used. Otherwise `dns_resolver` will be used as a fallback.
    /// If the URI contains a port, then that port will be used.
    async fn resolve_address(dns_resolver: &dyn DnsResolver, uri: Uri) -> io::Result<SocketAddr> {
        const DEFAULT_PORT: u16 = 443;

        let hostname = uri.host().ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidInput, "invalid url, missing host")
        })?;
        let port = uri.port_u16();
        if let Ok(addr) = hostname.parse::<IpAddr>() {
            return Ok(SocketAddr::new(addr, port.unwrap_or(DEFAULT_PORT)));
        }

        let addrs = dns_resolver.resolve(hostname.to_owned()).await?;
        let addr = addrs
            .first()
            .ok_or_else(|| io::Error::other("Empty DNS response"))?;
        let port = match (addr.port(), port) {
            (_, Some(port)) => port,
            (0, None) => DEFAULT_PORT,
            (addr_port, None) => addr_port,
        };
        Ok(SocketAddr::new(addr.ip(), port))
    }
}

impl fmt::Debug for HttpsConnector {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HttpsConnector").finish()
    }
}

impl Service<Uri> for HttpsConnector {
    type Response = TokioIo<AbortableStream<ApiConnection>>;
    type Error = io::Error;
    type Future =
        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let mut inner = self.inner.lock().unwrap();
        inner.stream_handles.retain(|handle| !handle.is_closed());
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, uri: Uri) -> Self::Future {
        let inner = self.inner.clone();
        let abort_notify = self.abort_notify.clone();
        #[cfg(target_os = "android")]
        let socket_bypass_tx = self.socket_bypass_tx.clone();
        let dns_resolver = self.dns_resolver.clone();

        #[cfg(any(feature = "api-override", test))]
        let disable_tls = self.disable_tls;

        let fut = async move {
            if uri.scheme() != Some(&Scheme::HTTPS) {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "invalid url, not https",
                ));
            }
            let Some(hostname) = uri.host().map(str::to_owned) else {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "invalid url, missing host",
                ));
            };
            let addr = Self::resolve_address(&*dns_resolver, uri).await?;

            // Loop until we have established a connection. This starts over if a new endpoint
            // is selected while connecting.
            let stream = loop {
                let notify = abort_notify.notified();
                let proxy_config = { inner.lock().unwrap().proxy_config.clone() };
                let stream_fut = proxy_config.connect(
                    &hostname,
                    &addr,
                    #[cfg(target_os = "android")]
                    socket_bypass_tx.clone(),
                    #[cfg(any(feature = "api-override", test))]
                    disable_tls,
                );

                pin_mut!(stream_fut);
                pin_mut!(notify);

                // Wait for connection. Abort and retry if we switched to a different server.
                if let future::Either::Left((stream, _)) = future::select(stream_fut, notify).await
                {
                    break stream?;
                }
            };

            let (stream, socket_handle) = AbortableStream::new(stream);

            {
                let mut inner = inner.lock().unwrap();
                inner.stream_handles.push(socket_handle);
            }

            Ok(TokioIo::new(stream))
        };

        Box::pin(fut)
    }
}

struct RequestHandler {
    connector: Arc<Mutex<HttpsConnectorInner>>,
    notify: Arc<Notify>,
    requests: mpsc::UnboundedReceiver<HttpsConnectorRequest>,
}

impl RequestHandler {
    /// Process requests from an associated [HttpsConnectorHandle].
    async fn run(mut self) {
        while let Some(request) = self.requests.next().await {
            self.handle(request);
        }
    }

    /// Process a single request from a connected [HttpsConnectorHandle].
    fn handle(&mut self, request: HttpsConnectorRequest) {
        let handles = {
            let mut inner = self.connector.lock().unwrap();
            match request {
                HttpsConnectorRequest::Reset => std::mem::take(&mut inner.stream_handles),
                HttpsConnectorRequest::SetConnectionMode(config) => {
                    inner.proxy_config = InnerConnectionMode::from(config);
                    std::mem::take(&mut inner.stream_handles)
                }
            }
        };
        for handle in handles {
            handle.close();
        }
        self.notify.notify_waiters();
    }
}
