datum-agent 0.10.3

use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::{Duration, Instant},
};

use datum::{Channel, Keep, Sink, Source};
use datum_agent::{
    Agent, AgentConfig, JobMat, JobSpec,
    dcp::{
        ClientKind, DcpClient, DcpJobFactories, DcpQuicServerConfig, DcpServer, DcpServerConfig,
        DcpTcpServerConfig, Hello, ResponseStatus, proto::DcpFrame,
    },
};
use datum_net::quic::{
    crypto::rustls::{QuicClientConfig, QuicServerConfig},
    quinn,
    rustls::{
        ClientConfig as RustlsClientConfig, RootCertStore, ServerConfig as RustlsServerConfig,
        pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer},
        server::WebPkiClientVerifier,
    },
};
use prost::Message as ProstMessage;
use rcgen::{CertifiedKey, generate_simple_self_signed};
use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
    net::{TcpListener, TcpStream},
};

fn agent() -> datum_agent::AgentHandle {
    Agent::start_with_config(AgentConfig {
        poll_interval: Duration::from_millis(5),
        event_buffer: 128,
        ..AgentConfig::default()
    })
    .expect("agent starts")
}

fn never_job(instance_name: String) -> JobSpec {
    JobSpec::new(instance_name, |context| {
        let control = context.control();
        Ok(
            Source::tick(Duration::ZERO, Duration::from_millis(10), 1_u64)
                .instrumented(
                    format!("{}:{}", context.name(), context.generation()),
                    context.instrumentation_registry(),
                )
                .via_mat(context.drain_flow(), Keep::right)
                .to_mat(Sink::ignore(), move |_switch, completion| {
                    JobMat::new(completion, control.clone())
                }),
        )
    })
}

fn never_factories() -> DcpJobFactories {
    let factories = DcpJobFactories::new();
    factories
        .register("never", |instance_name, _params| {
            Ok(never_job(instance_name))
        })
        .expect("factory registers");
    factories
}

async fn start_tcp_server(
    factories: DcpJobFactories,
    metrics_interval: Duration,
) -> (
    datum_agent::AgentHandle,
    datum_agent::dcp::DcpServerHandle,
    SocketAddr,
) {
    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        factories,
        DcpServerConfig {
            node_id: "test-agent".to_owned(),
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
            }),
            quic: None,
            auth_token: None,
            metrics_interval,
            frame_buffer: 128,
        },
    );
    let handle = server.start().await.expect("DCP server starts");
    let addr = handle.tcp_addr().expect("TCP address");
    (agent, handle, addr)
}

async fn connect_tcp(addr: SocketAddr) -> DcpClient {
    DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("DCP client connects")
}

async fn wait_status(
    client: &DcpClient,
    name: &str,
    predicate: impl Fn(&datum_agent::dcp::proto::JobStatus) -> bool,
) -> datum_agent::dcp::proto::JobStatus {
    wait_until(Duration::from_secs(3), || async {
        let status = client.job_status(name).await.expect("job status");
        predicate(&status).then_some(status)
    })
    .await
}

async fn wait_until<T, Fut>(timeout: Duration, mut condition: impl FnMut() -> Fut) -> T
where
    Fut: std::future::Future<Output = Option<T>>,
{
    let deadline = Instant::now() + timeout;
    loop {
        if let Some(value) = condition().await {
            return value;
        }
        assert!(Instant::now() < deadline, "condition timed out");
        tokio::time::sleep(Duration::from_millis(5)).await;
    }
}

#[tokio::test]
async fn hello_version_negotiation_accepts_and_rejects_unknown_major() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;

    let mut bad_hello = Hello::new("bad-client", ClientKind::Cli);
    bad_hello.protocol_version = "1.0.0".to_owned();
    let rejected = match DcpClient::connect_tcp(addr, bad_hello).await {
        Ok(_) => panic!("unknown major is rejected"),
        Err(error) => error,
    };
    assert!(
        matches!(
            rejected,
            datum_agent::dcp::DcpError::Response {
                status: ResponseStatus::ProtocolMismatch,
                ..
            }
        ),
        "unexpected rejection: {rejected:?}"
    );

    let client = connect_tcp(addr).await;
    assert!(client.list_jobs().await.expect("list jobs").is_empty());

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn commands_cover_start_list_status_drain_stop_and_restart() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;

    let started = client
        .start_job("never", "cmd-job", HashMap::new())
        .await
        .expect("job starts");
    assert_eq!(started.name, "cmd-job");
    assert_eq!(started.state, "Running");

    let listed = client.list_jobs().await.expect("jobs list");
    assert_eq!(listed.len(), 1);
    assert_eq!(listed[0].name, "cmd-job");

    let status = client.job_status("cmd-job").await.expect("status");
    assert_eq!(status.state, "Running");

    let restarted = client.restart_job("cmd-job").await.expect("restart");
    assert_eq!(restarted.state, "Running");
    assert_eq!(restarted.generation, 2);

    client.drain_job("cmd-job").await.expect("drain starts");
    let drained = wait_status(&client, "cmd-job", |status| status.state == "Drained").await;
    assert_eq!(drained.state, "Drained");

    let stopped_job = client
        .start_job("never", "stop-job", HashMap::new())
        .await
        .expect("second job starts");
    assert_eq!(stopped_job.state, "Running");
    let stopped = client.stop_job("stop-job").await.expect("stop");
    assert_eq!(stopped.state, "Stopped");

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn subscribe_events_preserves_lifecycle_ordering() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;
    let mut events = client.subscribe_events().await.expect("events subscribe");

    client
        .start_job("never", "event-job", HashMap::new())
        .await
        .expect("job starts");

    let submitted = tokio::time::timeout(Duration::from_secs(2), events.recv())
        .await
        .expect("submitted timeout")
        .expect("submitted event");
    let started = tokio::time::timeout(Duration::from_secs(2), events.recv())
        .await
        .expect("started timeout")
        .expect("started event");

    assert_eq!(submitted.kind, "Submitted");
    assert_eq!(started.kind, "Started");
    assert_eq!(submitted.name, "event-job");
    assert_eq!(started.name, "event-job");
    assert!(submitted.sequence < started.sequence);

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn subscribe_metrics_delivers_samples_while_job_runs() {
    let channel = Channel::bounded(16);
    let producer = channel.clone();
    let factories = DcpJobFactories::new();
    factories
        .register("metered", move |instance_name, _params| {
            let channel = channel.clone();
            Ok(JobSpec::new(instance_name, move |context| {
                let control = context.control();
                Ok(channel
                    .source()
                    .instrumented(
                        format!("{}:{}", context.name(), context.generation()),
                        context.instrumentation_registry(),
                    )
                    .via_mat(context.drain_flow(), Keep::right)
                    .to_mat(Sink::ignore(), move |_switch, completion| {
                        JobMat::new(completion, control.clone())
                    }))
            }))
        })
        .expect("factory registers");

    let (agent, handle, addr) = start_tcp_server(factories, Duration::from_millis(10)).await;
    let client = connect_tcp(addr).await;
    let mut metrics = client
        .subscribe_metrics(10)
        .await
        .expect("metrics subscribe");

    client
        .start_job("metered", "metrics-job", HashMap::new())
        .await
        .expect("job starts");
    for value in 0_u64..4 {
        producer.try_send(value).expect("metric item sends");
    }

    let deadline = Instant::now() + Duration::from_secs(3);
    let sample =
        loop {
            let sample = tokio::time::timeout(Duration::from_millis(200), metrics.recv())
                .await
                .ok()
                .flatten()
                .expect("metrics stream remains open");
            if sample.streams.iter().any(|metric| {
                metric.name.starts_with("metrics-job:") && metric.elements_through >= 4
            }) {
                break sample;
            }
            assert!(Instant::now() < deadline, "timed out waiting for metrics");
        };
    assert!(!sample.streams.is_empty());

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn loopback_tcp_mode_works_and_non_loopback_plaintext_is_refused() {
    let (agent_handle, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;
    assert!(client.list_jobs().await.expect("list").is_empty());
    handle.shutdown().await;
    agent_handle
        .registry()
        .shutdown()
        .expect("registry shuts down");

    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        never_factories(),
        DcpServerConfig {
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
            }),
            quic: None,
            ..DcpServerConfig::default()
        },
    );
    assert!(server.start().await.is_err());
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn malformed_frame_closes_connection_without_panicking() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let mut stream = TcpStream::connect(addr).await.expect("raw client connects");
    stream
        .write_all(&0_u32.to_be_bytes())
        .await
        .expect("malformed frame writes");
    let mut header = [0_u8; 4];
    let _ = tokio::time::timeout(Duration::from_secs(1), stream.read_exact(&mut header)).await;

    let client = connect_tcp(addr).await;
    assert!(
        client
            .list_jobs()
            .await
            .expect("server survived")
            .is_empty()
    );

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn concurrent_clients_can_issue_requests() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let mut tasks = Vec::new();
    for index in 0..8 {
        tasks.push(tokio::spawn(async move {
            let client = DcpClient::connect_tcp(
                addr,
                Hello::new(format!("client-{index}"), ClientKind::Cli),
            )
            .await
            .expect("client connects");
            client.list_jobs().await.expect("list jobs")
        }));
    }

    for task in tasks {
        assert!(task.await.expect("client task joins").is_empty());
    }

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn quic_mtls_handshake_failure_surfaces_cleanly() {
    let (server_config, client_config, no_cert_client_config) = mtls_quic_configs();
    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        never_factories(),
        DcpServerConfig {
            node_id: "quic-agent".to_owned(),
            tcp: None,
            quic: Some(DcpQuicServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
                server_config,
            }),
            auth_token: None,
            metrics_interval: Duration::from_millis(20),
            frame_buffer: 128,
        },
    );
    let handle = server.start().await.expect("QUIC server starts");
    let addr = handle.quic_addr().expect("QUIC addr");

    let failed = DcpClient::connect_quic(
        addr,
        "localhost",
        no_cert_client_config,
        Hello::new("bad-quic", ClientKind::Cli),
    )
    .await;
    assert!(failed.is_err(), "client without cert should fail");

    let client = DcpClient::connect_quic(
        addr,
        "localhost",
        client_config,
        Hello::new("good-quic", ClientKind::Cli),
    )
    .await
    .expect("mTLS client connects");
    assert!(client.list_jobs().await.expect("list over QUIC").is_empty());

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

fn mtls_quic_configs() -> (
    quinn::ServerConfig,
    quinn::ClientConfig,
    quinn::ClientConfig,
) {
    let CertifiedKey {
        cert: server_cert,
        key_pair: server_key,
    } = generate_simple_self_signed(["localhost".to_owned()]).expect("server cert");
    let CertifiedKey {
        cert: client_cert,
        key_pair: client_key,
    } = generate_simple_self_signed(["datum-client".to_owned()]).expect("client cert");

    let mut client_roots = RootCertStore::empty();
    client_roots
        .add(client_cert.der().clone())
        .expect("trust client cert");
    let verifier = WebPkiClientVerifier::builder(Arc::new(client_roots))
        .build()
        .expect("client verifier");
    let server_crypto = RustlsServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(
            vec![server_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der())),
        )
        .expect("server crypto");

    let mut server_roots = RootCertStore::empty();
    server_roots
        .add(server_cert.der().clone())
        .expect("trust server cert");
    let client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(server_roots.clone())
        .with_client_auth_cert(
            vec![client_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(client_key.serialize_der())),
        )
        .expect("client crypto");
    let no_cert_client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(server_roots)
        .with_no_client_auth();

    (
        quinn::ServerConfig::with_crypto(Arc::new(
            QuicServerConfig::try_from(server_crypto).expect("QUIC server config"),
        )),
        quinn::ClientConfig::new(Arc::new(
            QuicClientConfig::try_from(client_crypto).expect("QUIC client config"),
        )),
        quinn::ClientConfig::new(Arc::new(
            QuicClientConfig::try_from(no_cert_client_crypto).expect("QUIC client config"),
        )),
    )
}

async fn recv_raw_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Option<DcpFrame>> {
    let mut header = [0_u8; 4];
    match reader.read_exact(&mut header).await {
        Ok(_) => {}
        Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(error) => return Err(error),
    }
    let len = u32::from_be_bytes(header) as usize;
    if len == 0 || len > 16 * 1024 * 1024 {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            "bad frame length",
        ));
    }
    let mut payload = vec![0_u8; len];
    reader.read_exact(&mut payload).await?;
    Ok(Some(DcpFrame::decode(payload.as_slice()).map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, e)
    })?))
}

async fn send_raw_frame<W: AsyncWrite + Unpin>(
    writer: &mut W,
    frame: &DcpFrame,
) -> std::io::Result<()> {
    let payload = frame.encode_to_vec();
    let len = u32::try_from(payload.len())
        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "frame too large"))?;
    writer.write_all(&len.to_be_bytes()).await?;
    writer.write_all(&payload).await?;
    writer.flush().await
}

#[tokio::test]
async fn event_subscription_terminates_on_server_disconnect() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind listener");
    let addr = listener.local_addr().expect("listener addr");

    let server_task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("accept");
        let hello = recv_raw_frame(&mut stream)
            .await
            .expect("read hello")
            .expect("hello frame");
        assert!(matches!(
            hello.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Hello(_))
        ));
        let hello_resp = datum_agent::dcp::proto::Response::ok(0, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(hello_resp))
            .await
            .expect("write hello response");

        let subscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read subscribe")
            .expect("subscribe frame");
        assert!(matches!(
            subscribe.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(_))
        ));

        let subscribe_resp = datum_agent::dcp::proto::Response::ok(1, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(subscribe_resp))
            .await
            .expect("write subscribe response");

        drop(stream);
    });

    let client = DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("client connects");

    let mut events = client.subscribe_events().await.expect("events subscribe");

    server_task.await.expect("server task");

    let deadline = Instant::now() + Duration::from_secs(5);
    loop {
        match tokio::time::timeout(Duration::from_millis(100), events.recv()).await {
            Ok(None) => break,
            Ok(Some(_)) => {
                if Instant::now() > deadline {
                    panic!("subscription still receiving events after server disconnect");
                }
            }
            Err(_elapsed) => {
                if Instant::now() > deadline {
                    panic!("subscription recv() never completed after server disconnect");
                }
            }
        }
    }
    assert!(
        events.disconnect_reason().is_some(),
        "disconnect reason should be set after server drops connection"
    );
}

#[tokio::test]
async fn metric_subscription_terminates_on_server_disconnect() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind listener");
    let addr = listener.local_addr().expect("listener addr");

    let server_task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("accept");
        let hello = recv_raw_frame(&mut stream)
            .await
            .expect("read hello")
            .expect("hello frame");
        assert!(matches!(
            hello.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Hello(_))
        ));
        let hello_resp = datum_agent::dcp::proto::Response::ok(0, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(hello_resp))
            .await
            .expect("write hello response");

        let subscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read subscribe")
            .expect("subscribe frame");
        assert!(matches!(
            subscribe.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(_))
        ));

        let subscribe_resp = datum_agent::dcp::proto::Response::ok(1, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(subscribe_resp))
            .await
            .expect("write subscribe response");

        drop(stream);
    });

    let client = DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("client connects");

    let mut metrics = client
        .subscribe_metrics(50)
        .await
        .expect("metrics subscribe");

    server_task.await.expect("server task");

    let deadline = Instant::now() + Duration::from_secs(5);
    loop {
        match tokio::time::timeout(Duration::from_millis(100), metrics.recv()).await {
            Ok(None) => break,
            Ok(Some(_)) => {
                if Instant::now() > deadline {
                    panic!("subscription still receiving metrics after server disconnect");
                }
            }
            Err(_elapsed) => {
                if Instant::now() > deadline {
                    panic!("subscription recv() never completed after server disconnect");
                }
            }
        }
    }
    assert!(
        metrics.disconnect_reason().is_some(),
        "disconnect reason should be set after server drops connection"
    );
}