use std::collections::BTreeSet;
use std::env;
use std::io;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Instant;

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

use rpki_rs::rtr::client::{Client, PayloadError, PayloadTarget};
use rpki_rs::rtr::payload::{Action, Payload, Timing};

const DEFAULT_STEPS: usize = 1;
const DEFAULT_STEP_TIMEOUT_SECS: u64 = 300;
const DEFAULT_PROGRESS_EVERY: u64 = 10_000;

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

type DynStream = Box<dyn AsyncStream>;

#[derive(Debug, Clone)]
struct Config {
    addr: String,
    steps: usize,
    follow: bool,
    transport: TransportConfig,
    assert_substr: Vec<String>,
    assert_min_records: Option<usize>,
    print_records: bool,
    step_timeout_secs: u64,
    progress_every: u64,
}

impl Config {
    fn from_args() -> io::Result<Self> {
        let mut args = env::args().skip(1);
        let mut positional = Vec::new();

        let mut steps = DEFAULT_STEPS;
        let mut follow = false;
        let mut transport = TransportConfig::Tcp;
        let mut assert_substr = Vec::new();
        let mut assert_min_records = None;
        let mut print_records = false;
        let mut step_timeout_secs = DEFAULT_STEP_TIMEOUT_SECS;
        let mut progress_every = DEFAULT_PROGRESS_EVERY;

        while let Some(arg) = args.next() {
            match arg.as_str() {
                "-h" | "--help" => {
                    print_usage();
                    std::process::exit(0);
                }

                "--steps" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--steps requires value")
                    })?;
                    steps = parse_usize_arg(&v, "--steps")?;
                    if steps == 0 {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--steps must be >= 1",
                        ));
                    }
                }

                "--follow" => {
                    follow = true;
                }

                "--tls" => {
                    ensure_tls(&mut transport)?;
                }

                "--ca-cert" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--ca-cert requires path")
                    })?;
                    ensure_tls(&mut transport)?.ca_cert = Some(PathBuf::from(v));
                }

                "--server-name" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--server-name requires value",
                        )
                    })?;
                    ensure_tls(&mut transport)?.server_name = Some(v);
                }

                "--client-cert" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--client-cert requires path",
                        )
                    })?;
                    ensure_tls(&mut transport)?.client_cert = Some(PathBuf::from(v));
                }

                "--client-key" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--client-key requires path",
                        )
                    })?;
                    ensure_tls(&mut transport)?.client_key = Some(PathBuf::from(v));
                }

                "--assert-substr" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--assert-substr requires value",
                        )
                    })?;
                    assert_substr.push(v);
                }

                "--assert-min-records" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--assert-min-records requires value",
                        )
                    })?;
                    assert_min_records = Some(parse_usize_arg(&v, "--assert-min-records")?);
                }

                "--print-records" => {
                    print_records = true;
                }

                "--step-timeout-secs" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--step-timeout-secs requires value",
                        )
                    })?;
                    step_timeout_secs = parse_u64_arg(&v, "--step-timeout-secs")?;
                    if step_timeout_secs == 0 {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--step-timeout-secs must be >= 1",
                        ));
                    }
                }

                "--progress-every" => {
                    let v = args.next().ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--progress-every requires value",
                        )
                    })?;
                    progress_every = parse_u64_arg(&v, "--progress-every")?;
                    if progress_every == 0 {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "--progress-every must be >= 1",
                        ));
                    }
                }

                // 明确拒绝当前 wrapper 不支持的能力
                "--version" => {
                    let _ = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--version requires value")
                    })?;
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "--version is not exposed by this rpki client wrapper",
                    ));
                }

                "--timeout" => {
                    let _ = args.next().ok_or_else(|| {
                        io::Error::new(io::ErrorKind::InvalidInput, "--timeout requires value")
                    })?;
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "--timeout is not exposed by this rpki client wrapper; use --step-timeout-secs instead",
                    ));
                }

                _ 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());

        if let Some(extra) = positional.next() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "unexpected positional argument '{}'; only optional [addr] is supported",
                    extra
                ),
            ));
        }

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

        Ok(Self {
            addr,
            steps,
            follow,
            transport,
            assert_substr,
            assert_min_records,
            print_records,
            step_timeout_secs,
            progress_every,
        })
    }
}

fn print_usage() {
    eprintln!(
        "\
rpki-rs-test-client

Usage:
  rpki-rs-test-client [OPTIONS] [ADDR]

Examples:
  rpki-rs-test-client
  rpki-rs-test-client 127.0.0.1:323 --steps 1
  rpki-rs-test-client 127.0.0.1:323 --steps 1 --step-timeout-secs 600
  rpki-rs-test-client 127.0.0.1:323 --follow
  rpki-rs-test-client 127.0.0.1:323 --assert-min-records 1 --assert-substr 192.0.2.
  rpki-rs-test-client 127.0.0.1:3324 --tls --ca-cert certs/ca.pem --server-name localhost

Options:
  --steps <N>                 Number of client.step() calls to perform (default: 1)
  --follow                    Keep calling step() forever
  --tls                       Enable TLS
  --ca-cert <PATH>            CA certificate PEM file (required in TLS mode)
  --server-name <NAME>        TLS server name; required when ADDR host is an IP
  --client-cert <PATH>        Client certificate PEM file (optional, with --client-key)
  --client-key <PATH>         Client private key PEM file (optional, with --client-cert)
  --assert-substr <TEXT>      Assert final stable record dump contains substring
  --assert-min-records <N>    Assert final record count >= N
  --print-records             Print records after each successful step
  --step-timeout-secs <N>     Timeout for each step() call in seconds (default: 300)
  --progress-every <N>        Print apply progress every N updates (default: 10000)
  -h, --help                  Show this help

Not supported by this wrapper:
  --version
  --timeout
  explicit serial bootstrap via session_id/serial
"
    );
}

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

#[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(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 => Err(io::Error::other("tls configuration unavailable")),
    }
}

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 auth requires both --client-cert and --client-key",
                    ));
                }
            }

            let server_name = match cfg.server_name.take() {
                Some(name) => name,
                None => {
                    let host = parse_host_from_addr(addr).ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "failed to parse host from address",
                        )
                    })?;

                    if host.parse::<IpAddr>().is_ok() {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "TLS with IP address requires explicit --server-name",
                        ));
                    }

                    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,
            }))
        }
    }
}

#[derive(Debug, Default, Clone)]
struct SharedTarget {
    inner: Arc<Mutex<TargetState>>,
}

#[derive(Debug)]
struct TargetState {
    records: BTreeSet<Payload>,
    timing: Option<Timing>,
    announced_seen: u64,
    withdrawn_seen: u64,
    updates_applied_total: u64,
    progress_every: u64,
    apply_batches: u64,
    last_apply_started_at: Option<Instant>,
}

impl Default for TargetState {
    fn default() -> Self {
        Self {
            records: BTreeSet::new(),
            timing: None,
            announced_seen: 0,
            withdrawn_seen: 0,
            updates_applied_total: 0,
            progress_every: DEFAULT_PROGRESS_EVERY,
            apply_batches: 0,
            last_apply_started_at: None,
        }
    }
}

#[derive(Debug, Clone)]
struct TargetSnapshot {
    records: Vec<Payload>,
    timing: Option<Timing>,
    announced_seen: u64,
    withdrawn_seen: u64,
    updates_applied_total: u64,
    apply_batches: u64,
}

impl SharedTarget {
    fn new(progress_every: u64) -> Self {
        let state = TargetState {
            progress_every,
            ..TargetState::default()
        };

        Self {
            inner: Arc::new(Mutex::new(state)),
        }
    }

    fn snapshot(&self) -> TargetSnapshot {
        let guard = self.inner.lock().expect("target mutex poisoned");
        TargetSnapshot {
            records: guard.records.iter().cloned().collect(),
            timing: guard.timing,
            announced_seen: guard.announced_seen,
            withdrawn_seen: guard.withdrawn_seen,
            updates_applied_total: guard.updates_applied_total,
            apply_batches: guard.apply_batches,
        }
    }

    fn payload_to_stable_text(payload: &Payload) -> String {
        format!("{:?}", payload)
    }
}

impl TargetSnapshot {
    fn dump_text(&self) -> String {
        self.records
            .iter()
            .map(SharedTarget::payload_to_stable_text)
            .collect::<Vec<_>>()
            .join("\n")
    }
}

impl PayloadTarget for SharedTarget {
    type Update = Vec<(Action, Payload)>;

    fn start(&mut self, reset: bool) -> Self::Update {
        if reset {
            let mut guard = self.inner.lock().expect("target mutex poisoned");
            println!(
                "[target] start reset=true | clearing existing records={}",
                guard.records.len()
            );
            guard.records.clear();
        } else {
            let guard = self.inner.lock().expect("target mutex poisoned");
            println!(
                "[target] start reset=false | current records={}",
                guard.records.len()
            );
        }

        Vec::new()
    }

    fn apply(&mut self, update: Self::Update, timing: Timing) -> Result<(), PayloadError> {
        let total = update.len() as u64;

        {
            let mut guard = self.inner.lock().expect("target mutex poisoned");
            guard.apply_batches += 1;
            guard.last_apply_started_at = Some(Instant::now());
            println!(
                "[target] apply batch #{} started | updates={} | current records={}",
                guard.apply_batches,
                total,
                guard.records.len()
            );
        }

        let started = Instant::now();
        let mut local_processed: u64 = 0;

        let progress_every = {
            let guard = self.inner.lock().expect("target mutex poisoned");
            guard.progress_every
        };

        let mut guard = self.inner.lock().expect("target mutex poisoned");

        for (action, payload) in update {
            match action {
                Action::Announce => {
                    guard.announced_seen += 1;
                    if !guard.records.insert(payload) {
                        return Err(PayloadError::DuplicateAnnounce);
                    }
                }
                Action::Withdraw => {
                    guard.withdrawn_seen += 1;
                    if !guard.records.remove(&payload) {
                        return Err(PayloadError::UnknownWithdraw);
                    }
                }
            }

            local_processed += 1;
            guard.updates_applied_total += 1;

            if local_processed % progress_every == 0 || local_processed == total {
                println!(
                    "[target] apply progress | batch_processed={}/{} | total_updates_seen={} | records={} | elapsed={:.2?}",
                    local_processed,
                    total,
                    guard.updates_applied_total,
                    guard.records.len(),
                    started.elapsed(),
                );
            }
        }

        guard.timing = Some(timing);

        println!(
            "[target] apply batch complete | updates={} | records={} | announced_seen={} | withdrawn_seen={} | elapsed={:.2?}",
            total,
            guard.records.len(),
            guard.announced_seen,
            guard.withdrawn_seen,
            started.elapsed(),
        );

        Ok(())
    }
}

fn print_step_summary(
    step_no: usize,
    before: &TargetSnapshot,
    after: &TargetSnapshot,
    print_records: bool,
) {
    println!(
        "[step] {} ok | records: {} -> {} | delta announce={} withdraw={} | delta updates={} | apply_batches={}",
        step_no,
        before.records.len(),
        after.records.len(),
        after.announced_seen.saturating_sub(before.announced_seen),
        after.withdrawn_seen.saturating_sub(before.withdrawn_seen),
        after
            .updates_applied_total
            .saturating_sub(before.updates_applied_total),
        after.apply_batches,
    );

    if let Some(timing) = after.timing {
        println!(
            "[step] {} timing | refresh={} retry={} expire={}",
            step_no, timing.refresh, timing.retry, timing.expire
        );
    }

    if print_records {
        println!("-- records after step {} --", step_no);
        if after.records.is_empty() {
            println!("(empty)");
        } else {
            for rec in &after.records {
                println!("{}", SharedTarget::payload_to_stable_text(rec));
            }
        }
    }
}

fn run_assertions(config: &Config, snapshot: &TargetSnapshot) -> io::Result<()> {
    if let Some(min) = config.assert_min_records
        && snapshot.records.len() < min
    {
        return Err(io::Error::other(format!(
            "assertion failed: records {} < {}",
            snapshot.records.len(),
            min
        )));
    }

    if !config.assert_substr.is_empty() {
        let dump = snapshot.dump_text();
        for needle in &config.assert_substr {
            if !dump.contains(needle) {
                return Err(io::Error::other(format!(
                    "assertion failed: missing substring '{}'",
                    needle
                )));
            }
        }
    }

    Ok(())
}

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

    println!("== rpki-rs-test-client ==");
    println!("target             : {}", config.addr);
    println!("steps              : {}", config.steps);
    println!("follow             : {}", config.follow);
    println!("step_timeout_secs  : {}", config.step_timeout_secs);
    println!("progress_every     : {}", config.progress_every);

    match &config.transport {
        TransportConfig::Tcp => {
            println!("transport          : tcp");
        }
        TransportConfig::Tls(tls) => {
            println!("transport          : tls");
            println!(
                "server_name        : {}",
                tls.server_name.as_deref().unwrap_or("<unset>")
            );
            println!(
                "ca_cert            : {}",
                tls.ca_cert
                    .as_ref()
                    .map(|p| p.display().to_string())
                    .unwrap_or_else(|| "<unset>".to_string())
            );
        }
    }

    let stream = connect_stream(&config).await?;
    let target = SharedTarget::new(config.progress_every);
    let inspect = target.clone();

    let mut client = Client::new(stream, target, None);

    for idx in 0..config.steps {
        let step_no = idx + 1;
        let before = inspect.snapshot();

        println!("[step] {} begin", step_no);
        let step_started = Instant::now();

        timeout(Duration::from_secs(config.step_timeout_secs), client.step())
            .await
            .map_err(|_| {
                io::Error::new(
                    io::ErrorKind::TimedOut,
                    format!(
                        "step {} timed out after {}s",
                        step_no, config.step_timeout_secs
                    ),
                )
            })?
            .map_err(|err| io::Error::new(err.kind(), format!("step {} failed: {}", step_no, err)))?;

        let after = inspect.snapshot();
        print_step_summary(step_no, &before, &after, config.print_records);
        println!(
            "[step] {} finished in {:.2?}",
            step_no,
            step_started.elapsed()
        );
    }

    if config.follow {
        let mut step_index = config.steps;
        println!("[follow] enabled, entering continuous step loop");

        loop {
            step_index += 1;
            let before = inspect.snapshot();

            println!("[step] {} begin", step_index);
            let step_started = Instant::now();

            timeout(Duration::from_secs(config.step_timeout_secs), client.step())
                .await
                .map_err(|_| {
                    io::Error::new(
                        io::ErrorKind::TimedOut,
                        format!(
                            "step {} timed out after {}s",
                            step_index, config.step_timeout_secs
                        ),
                    )
                })?
                .map_err(|err| {
                    io::Error::new(err.kind(), format!("step {} failed: {}", step_index, err))
                })?;

            let after = inspect.snapshot();
            print_step_summary(step_index, &before, &after, config.print_records);
            println!(
                "[step] {} finished in {:.2?}",
                step_index,
                step_started.elapsed()
            );
        }
    }

    let state = client.state();
    let final_snapshot = inspect.snapshot();

    println!("state              : {:?}", state);
    if let Some(timing) = final_snapshot.timing {
        println!(
            "timing             : refresh={} retry={} expire={}",
            timing.refresh, timing.retry, timing.expire
        );
    }
    println!("records            : {}", final_snapshot.records.len());
    println!(
        "updates_seen       : announce={} withdraw={}",
        final_snapshot.announced_seen, final_snapshot.withdrawn_seen
    );
    println!(
        "updates_applied    : {}",
        final_snapshot.updates_applied_total
    );
    println!("apply_batches      : {}", final_snapshot.apply_batches);

    run_assertions(&config, &final_snapshot)?;
    println!("[assert] passed");

    Ok(())
}

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 certs in {}", ca_cert_path.display()),
        ));
    }

    let builder = RustlsClientConfig::builder().with_root_certificates(roots);

    let cfg = 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 client cert/key: {}", err),
                )
            })?
        }
        (None, None) => builder.with_no_client_auth(),
        _ => unreachable!(),
    };

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

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 certs 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 in {}", path.display()),
            )
        })
}

fn parse_host_from_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, _)| host.to_string())
}

fn parse_usize_arg(value: &str, name: &str) -> io::Result<usize> {
    value.parse::<usize>().map_err(|err| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid {} '{}': {}", name, value, err),
        )
    })
}

fn parse_u64_arg(value: &str, name: &str) -> io::Result<u64> {
    value.parse::<u64>().map_err(|err| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid {} '{}': {}", name, value, err),
        )
    })
}