use arc_swap::ArcSwap;
use std::fs;
use std::io::BufReader;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use rustls::{ClientConfig, RootCertStore};
use rustls_pki_types::{CertificateDer, PrivateKeyDer, ServerName};
use tokio::io::AsyncBufReadExt;
use tokio::net::TcpStream;
use tokio::time::{Instant, sleep};
use tokio_rustls::TlsConnector;

use rpki::rtr::cache::{RtrCacheBuilder, SessionIds, SharedRtrCache};
use rpki::rtr::payload::Timing;
use rpki::rtr::pdu::{CacheResponse, EndOfDataV1, ResetQuery};
use rpki::rtr::server::ssh::SshAuthMode;
use rpki::rtr::server::RtrService;
use russh::client;
use russh::keys;
use russh::keys::ssh_key::LineEnding;

fn fixture_path(name: &str) -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join("tls")
        .join(name)
}

fn load_pem_certs(path: &Path) -> Vec<CertificateDer<'static>> {
    let file = fs::File::open(path).expect("open cert file");
    let mut reader = BufReader::new(file);
    rustls_pemfile::certs(&mut reader)
        .collect::<Result<Vec<_>, _>>()
        .expect("parse certs")
}

fn load_pem_key(path: &Path) -> PrivateKeyDer<'static> {
    let file = fs::File::open(path).expect("open key file");
    let mut reader = BufReader::new(file);
    rustls_pemfile::private_key(&mut reader)
        .expect("read private key")
        .expect("missing private key")
}

fn test_cache() -> SharedRtrCache {
    Arc::new(ArcSwap::from_pointee(
        RtrCacheBuilder::new()
            .session_ids(SessionIds::from_array([42, 42, 42]))
            .serials([100, 100, 100])
            .timing(Timing::new(600, 600, 7200))
            .build(),
    ))
}

fn reserve_local_addr() -> SocketAddr {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind temp listener");
    listener.local_addr().expect("local addr")
}

async fn wait_for_port(addr: SocketAddr) {
    let deadline = Instant::now() + Duration::from_secs(2);
    loop {
        if TcpStream::connect(addr).await.is_ok() {
            return;
        }
        assert!(
            Instant::now() < deadline,
            "port {} did not open in time",
            addr
        );
        sleep(Duration::from_millis(20)).await;
    }
}

async fn connect_tls_client(addr: SocketAddr) -> tokio_rustls::client::TlsStream<TcpStream> {
    let mut roots = RootCertStore::empty();
    for cert in load_pem_certs(&fixture_path("client-ca.crt")) {
        roots.add(cert).expect("add root cert");
    }

    let certs = load_pem_certs(&fixture_path("client-good.crt"));
    let key = load_pem_key(&fixture_path("client-good.key"));
    let cfg = ClientConfig::builder()
        .with_root_certificates(roots)
        .with_client_auth_cert(certs, key)
        .expect("build tls client auth");
    let connector = TlsConnector::from(Arc::new(cfg));

    let tcp = TcpStream::connect(addr).await.expect("connect tls tcp");
    connector
        .connect(ServerName::IpAddress(addr.ip().into()), tcp)
        .await
        .expect("tls connect")
}

struct TestSshClientHandler;

impl client::Handler for TestSshClientHandler {
    type Error = anyhow::Error;

    async fn check_server_key(
        &mut self,
        _server_public_key: &russh::keys::ssh_key::PublicKey,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

#[tokio::test]
async fn unified_server_tcp_handles_reset_query() {
    let service = RtrService::new(test_cache());
    let tcp_addr = reserve_local_addr();

    let running = service.spawn_tcp_only(tcp_addr);
    wait_for_port(tcp_addr).await;

    let mut client = TcpStream::connect(tcp_addr).await.expect("connect tcp");
    ResetQuery::new(1)
        .write(&mut client)
        .await
        .expect("send reset");

    let response = CacheResponse::read(&mut client)
        .await
        .expect("read cache response");
    assert_eq!(response.version(), 1);
    assert_eq!(response.session_id(), 42);

    let eod = EndOfDataV1::read(&mut client).await.expect("read eod");
    assert_eq!(eod.version(), 1);
    assert_eq!(eod.session_id(), 42);
    assert_eq!(eod.serial_number(), 100);

    running.shutdown();
    running.wait().await;
}

#[tokio::test]
async fn unified_server_tls_handles_reset_query() {
    let service = RtrService::new(test_cache());
    let tcp_addr = reserve_local_addr();
    let tls_addr = reserve_local_addr();

    let running = service.spawn_tcp_and_tls_from_pem(
        tcp_addr,
        tls_addr,
        fixture_path("server.crt"),
        fixture_path("server.key"),
        fixture_path("client-ca.crt"),
    );
    wait_for_port(tls_addr).await;

    let mut client = connect_tls_client(tls_addr).await;
    ResetQuery::new(1)
        .write(&mut client)
        .await
        .expect("send reset tls");

    let response = CacheResponse::read(&mut client)
        .await
        .expect("read tls cache response");
    assert_eq!(response.version(), 1);
    assert_eq!(response.session_id(), 42);

    let eod = EndOfDataV1::read(&mut client).await.expect("read tls eod");
    assert_eq!(eod.version(), 1);
    assert_eq!(eod.session_id(), 42);
    assert_eq!(eod.serial_number(), 100);

    running.shutdown();
    running.wait().await;
}

#[tokio::test]
async fn unified_server_ssh_opens_listener_and_emits_banner() {
    let service = RtrService::new(test_cache());
    let tcp_addr = reserve_local_addr();
    let ssh_addr = reserve_local_addr();

    let tmp = tempfile::tempdir().expect("tempdir");
    let host_key_path = tmp.path().join("ssh_host_ed25519_key");
    let authorized_keys_path = tmp.path().join("authorized_keys");

    let host_key =
        keys::PrivateKey::random(&mut rand::rng(), keys::Algorithm::Ed25519).expect("gen host key");
    let host_key_pem = host_key
        .to_openssh(LineEnding::LF)
        .expect("encode host key");
    fs::write(&host_key_path, host_key_pem).expect("write host key");

    let pubkey_line = host_key.public_key().to_openssh().expect("encode pubkey");
    fs::write(&authorized_keys_path, format!("{pubkey_line}\n")).expect("write authorized_keys");

    let running = service.spawn_tcp_and_ssh_from_openssh(
        tcp_addr,
        ssh_addr,
        &host_key_path,
        &authorized_keys_path,
        "rpki-rtr",
        "rpki-rtr",
        SshAuthMode::Key,
        None,
    );
    wait_for_port(ssh_addr).await;

    let stream = TcpStream::connect(ssh_addr).await.expect("connect ssh");
    let mut reader = tokio::io::BufReader::new(stream);
    let mut banner = String::new();
    reader
        .read_line(&mut banner)
        .await
        .expect("read ssh banner");
    assert!(
        banner.starts_with("SSH-2.0-"),
        "unexpected ssh banner: {}",
        banner
    );

    running.shutdown();
    running.wait().await;
}

#[tokio::test]
async fn unified_server_ssh_accepts_password_when_configured() {
    let service = RtrService::new(test_cache());
    let tcp_addr = reserve_local_addr();
    let ssh_addr = reserve_local_addr();

    let tmp = tempfile::tempdir().expect("tempdir");
    let host_key_path = tmp.path().join("ssh_host_ed25519_key");
    let authorized_keys_path = tmp.path().join("authorized_keys");

    let host_key =
        keys::PrivateKey::random(&mut rand::rng(), keys::Algorithm::Ed25519).expect("gen host key");
    let host_key_pem = host_key
        .to_openssh(LineEnding::LF)
        .expect("encode host key");
    fs::write(&host_key_path, host_key_pem).expect("write host key");

    let pubkey_line = host_key.public_key().to_openssh().expect("encode pubkey");
    fs::write(&authorized_keys_path, format!("{pubkey_line}\n")).expect("write authorized_keys");

    let running = service.spawn_tcp_and_ssh_from_openssh(
        tcp_addr,
        ssh_addr,
        &host_key_path,
        &authorized_keys_path,
        "rpki-rtr",
        "rpki-rtr",
        SshAuthMode::Both,
        Some("test-password"),
    );
    wait_for_port(ssh_addr).await;

    let mut session = client::connect(
        Arc::new(client::Config::default()),
        ssh_addr,
        TestSshClientHandler,
    )
    .await
    .expect("connect ssh client");
    let auth_result = session
        .authenticate_password("rpki-rtr", "test-password")
        .await
        .expect("password auth result");
    assert!(auth_result.success(), "password auth should succeed");

    let channel = session
        .channel_open_session()
        .await
        .expect("open session channel");
    channel
        .request_subsystem(true, "rpki-rtr")
        .await
        .expect("request subsystem");
    let mut stream = channel.into_stream();

    ResetQuery::new(1)
        .write(&mut stream)
        .await
        .expect("send reset over ssh subsystem");
    let response = CacheResponse::read(&mut stream)
        .await
        .expect("read cache response over ssh subsystem");
    assert_eq!(response.version(), 1);
    assert_eq!(response.session_id(), 42);
    let eod = EndOfDataV1::read(&mut stream)
        .await
        .expect("read eod over ssh subsystem");
    assert_eq!(eod.version(), 1);
    assert_eq!(eod.session_id(), 42);
    assert_eq!(eod.serial_number(), 100);

    running.shutdown();
    running.wait().await;
}

#[tokio::test]
async fn unified_server_ssh_rejects_password_when_not_configured() {
    let service = RtrService::new(test_cache());
    let tcp_addr = reserve_local_addr();
    let ssh_addr = reserve_local_addr();

    let tmp = tempfile::tempdir().expect("tempdir");
    let host_key_path = tmp.path().join("ssh_host_ed25519_key");
    let authorized_keys_path = tmp.path().join("authorized_keys");

    let host_key =
        keys::PrivateKey::random(&mut rand::rng(), keys::Algorithm::Ed25519).expect("gen host key");
    let host_key_pem = host_key
        .to_openssh(LineEnding::LF)
        .expect("encode host key");
    fs::write(&host_key_path, host_key_pem).expect("write host key");

    let pubkey_line = host_key.public_key().to_openssh().expect("encode pubkey");
    fs::write(&authorized_keys_path, format!("{pubkey_line}\n")).expect("write authorized_keys");

    let running = service.spawn_tcp_and_ssh_from_openssh(
        tcp_addr,
        ssh_addr,
        &host_key_path,
        &authorized_keys_path,
        "rpki-rtr",
        "rpki-rtr",
        SshAuthMode::Key,
        None,
    );
    wait_for_port(ssh_addr).await;

    let mut session = client::connect(
        Arc::new(client::Config::default()),
        ssh_addr,
        TestSshClientHandler,
    )
    .await
    .expect("connect ssh client");
    let auth_result = session
        .authenticate_password("rpki-rtr", "test-password")
        .await
        .expect("password auth result");
    assert!(!auth_result.success(), "password auth should be rejected");

    running.shutdown();
    running.wait().await;
}
