use std::sync::Arc;

use anyhow::Context;
use tokio::io::{self, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

use crate::relay;

pub async fn run(destination: &str, server: &str, no_tls: bool) -> anyhow::Result<()> {
    let server_addr = if server.contains(':') {
        server.to_string()
    } else {
        format!("{server}:443")
    };

    let tcp_stream = TcpStream::connect(&server_addr)
        .await
        .with_context(|| format!("failed to connect to relay at {server_addr}"))?;
    tcp_stream.set_nodelay(true)?;
    relay::tune_socket(&tcp_stream);

    if no_tls {
        handshake_and_relay(tcp_stream, destination).await
    } else {
        let host = server_addr.rsplit_once(':').unwrap().0;
        let tls_stream = tls_connect(tcp_stream, host).await?;
        handshake_and_relay(tls_stream, destination).await
    }
}

async fn handshake_and_relay<S: AsyncRead + AsyncWrite + Unpin>(
    mut stream: S,
    destination: &str,
) -> anyhow::Result<()> {
    relay::write_line(&mut stream, &format!("V2 {destination}")).await?;

    let response = relay::read_line(&mut stream)
        .await
        .context("failed to read server response")?;

    if response != "OK" {
        eprintln!("layover: {response}");
        std::process::exit(1);
    }

    // Bridge stdin/stdout <-> relay server.
    // Use join! (not select!) so both directions drain fully before exiting.
    // BufReader(64KB) matches relay buffer size — avoids 8x more syscalls from
    // io::copy's default 8KB buffer.
    let (stream_read, mut stream_write) = io::split(stream);
    let mut stream_read = BufReader::with_capacity(65536, stream_read);
    let stdin = io::stdin();
    let mut stdin = BufReader::with_capacity(65536, stdin);
    let mut stdout = io::stdout();

    let up = async {
        let r = io::copy_buf(&mut stdin, &mut stream_write).await;
        let _ = stream_write.shutdown().await;
        r
    };
    let down = io::copy_buf(&mut stream_read, &mut stdout);

    let (up_result, down_result) = tokio::join!(up, down);
    let _ = up_result;
    let _ = down_result;

    Ok(())
}

async fn tls_connect(
    tcp_stream: TcpStream,
    host: &str,
) -> anyhow::Result<tokio_rustls::client::TlsStream<TcpStream>> {
    use tokio_rustls::rustls::pki_types::ServerName;
    use tokio_rustls::rustls::{ClientConfig, RootCertStore};

    let mut root_store = RootCertStore::empty();
    root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());

    let config = ClientConfig::builder()
        .with_root_certificates(root_store)
        .with_no_client_auth();

    let connector = TlsConnector::from(Arc::new(config));
    let server_name = ServerName::try_from(host.to_string())
        .map_err(|e| anyhow::anyhow!("invalid server hostname: {e}"))?;

    connector
        .connect(server_name, tcp_stream)
        .await
        .context("TLS handshake with relay server failed")
}
