datum-agent 0.10.3

use std::{
    collections::HashMap,
    error::Error,
    fs,
    path::Path,
    sync::Arc,
    time::{Duration, Instant},
};

use datum::{Keep, Sink, Source};
use datum_agent::{
    Agent, AgentConfig, ClusterAgent, ClusterAgentConfig, DesiredJobState, JobMat,
    JobRegistryHandle, JobSpec, JobState, NodeSessionConfig, NodeSessionTransport,
    dcp::{DcpJobFactories, DcpQuicServerConfig, DcpServer, DcpServerConfig, DcpTcpServerConfig},
};
use datum_cluster::ClusterConfig;
use datum_net::quic::{
    crypto::rustls::{QuicClientConfig, QuicServerConfig},
    quinn,
    rustls::{
        ClientConfig as RustlsClientConfig, RootCertStore, ServerConfig as RustlsServerConfig,
        pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
        server::WebPkiClientVerifier,
    },
};

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let config_values = load_config_values()?;
    let factories = registered_factories();
    if cluster_enabled(&config_values) {
        let config = cluster_agent_config(&config_values)?;
        let handle = ClusterAgent::start(config, factories).await?;
        if let Some(addr) = handle.tcp_addr() {
            eprintln!("datum-agent DCP TCP listening on {addr}");
        }
        if let Some(addr) = handle.quic_addr() {
            eprintln!("datum-agent DCP QUIC listening on {addr}");
        }
        eprintln!(
            "datum-agent cluster node {} gossip listening on {}",
            handle.cluster().node_id(),
            handle.cluster().advertise_addr()
        );

        tokio::signal::ctrl_c().await?;
        drain_jobs(handle.agent().registry().clone(), Duration::from_secs(30)).await;
        handle.shutdown().await?;
    } else {
        let agent = Agent::start_with_config(AgentConfig::default())?;
        let server_config = dcp_server_config(&config_values)?;
        let server = DcpServer::from_agent(&agent, factories, server_config);
        let handle = server.start().await?;

        if let Some(addr) = handle.tcp_addr() {
            eprintln!("datum-agent DCP TCP listening on {addr}");
        }
        if let Some(addr) = handle.quic_addr() {
            eprintln!("datum-agent DCP QUIC listening on {addr}");
        }

        tokio::signal::ctrl_c().await?;
        drain_jobs(agent.registry().clone(), Duration::from_secs(30)).await;
        handle.shutdown().await;
        agent.registry().shutdown()?;
    }
    Ok(())
}

fn load_config_values() -> Result<HashMap<String, String>, Box<dyn Error>> {
    let args = daemon_args()?;
    let mut values = HashMap::new();
    if let Some(path) = args
        .config_path
        .or_else(|| std::env::var("DATUM_AGENT_CONFIG").ok())
    {
        let content = fs::read_to_string(path)?;
        let mut section = String::new();
        for line in content.lines() {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                continue;
            }
            if line.starts_with('[') && line.ends_with(']') {
                section = line[1..line.len() - 1].trim().to_ascii_lowercase();
                continue;
            }
            if let Some((key, value)) = line.split_once('=') {
                values.insert(config_key(&section, key.trim()), value.trim().to_owned());
            }
        }
    }

    for (key, value) in std::env::vars() {
        if key.starts_with("DATUM_AGENT_") {
            values.insert(key, value);
        }
    }
    if args.cluster {
        values.insert("DATUM_AGENT_CLUSTER".to_owned(), "1".to_owned());
    }
    Ok(values)
}

#[derive(Default)]
struct DaemonArgs {
    config_path: Option<String>,
    cluster: bool,
}

fn daemon_args() -> Result<DaemonArgs, Box<dyn Error>> {
    let mut parsed = DaemonArgs::default();
    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--cluster" => parsed.cluster = true,
            "--config" => {
                let Some(path) = args.next() else {
                    return Err("--config requires a following config file path".into());
                };
                parsed.config_path = Some(path);
            }
            other if !other.starts_with('-') && parsed.config_path.is_none() => {
                parsed.config_path = Some(other.to_owned());
            }
            other => return Err(format!("unknown datum-agent argument: {other}").into()),
        }
    }
    Ok(parsed)
}

fn config_key(section: &str, key: &str) -> String {
    if section.is_empty() {
        return key.to_owned();
    }
    let normalized = key
        .chars()
        .map(|ch| {
            if ch == '-' || ch == '.' {
                '_'
            } else {
                ch.to_ascii_uppercase()
            }
        })
        .collect::<String>();
    match section {
        "cluster" => format!("DATUM_AGENT_CLUSTER_{normalized}"),
        "dcp" => format!("DATUM_AGENT_DCP_{normalized}"),
        _ => key.to_owned(),
    }
}

fn dcp_server_config(values: &HashMap<String, String>) -> Result<DcpServerConfig, Box<dyn Error>> {
    let mut config = DcpServerConfig::default();
    if let Some(node_id) = values.get("DATUM_AGENT_DCP_NODE_ID") {
        config.node_id = node_id.clone();
    }
    if let Some(token) = values.get("DATUM_AGENT_DCP_AUTH_TOKEN") {
        config.auth_token = Some(token.clone());
    }
    if let Some(interval) = values.get("DATUM_AGENT_DCP_METRICS_INTERVAL_MS") {
        config.metrics_interval = Duration::from_millis(interval.parse()?);
    }

    let tcp_enabled = values
        .get("DATUM_AGENT_DCP_TCP")
        .map(|value| value != "0" && value != "false")
        .unwrap_or(true);
    config.tcp = if tcp_enabled {
        let addr = values
            .get("DATUM_AGENT_DCP_TCP_ADDR")
            .map(String::as_str)
            .unwrap_or("127.0.0.1:9555")
            .parse()?;
        Some(DcpTcpServerConfig { addr })
    } else {
        None
    };

    config.quic = match values.get("DATUM_AGENT_DCP_QUIC_ADDR") {
        Some(addr) => Some(DcpQuicServerConfig {
            addr: addr.parse()?,
            server_config: load_quic_server_config(values)?,
        }),
        None => None,
    };

    Ok(config)
}

fn cluster_enabled(values: &HashMap<String, String>) -> bool {
    values
        .get("DATUM_AGENT_CLUSTER")
        .or_else(|| values.get("DATUM_AGENT_CLUSTER_ENABLED"))
        .map(|value| value != "0" && value != "false")
        .unwrap_or(false)
}

fn cluster_agent_config(
    values: &HashMap<String, String>,
) -> Result<ClusterAgentConfig, Box<dyn Error>> {
    let dcp = dcp_server_config(values)?;
    let mut cluster = ClusterConfig {
        node_id: values
            .get("DATUM_AGENT_CLUSTER_NODE_ID")
            .or_else(|| values.get("DATUM_AGENT_DCP_NODE_ID"))
            .cloned()
            .unwrap_or_else(|| dcp.node_id.clone()),
        ..ClusterConfig::default()
    };
    if let Some(roles) = values.get("DATUM_AGENT_CLUSTER_ROLES") {
        cluster.roles = split_csv(roles);
    }
    if let Some(addr) = values.get("DATUM_AGENT_CLUSTER_BIND_ADDR") {
        cluster.bind_addr = addr.parse()?;
    }
    if let Some(addr) = values.get("DATUM_AGENT_CLUSTER_ADVERTISE_ADDR") {
        cluster.advertise_addr = addr.parse()?;
    }
    if let Some(seeds) = values.get("DATUM_AGENT_CLUSTER_SEEDS") {
        cluster.seed_nodes = split_csv(seeds)
            .into_iter()
            .map(|seed| seed.parse())
            .collect::<Result<Vec<_>, _>>()?;
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_GOSSIP_INTERVAL_MS") {
        cluster.gossip_interval = Duration::from_millis(ms.parse()?);
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_PROBE_TIMEOUT_MS") {
        cluster.probe_timeout = Duration::from_millis(ms.parse()?);
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_DOWNING_TIMEOUT_MS") {
        cluster.downing_timeout = Duration::from_millis(ms.parse()?);
    }

    let mut sessions = NodeSessionConfig::default();
    if let Some(role) = values.get("DATUM_AGENT_CLUSTER_AGENT_ROLE") {
        sessions.agent_role = role.clone();
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_REQUEST_TIMEOUT_MS") {
        sessions.request_timeout = Duration::from_millis(ms.parse()?);
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_RECONNECT_MIN_BACKOFF_MS") {
        sessions.reconnect_min_backoff = Duration::from_millis(ms.parse()?);
    }
    if let Some(ms) = values.get("DATUM_AGENT_CLUSTER_RECONNECT_MAX_BACKOFF_MS") {
        sessions.reconnect_max_backoff = Duration::from_millis(ms.parse()?);
    }
    if values
        .get("DATUM_AGENT_CLUSTER_DCP_TRANSPORT")
        .map(|value| value.eq_ignore_ascii_case("quic"))
        .unwrap_or(false)
    {
        sessions.transport = NodeSessionTransport::QuicMtls {
            server_name: values
                .get("DATUM_AGENT_CLUSTER_DCP_QUIC_SERVER_NAME")
                .cloned()
                .unwrap_or_else(|| "localhost".to_owned()),
            client_config: load_quic_client_config(values)?,
        };
    }

    Ok(ClusterAgentConfig {
        agent: AgentConfig::default(),
        cluster,
        dcp,
        sessions,
    })
}

fn split_csv(value: &str) -> Vec<String> {
    value
        .split(',')
        .map(str::trim)
        .filter(|item| !item.is_empty())
        .map(ToOwned::to_owned)
        .collect()
}

fn load_quic_server_config(
    values: &HashMap<String, String>,
) -> Result<quinn::ServerConfig, Box<dyn Error>> {
    let cert = read_required_der(values, "DATUM_AGENT_DCP_QUIC_CERT_DER")?;
    let key = read_required_der(values, "DATUM_AGENT_DCP_QUIC_KEY_DER")?;
    let client_ca = read_required_der(values, "DATUM_AGENT_DCP_QUIC_CLIENT_CA_DER")?;

    let mut client_roots = RootCertStore::empty();
    client_roots.add(CertificateDer::from(client_ca))?;
    let verifier = WebPkiClientVerifier::builder(Arc::new(client_roots)).build()?;

    let server_crypto = RustlsServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(
            vec![CertificateDer::from(cert)],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key)),
        )?;
    Ok(quinn::ServerConfig::with_crypto(Arc::new(
        QuicServerConfig::try_from(server_crypto)?,
    )))
}

fn load_quic_client_config(
    values: &HashMap<String, String>,
) -> Result<quinn::ClientConfig, Box<dyn Error>> {
    let server_ca = read_required_der(values, "DATUM_AGENT_CLUSTER_DCP_QUIC_SERVER_CA_DER")?;
    let client_cert = read_required_der(values, "DATUM_AGENT_CLUSTER_DCP_QUIC_CLIENT_CERT_DER")?;
    let client_key = read_required_der(values, "DATUM_AGENT_CLUSTER_DCP_QUIC_CLIENT_KEY_DER")?;

    let mut roots = RootCertStore::empty();
    roots.add(CertificateDer::from(server_ca))?;
    let client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(roots)
        .with_client_auth_cert(
            vec![CertificateDer::from(client_cert)],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(client_key)),
        )?;
    Ok(quinn::ClientConfig::new(Arc::new(
        QuicClientConfig::try_from(client_crypto)?,
    )))
}

fn read_required_der(
    values: &HashMap<String, String>,
    key: &str,
) -> Result<Vec<u8>, Box<dyn Error>> {
    let path = values
        .get(key)
        .ok_or_else(|| format!("missing required config key: {key}"))?;
    Ok(fs::read(Path::new(path))?)
}

fn registered_factories() -> DcpJobFactories {
    let factories = DcpJobFactories::new();
    factories
        .register("empty", |instance_name, _params| {
            Ok(JobSpec::new(instance_name, |context| {
                let control = context.control();
                Ok(Source::<u64>::empty()
                    .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("built-in empty factory registers");
    factories
}

async fn drain_jobs(registry: JobRegistryHandle, timeout: Duration) {
    let _ = tokio::task::spawn_blocking({
        let registry = registry.clone();
        move || {
            if let Ok(statuses) = registry.list() {
                for status in statuses {
                    if matches!(
                        status.desired_state,
                        DesiredJobState::Running | DesiredJobState::Draining
                    ) {
                        let _ = registry.drain(status.name);
                    }
                }
            }
        }
    })
    .await;

    let deadline = Instant::now() + timeout;
    loop {
        let registry = registry.clone();
        let done = tokio::task::spawn_blocking(move || {
            registry
                .list()
                .map(|statuses| {
                    statuses.iter().all(|status| {
                        matches!(
                            status.state,
                            JobState::Completed
                                | JobState::Drained
                                | JobState::Stopped
                                | JobState::Failed
                        )
                    })
                })
                .unwrap_or(true)
        })
        .await
        .unwrap_or(true);
        if done || Instant::now() >= deadline {
            break;
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
}