use std::env;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use rustls::{ClientConfig as RustlsClientConfig, RootCertStore};
use rustls_pki_types::{CertificateDer, PrivateKeyDer, ServerName};
use tokio::io::{self as tokio_io, AsyncBufReadExt, AsyncRead, AsyncWrite, BufReader, WriteHalf};
use tokio::net::TcpStream;
use tokio::time::{Duration, Instant, timeout};
use tokio_rustls::TlsConnector;

mod pretty;
mod protocol;
mod wire;

use crate::pretty::{parse_end_of_data_info, parse_serial_notify_serial, print_pdu, print_raw_pdu};
use crate::protocol::{PduHeader, PduType, QueryMode};
use crate::wire::{read_pdu, send_reset_query, send_serial_query};

const DEFAULT_READ_TIMEOUT_SECS: u64 = 30;
const DEFAULT_POLL_INTERVAL_SECS: u64 = 600;

trait AsyncStream: AsyncRead + AsyncWrite + Unpin + Send {}
impl<T> AsyncStream for T where T: AsyncRead + AsyncWrite + Unpin + Send {}

type DynStream = Box<dyn AsyncStream>;
type ClientWriter = WriteHalf<DynStream>;

#[tokio::main]
async fn main() -> io::Result<()> {
    let config = Config::from_args()?;

    println!("== RTR debug client ==");
    println!("target  : {}", config.addr);
    println!("transport: {}", config.transport.describe());
    println!("version : {}", config.version);
    println!("timeout : {}s", config.read_timeout_secs);
    println!(
        "poll    : {}s (default before EndOfData refresh is known)",
        config.default_poll_secs
    );
    println!("keep-after-error: {}", config.keep_after_error);
    match &config.mode {
        QueryMode::Reset => {
            println!("mode    : reset");
        }
        QueryMode::Serial { session_id, serial } => {
            println!("mode    : serial");
            println!("session : {}", session_id);
            println!("serial  : {}", serial);
        }
    }
    println!();
    print_help();

    let mut state = ClientState::new(
        config.version,
        config.read_timeout_secs,
        config.default_poll_secs,
        config.keep_after_error,
    );

    let stdin = tokio::io::stdin();
    let mut stdin_lines = BufReader::new(stdin).lines();

    loop {
        let stream = loop {
            match connect_stream(&config).await {
                Ok(stream) => {
                    println!("connected to {}", config.addr);
                    break stream;
                }
                Err(err) => {
                    let delay = state.reconnect_delay_secs();
                    eprintln!("connect failed: {}. retry after {}s", err, delay);
                    tokio::time::sleep(Duration::from_secs(delay)).await;
                }
            }
        };

        let (mut reader, mut writer) = tokio_io::split(stream);
        send_resume_query(&mut writer, &mut state, &config.mode).await?;
        state.schedule_next_poll();
        println!();

        let reconnect = loop {
            let poll_sleep = tokio::time::sleep_until(state.poll_deadline());
            tokio::pin!(poll_sleep);

            tokio::select! {
                line = stdin_lines.next_line() => {
                    match line {
                        Ok(Some(line)) => {
                            match handle_console_command(
                                &line,
                                &mut writer,
                                &mut state,
                            ).await {
                                Ok(should_quit) => {
                                    if should_quit {
                                        println!("quit requested, closing client.");
                                        return Ok(());
                                    }
                                }
                                Err(err) if should_reconnect(&err) => {
                                    eprintln!("command failed due to disconnected transport: {}", err);
                                    break true;
                                }
                                Err(err) => return Err(err),
                            }
                        }
                        Ok(None) => {
                            println!("stdin closed, continue network loop.");
                        }
                        Err(err) => {
                            eprintln!("read stdin failed: {}", err);
                        }
                    }
                }

                _ = &mut poll_sleep => {
                    match handle_poll_tick(&mut writer, &mut state).await {
                        Ok(()) => state.schedule_next_poll(),
                        Err(err) if should_reconnect(&err) => {
                            eprintln!("auto poll failed due to disconnected transport: {}", err);
                            break true;
                        }
                        Err(err) => return Err(err),
                    }
                }

                read_result = timeout(
                    Duration::from_secs(state.read_timeout_secs),
                    read_pdu(&mut reader)
                ) => {
                    match read_result {
                        Ok(Ok(pdu)) => {
                            print_raw_pdu(&pdu.header, &pdu.body);
                            print_pdu(&pdu.header, &pdu.body);
                            match handle_incoming_pdu(&mut writer, &mut state, &pdu.header, &pdu.body).await {
                                Ok(()) => {}
                                Err(err) if should_reconnect(&err) => {
                                    eprintln!("connection dropped while handling incoming PDU: {}", err);
                                    break true;
                                }
                                Err(err) => return Err(err),
                            }
                        }
                        Ok(Err(err)) => {
                            eprintln!("read PDU failed: {}", err);
                            if should_reconnect(&err) {
                                break true;
                            }
                            return Err(err);
                        }
                        Err(_) => {
                            println!(
                                "[timeout] no PDU received in {}s, connection kept open.",
                                state.read_timeout_secs
                            );
                        }
                    }
                }
            }
        };

        if reconnect {
            let delay = state.reconnect_delay_secs();
            state.current_session_id = None;
            println!("[reconnect] transport disconnected, retry after {}s", delay);
            tokio::time::sleep(Duration::from_secs(delay)).await;
        }
    }
}

async fn send_resume_query(
    writer: &mut ClientWriter,
    state: &mut ClientState,
    mode: &QueryMode,
) -> io::Result<()> {
    match (state.session_id, state.serial) {
        (Some(session_id), Some(serial)) => {
            println!(
                "reconnected, send Serial Query with session_id={}, serial={}",
                session_id, serial
            );
            send_serial_query(writer, state.version, session_id, serial).await?;
        }
        _ => match mode {
            QueryMode::Reset => {
                send_reset_query(writer, state.version).await?;
                println!("sent Reset Query");
            }
            QueryMode::Serial { session_id, serial } => {
                state.session_id = Some(*session_id);
                state.serial = Some(*serial);
                send_serial_query(writer, state.version, *session_id, *serial).await?;
                println!("sent Serial Query");
            }
        },
    }

    Ok(())
}

fn should_reconnect(err: &io::Error) -> bool {
    matches!(
        err.kind(),
        io::ErrorKind::UnexpectedEof
            | io::ErrorKind::ConnectionAborted
            | io::ErrorKind::ConnectionReset
            | io::ErrorKind::BrokenPipe
            | io::ErrorKind::NotConnected
    )
}

async fn handle_incoming_pdu(
    writer: &mut ClientWriter,
    state: &mut ClientState,
    header: &PduHeader,
    body: &[u8],
) -> io::Result<()> {
    match header.pdu_type() {
        PduType::CacheResponse => {
            state.current_session_id = Some(header.session_id());
            state.last_error_code = None;
        }

        PduType::Ipv4Prefix | PduType::Ipv6Prefix | PduType::RouterKey | PduType::Aspa => {
            if state.current_session_id.is_none() {
                state.current_session_id = Some(header.session_id());
            }
        }

        PduType::EndOfData => {
            let session_id = header.session_id();
            let eod = parse_end_of_data_info(body);

            state.session_id = Some(session_id);
            state.current_session_id = Some(session_id);

            println!();

            if let Some(eod) = eod {
                state.serial = Some(eod.serial);
                state.refresh = eod.refresh;
                state.retry = eod.retry;
                state.expire = eod.expire;
                state.last_error_code = None;

                println!(
                    "updated client state: session_id={}, serial={}",
                    session_id, eod.serial
                );

                if let Some(refresh) = eod.refresh {
                    println!("refresh    : {}", refresh);
                }
                if let Some(retry) = eod.retry {
                    println!("retry      : {}", retry);
                }
                if let Some(expire) = eod.expire {
                    println!("expire     : {}", expire);
                }

                state.schedule_next_poll();
                println!(
                    "next auto poll scheduled after {}s",
                    state.effective_poll_secs()
                );
            } else {
                println!(
                    "updated client state: session_id={}, serial=<unknown>",
                    session_id
                );
            }

            println!("received EndOfData, keep connection open.");
            println!();
        }

        PduType::SerialNotify => {
            let notify_session_id = header.session_id();
            let notify_serial = parse_serial_notify_serial(body);

            println!();
            println!(
                "[notify] received Serial Notify: session_id={}, notify_serial={:?}",
                notify_session_id, notify_serial
            );

            match (state.session_id, state.serial, notify_serial) {
                (Some(current_session_id), Some(current_serial), Some(_new_serial))
                    if current_session_id == notify_session_id =>
                {
                    println!(
                        "received Serial Notify for current session {}, send Serial Query with serial {}",
                        current_session_id, current_serial
                    );
                    send_serial_query(writer, state.version, current_session_id, current_serial)
                        .await?;
                }

                _ => {
                    println!(
                        "received Serial Notify but local session/serial state is not usable, send Reset Query"
                    );
                    send_reset_query(writer, state.version).await?;
                }
            }

            state.schedule_next_poll();
            println!();
        }

        PduType::CacheReset => {
            println!();
            println!("received Cache Reset, send Reset Query");
            state.current_session_id = None;
            state.serial = None;
            state.last_error_code = None;
            send_reset_query(writer, state.version).await?;
            state.schedule_next_poll();
            println!();
        }

        PduType::ErrorReport => {
            println!();
            println!("received Error Report, pause auto polling for debugging.");
            state.last_error_code = Some(header.error_code());
            if let Some(retry) = state.retry {
                println!("server retry hint currently stored: {}s", retry);
                if state.should_prefer_retry_poll() {
                    println!("when resumed, auto polling will use retry instead of refresh.");
                }
            }
            if state.keep_after_error {
                println!("keep-after-error is enabled, auto polling will continue.");
                state.schedule_next_poll();
            } else {
                println!("use `reset`, `serial`, or `poll resume` to continue manually.");
                state.pause_auto_poll();
            }
            println!();
        }

        PduType::SerialQuery | PduType::ResetQuery | PduType::Unknown(_) => {
            // only print, no extra action
        }
    }

    Ok(())
}

async fn handle_poll_tick(writer: &mut ClientWriter, state: &mut ClientState) -> io::Result<()> {
    println!();
    println!(
        "[auto-poll] timer fired (interval={}s)",
        state.effective_poll_secs()
    );

    match (state.session_id, state.serial) {
        (Some(session_id), Some(serial)) => {
            println!(
                "[auto-poll] send Serial Query with session_id={}, serial={}",
                session_id, serial
            );
            send_serial_query(writer, state.version, session_id, serial).await?;
        }
        _ => {
            println!("[auto-poll] local state incomplete, send Reset Query");
            send_reset_query(writer, state.version).await?;
        }
    }

    println!();
    Ok(())
}

async fn handle_console_command(
    line: &str,
    writer: &mut ClientWriter,
    state: &mut ClientState,
) -> io::Result<bool> {
    let line = line.trim();

    if line.is_empty() {
        return Ok(false);
    }

    let parts: Vec<&str> = line.split_whitespace().collect();

    match parts.as_slice() {
        ["help"] => {
            print_help();
        }

        ["state"] => {
            print_state(state);
        }

        ["reset"] => {
            println!("manual command: send Reset Query");
            send_reset_query(writer, state.version).await?;
            state.schedule_next_poll();
        }

        ["serial"] => match (state.session_id, state.serial) {
            (Some(session_id), Some(serial)) => {
                println!(
                    "manual command: send Serial Query with current state: session_id={}, serial={}",
                    session_id, serial
                );
                send_serial_query(writer, state.version, session_id, serial).await?;
                state.schedule_next_poll();
            }
            _ => {
                println!(
                    "manual command failed: current session_id/serial not available, use `reset` or `serial <session_id> <serial>`"
                );
            }
        },

        ["serial", session_id, serial] => {
            let session_id = match session_id.parse::<u16>() {
                Ok(v) => v,
                Err(err) => {
                    println!("invalid session_id: {}", err);
                    return Ok(false);
                }
            };

            let serial = match serial.parse::<u32>() {
                Ok(v) => v,
                Err(err) => {
                    println!("invalid serial: {}", err);
                    return Ok(false);
                }
            };

            println!(
                "manual command: send Serial Query with explicit args: session_id={}, serial={}",
                session_id, serial
            );
            state.session_id = Some(session_id);
            state.serial = Some(serial);
            send_serial_query(writer, state.version, session_id, serial).await?;
            state.schedule_next_poll();
        }

        ["timeout"] => {
            println!("current read timeout: {}s", state.read_timeout_secs);
        }

        ["timeout", secs] => {
            let secs = match secs.parse::<u64>() {
                Ok(v) if v > 0 => v,
                Ok(_) => {
                    println!("timeout must be > 0");
                    return Ok(false);
                }
                Err(err) => {
                    println!("invalid timeout seconds: {}", err);
                    return Ok(false);
                }
            };

            state.read_timeout_secs = secs;
            println!("updated read timeout to {}s", state.read_timeout_secs);
        }

        ["poll"] => {
            println!(
                "current effective poll interval: {}s",
                state.effective_poll_secs()
            );
            println!(
                "poll interval source       : {}",
                state.poll_interval_source()
            );
            println!("stored refresh hint        : {:?}", state.refresh);
            println!("default poll interval      : {}s", state.default_poll_secs);
            println!("last_error_code            : {:?}", state.last_error_code);
            println!("auto polling paused        : {}", state.poll_paused);
        }

        ["poll", "pause"] => {
            state.pause_auto_poll();
            println!("auto polling paused");
        }

        ["poll", "resume"] => {
            state.resume_auto_poll();
            println!(
                "auto polling resumed, next poll scheduled after {}s",
                state.effective_poll_secs()
            );
        }

        ["poll", secs] => {
            let secs = match secs.parse::<u64>() {
                Ok(v) if v > 0 => v,
                Ok(_) => {
                    println!("poll interval must be > 0");
                    return Ok(false);
                }
                Err(err) => {
                    println!("invalid poll seconds: {}", err);
                    return Ok(false);
                }
            };

            state.refresh = Some(secs as u32);
            state.schedule_next_poll();
            println!("updated poll interval to {}s", secs);
        }

        ["quit"] | ["exit"] => {
            return Ok(true);
        }

        _ => {
            println!("unknown command: {}", line);
            print_help();
        }
    }

    Ok(false)
}

fn print_help() {
    println!("available commands:");
    println!("  help                     show this help");
    println!("  state                    print current client state");
    println!("  reset                    send Reset Query");
    println!("  serial                   send Serial Query with current session_id/serial");
    println!("  serial <sid> <serial>    send Serial Query with explicit values");
    println!("  timeout                  show current read timeout");
    println!("  timeout <secs>           update read timeout seconds");
    println!("  poll                     show current poll interval");
    println!("  poll <secs>              override poll interval seconds");
    println!("  poll pause               pause auto polling");
    println!("  poll resume              resume auto polling");
    println!("  keep-after-error         show current keep-after-error setting");
    println!("  quit                     exit client");
    println!();
}

fn print_state(state: &ClientState) {
    println!("client state:");
    println!("  version            : {}", state.version);
    println!("  session_id         : {:?}", state.session_id);
    println!("  serial             : {:?}", state.serial);
    println!("  current_session_id : {:?}", state.current_session_id);
    println!("  refresh            : {:?}", state.refresh);
    println!("  retry              : {:?}", state.retry);
    println!("  expire             : {:?}", state.expire);
    println!("  read_timeout_secs  : {}", state.read_timeout_secs);
    println!("  default_poll_secs  : {}", state.default_poll_secs);
    println!("  effective_poll_secs: {}", state.effective_poll_secs());
    println!("  poll_source        : {}", state.poll_interval_source());
    println!("  last_error_code    : {:?}", state.last_error_code);
    println!("  keep_after_error   : {}", state.keep_after_error);
    println!("  poll_paused        : {}", state.poll_paused);
    println!();
}

#[derive(Debug)]
struct ClientState {
    version: u8,
    session_id: Option<u16>,
    serial: Option<u32>,
    current_session_id: Option<u16>,

    refresh: Option<u32>,
    retry: Option<u32>,
    expire: Option<u32>,
    last_error_code: Option<u16>,
    keep_after_error: bool,

    read_timeout_secs: u64,
    default_poll_secs: u64,
    next_poll_deadline: Instant,
    poll_paused: bool,
}

impl ClientState {
    fn new(
        version: u8,
        read_timeout_secs: u64,
        default_poll_secs: u64,
        keep_after_error: bool,
    ) -> Self {
        Self {
            version,
            session_id: None,
            serial: None,
            current_session_id: None,
            refresh: None,
            retry: None,
            expire: None,
            last_error_code: None,
            keep_after_error,
            read_timeout_secs,
            default_poll_secs,
            next_poll_deadline: Instant::now() + Duration::from_secs(default_poll_secs),
            poll_paused: false,
        }
    }

    fn effective_poll_secs(&self) -> u64 {
        if self.should_prefer_retry_poll() {
            self.retry.map(|v| v as u64).unwrap_or_else(|| {
                self.refresh
                    .map(|v| v as u64)
                    .unwrap_or(self.default_poll_secs)
            })
        } else {
            self.refresh
                .map(|v| v as u64)
                .unwrap_or(self.default_poll_secs)
        }
    }

    fn schedule_next_poll(&mut self) {
        self.next_poll_deadline = Instant::now() + Duration::from_secs(self.effective_poll_secs());
    }

    fn pause_auto_poll(&mut self) {
        self.poll_paused = true;
    }

    fn resume_auto_poll(&mut self) {
        self.poll_paused = false;
        self.schedule_next_poll();
    }

    fn poll_deadline(&self) -> Instant {
        if self.poll_paused {
            Instant::now() + Duration::from_secs(365 * 24 * 60 * 60)
        } else {
            self.next_poll_deadline
        }
    }

    fn should_prefer_retry_poll(&self) -> bool {
        matches!(self.last_error_code, Some(2 | 10))
    }

    fn poll_interval_source(&self) -> &'static str {
        if self.should_prefer_retry_poll() && self.retry.is_some() {
            "retry"
        } else if self.refresh.is_some() {
            "refresh"
        } else {
            "default"
        }
    }

    fn reconnect_delay_secs(&self) -> u64 {
        if self.should_prefer_retry_poll() {
            self.retry
                .map(|v| v as u64)
                .unwrap_or(self.default_poll_secs)
        } else {
            self.default_poll_secs
        }
    }
}

#[derive(Debug)]
struct Config {
    addr: String,
    version: u8,
    mode: QueryMode,
    read_timeout_secs: u64,
    default_poll_secs: u64,
    transport: TransportConfig,
    keep_after_error: bool,
}

impl Config {
    fn from_args() -> io::Result<Self> {
        let mut args = env::args().skip(1);
        let mut positional = Vec::new();
        let mut transport = TransportConfig::Tcp;
        let mut read_timeout_secs = DEFAULT_READ_TIMEOUT_SECS;
        let mut default_poll_secs = DEFAULT_POLL_INTERVAL_SECS;
        let mut keep_after_error = false;

        while let Some(arg) = args.next() {
            match arg.as_str() {
                "--tls" => {
                    if matches!(transport, TransportConfig::Tcp) {
                        transport = TransportConfig::Tls(TlsConfig::default());
                    }
                }
                "--ca-cert" => {
                    let path = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--ca-cert requires a path")
                    })?;
                    ensure_tls_config(&mut transport)?.ca_cert = Some(PathBuf::from(path));
                }
                "--client-cert" => {
                    let path = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--client-cert requires a path")
                    })?;
                    ensure_tls_config(&mut transport)?.client_cert = Some(PathBuf::from(path));
                }
                "--client-key" => {
                    let path = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--client-key requires a path")
                    })?;
                    ensure_tls_config(&mut transport)?.client_key = Some(PathBuf::from(path));
                }
                "--server-name" => {
                    let name = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--server-name requires a value",
                        )
                    })?;
                    ensure_tls_config(&mut transport)?.server_name = Some(name);
                }
                "--timeout" => {
                    let secs = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--timeout requires seconds")
                    })?;
                    read_timeout_secs = parse_u64_arg(&secs, "--timeout")?;
                }
                "--poll" => {
                    let secs = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--poll requires seconds")
                    })?;
                    default_poll_secs = parse_u64_arg(&secs, "--poll")?;
                }
                "--keep-after-error" => {
                    keep_after_error = true;
                }
                _ if arg.starts_with("--") => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        format!("unknown option '{}'", arg),
                    ));
                }
                _ => positional.push(arg),
            }
        }

        let mut positional = positional.into_iter();

        let addr = positional
            .next()
            .unwrap_or_else(|| "127.0.0.1:323".to_string());

        let version = positional
            .next()
            .map(|s| {
                s.parse::<u8>().map_err(|e| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        format!("invalid version '{}': {}", s, e),
                    )
                })
            })
            .transpose()?
            .unwrap_or(1);

        if version > 2 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("unsupported RTR version {}, expected 0..=2", version),
            ));
        }

        let mode = match positional.next().as_deref() {
            None | Some("reset") => QueryMode::Reset,
            Some("serial") => {
                let session_id = positional
                    .next()
                    .ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "serial mode requires session_id and serial",
                        )
                    })?
                    .parse::<u16>()
                    .map_err(|e| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            format!("invalid session_id: {}", e),
                        )
                    })?;

                let serial = positional
                    .next()
                    .ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "serial mode requires serial")
                    })?
                    .parse::<u32>()
                    .map_err(|e| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            format!("invalid serial: {}", e),
                        )
                    })?;

                QueryMode::Serial { session_id, serial }
            }
            Some(other) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("invalid mode '{}', expected 'reset' or 'serial'", other),
                ));
            }
        };

        let transport = finalize_transport(transport, &addr)?;

        Ok(Self {
            addr,
            version,
            mode,
            read_timeout_secs,
            default_poll_secs,
            transport,
            keep_after_error,
        })
    }
}

#[derive(Debug, Clone)]
enum TransportConfig {
    Tcp,
    Tls(TlsConfig),
}

impl TransportConfig {
    fn describe(&self) -> String {
        match self {
            Self::Tcp => "tcp".to_string(),
            Self::Tls(cfg) => format!(
                "tls (server_name={}, ca_cert={}, client_cert={})",
                cfg.server_name.as_deref().unwrap_or("<unset>"),
                cfg.ca_cert
                    .as_ref()
                    .map(|path| path.display().to_string())
                    .unwrap_or_else(|| "<unset>".to_string()),
                cfg.client_cert
                    .as_ref()
                    .map(|path| path.display().to_string())
                    .unwrap_or_else(|| "<none>".to_string())
            ),
        }
    }
}

#[derive(Debug, Clone, Default)]
struct TlsConfig {
    server_name: Option<String>,
    ca_cert: Option<PathBuf>,
    client_cert: Option<PathBuf>,
    client_key: Option<PathBuf>,
}

fn ensure_tls_config(transport: &mut TransportConfig) -> io::Result<&mut TlsConfig> {
    if matches!(transport, TransportConfig::Tcp) {
        *transport = TransportConfig::Tls(TlsConfig::default());
    }

    match transport {
        TransportConfig::Tls(cfg) => Ok(cfg),
        TransportConfig::Tcp => unreachable!(),
    }
}

fn finalize_transport(transport: TransportConfig, addr: &str) -> io::Result<TransportConfig> {
    match transport {
        TransportConfig::Tcp => Ok(TransportConfig::Tcp),
        TransportConfig::Tls(mut cfg) => {
            let ca_cert = cfg.ca_cert.take().ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "TLS mode requires --ca-cert <path>",
                )
            })?;

            match (&cfg.client_cert, &cfg.client_key) {
                (Some(_), Some(_)) | (None, None) => {}
                _ => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "TLS client authentication requires both --client-cert and --client-key",
                    ));
                }
            }

            let server_name = cfg
                .server_name
                .take()
                .or_else(|| default_server_name_for_addr(addr))
                .ok_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "TLS mode requires --server-name or an address with a parsable host",
                    )
                })?;

            Ok(TransportConfig::Tls(TlsConfig {
                server_name: Some(server_name),
                ca_cert: Some(ca_cert),
                client_cert: cfg.client_cert,
                client_key: cfg.client_key,
            }))
        }
    }
}

async fn connect_stream(config: &Config) -> io::Result<DynStream> {
    match &config.transport {
        TransportConfig::Tcp => Ok(Box::new(TcpStream::connect(&config.addr).await?)),
        TransportConfig::Tls(tls) => connect_tls_stream(&config.addr, tls).await,
    }
}

async fn connect_tls_stream(addr: &str, tls: &TlsConfig) -> io::Result<DynStream> {
    let stream = TcpStream::connect(addr).await?;
    let connector = build_tls_connector(tls)?;
    let server_name_str = tls
        .server_name
        .as_ref()
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing TLS server name"))?;
    let server_name = ServerName::try_from(server_name_str.clone()).map_err(|err| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid TLS server name '{}': {}", server_name_str, err),
        )
    })?;
    let tls_stream = connector
        .connect(server_name, stream)
        .await
        .map_err(|err| {
            io::Error::new(
                io::ErrorKind::ConnectionAborted,
                format!("TLS handshake failed: {}", err),
            )
        })?;
    Ok(Box::new(tls_stream))
}

fn build_tls_connector(tls: &TlsConfig) -> io::Result<TlsConnector> {
    let ca_cert_path = tls
        .ca_cert
        .as_ref()
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing TLS CA cert"))?;
    let ca_certs = load_certs(ca_cert_path)?;
    let mut roots = RootCertStore::empty();
    let (added, _ignored) = roots.add_parsable_certificates(ca_certs);
    if added == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "no valid CA certificates found in {}",
                ca_cert_path.display()
            ),
        ));
    }

    let builder = RustlsClientConfig::builder().with_root_certificates(roots);
    let client_config = match (&tls.client_cert, &tls.client_key) {
        (Some(cert_path), Some(key_path)) => {
            let certs = load_certs(cert_path)?;
            let key = load_private_key(key_path)?;
            builder.with_client_auth_cert(certs, key).map_err(|err| {
                io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("invalid TLS client certificate/key: {}", err),
                )
            })?
        }
        (None, None) => builder.with_no_client_auth(),
        _ => unreachable!(),
    };

    Ok(TlsConnector::from(Arc::new(client_config)))
}

fn load_certs(path: &Path) -> io::Result<Vec<CertificateDer<'static>>> {
    let mut reader = std::io::BufReader::new(std::fs::File::open(path)?);
    let certs = rustls_pemfile::certs(&mut reader)
        .collect::<Result<Vec<_>, _>>()
        .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
    if certs.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("no certificates found in {}", path.display()),
        ));
    }
    Ok(certs)
}

fn load_private_key(path: &Path) -> io::Result<PrivateKeyDer<'static>> {
    let mut reader = std::io::BufReader::new(std::fs::File::open(path)?);
    rustls_pemfile::private_key(&mut reader)
        .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?
        .ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("no private key found in {}", path.display()),
            )
        })
}

fn default_server_name_for_addr(addr: &str) -> Option<String> {
    if let Some(rest) = addr.strip_prefix('[') {
        return rest.split(']').next().map(str::to_string);
    }
    addr.rsplit_once(':').map(|(host, _port)| host.to_string())
}

fn parse_u64_arg(value: &str, name: &str) -> io::Result<u64> {
    let parsed = value.parse::<u64>().map_err(|err| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid value for {}: {}", name, err),
        )
    })?;
    if parsed == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("{} must be > 0", name),
        ));
    }
    Ok(parsed)
}