nemo-relay 0.5.0

// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Agent Trajectory Observability Format (ATOF) JSONL exporter support for NeMo
//! Flow.
//!
//! The [`AtofExporter`] registers as an event subscriber and writes each
//! canonical NeMo Relay Agent Trajectory Observability Format (ATOF) event as
//! one JSON object per JSONL line.

use std::collections::HashMap;
use std::fs::{File, OpenOptions, create_dir_all};
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, mpsc as std_mpsc};
use std::time::Duration;

use chrono::Utc;
#[cfg(feature = "atof-streaming")]
use futures_util::{SinkExt, stream};
use serde::{Deserialize, Serialize};
use serde_json::Value as Json;
#[cfg(feature = "atof-streaming")]
use tokio_tungstenite::tungstenite::client::IntoClientRequest;

use crate::api::event::Event;
use crate::api::runtime::EventSubscriberFn;
use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber};
use crate::error::FlowError;

/// Result type for the ATOF JSONL exporter.
pub type Result<T> = std::result::Result<T, AtofExporterError>;

/// Errors produced while configuring or operating the ATOF JSONL exporter.
#[derive(Debug, thiserror::Error)]
pub enum AtofExporterError {
    /// Failed to resolve the current working directory for default config.
    #[error("failed to resolve current working directory: {0}")]
    CurrentDirectory(std::io::Error),
    /// Failed to open the output file.
    #[error("failed to open ATOF output file {path:?}: {source}")]
    OpenFile {
        /// Output path that failed to open.
        path: PathBuf,
        /// Underlying I/O error.
        source: std::io::Error,
    },
    /// Failed while flushing the output file.
    #[error("failed to flush ATOF output file {path:?}: {source}")]
    Flush {
        /// Output path that failed to flush.
        path: PathBuf,
        /// Underlying I/O error.
        source: std::io::Error,
    },
    /// The exporter recorded an earlier write or serialization error.
    #[error("previous ATOF export failed for {path:?}: {message}")]
    StoredFailure {
        /// Output path associated with the failure.
        path: PathBuf,
        /// Stored failure message.
        message: String,
    },
    /// A streaming endpoint configuration is invalid.
    #[error("invalid ATOF streaming endpoint: {0}")]
    InvalidEndpoint(String),
    /// The internal exporter state lock was poisoned.
    #[error("the ATOF exporter state lock was poisoned")]
    LockPoisoned,
    /// Runtime subscriber registration failed.
    #[error(transparent)]
    Runtime(#[from] FlowError),
}

/// File write behavior for [`AtofExporter`].
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AtofExporterMode {
    /// Append events to an existing file or create it if missing.
    #[default]
    Append,
    /// Truncate an existing file when the exporter is created.
    Overwrite,
}

impl AtofExporterMode {
    /// Parse a string mode used by language bindings.
    pub fn parse(value: &str) -> Option<Self> {
        match value {
            "append" => Some(Self::Append),
            "overwrite" => Some(Self::Overwrite),
            _ => None,
        }
    }

    /// Return the stable string representation used by language bindings.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Append => "append",
            Self::Overwrite => "overwrite",
        }
    }
}

/// Streaming transport used by an ATOF endpoint.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AtofEndpointTransport {
    /// POST each event as one JSONL record.
    #[default]
    HttpPost,
    /// Send each event as one WebSocket JSON text message.
    Websocket,
    /// Stream events over one long-lived HTTP NDJSON upload.
    Ndjson,
}

/// Field name transformation policy used before sending events to an endpoint.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AtofEndpointFieldNamePolicy {
    /// Preserve canonical ATOF field names exactly.
    #[default]
    Preserve,
    /// Replace dots in JSON object keys with underscores, recursively.
    ReplaceDots,
}

impl AtofEndpointFieldNamePolicy {
    /// Parse a string policy used by configuration and bindings.
    pub fn parse(value: &str) -> Option<Self> {
        match value {
            "preserve" => Some(Self::Preserve),
            "replace_dots" => Some(Self::ReplaceDots),
            _ => None,
        }
    }

    /// Return the stable string representation used by configuration and bindings.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Preserve => "preserve",
            Self::ReplaceDots => "replace_dots",
        }
    }
}

impl AtofEndpointTransport {
    /// Parse a string transport used by configuration and bindings.
    pub fn parse(value: &str) -> Option<Self> {
        match value {
            "http_post" => Some(Self::HttpPost),
            "websocket" => Some(Self::Websocket),
            "ndjson" => Some(Self::Ndjson),
            _ => None,
        }
    }

    /// Return the stable string representation used by configuration and bindings.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::HttpPost => "http_post",
            Self::Websocket => "websocket",
            Self::Ndjson => "ndjson",
        }
    }
}

/// Streaming destination for raw ATOF events.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AtofEndpointConfig {
    /// Endpoint URL.
    pub url: String,
    /// Endpoint transport.
    #[serde(default)]
    pub transport: AtofEndpointTransport,
    /// Headers applied to endpoint requests or handshakes.
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub headers: HashMap<String, String>,
    /// Per-endpoint timeout in milliseconds.
    #[serde(default = "default_endpoint_timeout_millis")]
    pub timeout_millis: u64,
    /// Field name transformation policy applied before sending events.
    #[serde(default)]
    pub field_name_policy: AtofEndpointFieldNamePolicy,
}

impl AtofEndpointConfig {
    /// Create a streaming endpoint with defaults.
    pub fn new(url: impl Into<String>, transport: AtofEndpointTransport) -> Self {
        Self {
            url: url.into(),
            transport,
            headers: HashMap::new(),
            timeout_millis: default_endpoint_timeout_millis(),
            field_name_policy: AtofEndpointFieldNamePolicy::Preserve,
        }
    }

    /// Add a header to this endpoint config.
    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.headers.insert(key.into(), value.into());
        self
    }

    /// Override the endpoint timeout.
    pub fn with_timeout_millis(mut self, timeout_millis: u64) -> Self {
        self.timeout_millis = timeout_millis;
        self
    }

    /// Override the endpoint field name policy.
    pub fn with_field_name_policy(
        mut self,
        field_name_policy: AtofEndpointFieldNamePolicy,
    ) -> Self {
        self.field_name_policy = field_name_policy;
        self
    }
}

/// Configuration for [`AtofExporter`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AtofExporterConfig {
    /// Directory that contains the JSONL output file.
    #[serde(default = "default_output_directory")]
    pub output_directory: PathBuf,
    /// Append or overwrite behavior used when opening the file.
    #[serde(default)]
    pub mode: AtofExporterMode,
    /// Output filename.
    #[serde(default = "default_filename")]
    pub filename: String,
    /// Optional streaming endpoints that receive every raw ATOF event.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub endpoints: Vec<AtofEndpointConfig>,
}

impl Default for AtofExporterConfig {
    fn default() -> Self {
        Self {
            output_directory: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
            mode: AtofExporterMode::Append,
            filename: default_filename(),
            endpoints: Vec::new(),
        }
    }
}

impl AtofExporterConfig {
    /// Create a config with defaults.
    pub fn new() -> Self {
        Self::default()
    }

    /// Override the output directory.
    pub fn with_output_directory(mut self, output_directory: impl Into<PathBuf>) -> Self {
        self.output_directory = output_directory.into();
        self
    }

    /// Override the output mode.
    pub fn with_mode(mut self, mode: AtofExporterMode) -> Self {
        self.mode = mode;
        self
    }

    /// Override the output filename.
    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = filename.into();
        self
    }

    /// Override streaming endpoints.
    pub fn with_endpoints(mut self, endpoints: Vec<AtofEndpointConfig>) -> Self {
        self.endpoints = endpoints;
        self
    }

    /// Add one streaming endpoint.
    pub fn with_endpoint(mut self, endpoint: AtofEndpointConfig) -> Self {
        self.endpoints.push(endpoint);
        self
    }

    /// Return the full output path for this config.
    pub fn path(&self) -> PathBuf {
        self.output_directory.join(&self.filename)
    }
}

struct AtofExporterState {
    writer: BufWriter<File>,
    last_error: Option<String>,
    endpoints: Vec<AtofEndpointWorker>,
    closed: bool,
}

/// Filesystem-backed Agent Trajectory Observability Format (ATOF) JSONL event exporter.
pub struct AtofExporter {
    path: PathBuf,
    state: Arc<Mutex<AtofExporterState>>,
}

impl AtofExporter {
    /// Create a new exporter from config and open its output file.
    pub fn new(config: AtofExporterConfig) -> Result<Self> {
        let path = config.path();
        create_dir_all(&config.output_directory).map_err(|source| AtofExporterError::OpenFile {
            path: path.clone(),
            source,
        })?;
        let file = open_file(&path, config.mode)?;
        let endpoints = start_endpoint_workers(&config.endpoints)?;
        Ok(Self {
            path,
            state: Arc::new(Mutex::new(AtofExporterState {
                writer: BufWriter::new(file),
                last_error: None,
                endpoints,
                closed: false,
            })),
        })
    }

    /// Return the output JSONL path.
    pub fn path(&self) -> &Path {
        self.path.as_path()
    }

    /// Return an event subscriber that writes one JSONL record per observed event.
    pub fn subscriber(&self) -> EventSubscriberFn {
        let state = Arc::clone(&self.state);
        Arc::new(move |event: &Event| {
            let Ok(mut state) = state.lock() else {
                return;
            };
            if state.closed || state.last_error.is_some() {
                return;
            }
            let Ok(value) = event.try_to_json_value() else {
                state.last_error = Some("failed to serialize ATOF event".to_string());
                return;
            };
            if let Err(error) = write_json_value(&mut state.writer, &value) {
                state.last_error = Some(error);
                return;
            }
            let Ok(raw_json) = serde_json::to_string(&value) else {
                state.last_error = Some("failed to serialize ATOF event".to_string());
                return;
            };
            for endpoint in &state.endpoints {
                endpoint.enqueue(raw_json.clone());
            }
        })
    }

    /// Register this exporter globally under the given subscriber name.
    pub fn register(&self, name: &str) -> Result<()> {
        register_subscriber(name, self.subscriber()).map_err(Into::into)
    }

    /// Deregister a global subscriber by name.
    pub fn deregister(&self, name: &str) -> Result<bool> {
        deregister_subscriber(name).map_err(Into::into)
    }

    /// Flush the underlying file and drain queued endpoint events.
    pub fn force_flush(&self) -> Result<()> {
        flush_subscribers()?;
        let mut state = self
            .state
            .lock()
            .map_err(|_| AtofExporterError::LockPoisoned)?;
        if state.closed {
            return stored_failure_result(&self.path, &state);
        }
        state
            .writer
            .flush()
            .map_err(|source| AtofExporterError::Flush {
                path: self.path.clone(),
                source,
            })?;
        for endpoint in &state.endpoints {
            endpoint.flush();
        }
        stored_failure_result(&self.path, &state)
    }

    /// Shut down the exporter by flushing buffered data and closing endpoints.
    pub fn shutdown(&self) -> Result<()> {
        flush_subscribers()?;
        let mut state = self
            .state
            .lock()
            .map_err(|_| AtofExporterError::LockPoisoned)?;
        if state.closed {
            return stored_failure_result(&self.path, &state);
        }
        state.closed = true;
        let flush_result = state
            .writer
            .flush()
            .map_err(|source| AtofExporterError::Flush {
                path: self.path.clone(),
                source,
            });
        for endpoint in &state.endpoints {
            endpoint.close();
        }
        flush_result?;
        stored_failure_result(&self.path, &state)
    }
}

fn default_filename() -> String {
    format!(
        "nemo-relay-events-{}.jsonl",
        Utc::now().format("%Y-%m-%d-%H.%M.%S")
    )
}

fn default_output_directory() -> PathBuf {
    std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
}

fn default_endpoint_timeout_millis() -> u64 {
    3_000
}

fn open_file(path: &Path, mode: AtofExporterMode) -> Result<File> {
    let mut options = OpenOptions::new();
    options.create(true);
    match mode {
        AtofExporterMode::Append => {
            options.append(true);
        }
        AtofExporterMode::Overwrite => {
            options.write(true).truncate(true);
        }
    }
    options
        .open(path)
        .map_err(|source| AtofExporterError::OpenFile {
            path: path.to_path_buf(),
            source,
        })
}

fn write_json_value(writer: &mut BufWriter<File>, value: &Json) -> std::result::Result<(), String> {
    serde_json::to_writer(&mut *writer, value).map_err(|error| error.to_string())?;
    writer.write_all(b"\n").map_err(|error| error.to_string())?;
    writer.flush().map_err(|error| error.to_string())
}

fn stored_failure_result(path: &Path, state: &AtofExporterState) -> Result<()> {
    if let Some(message) = &state.last_error {
        return Err(AtofExporterError::StoredFailure {
            path: path.to_path_buf(),
            message: message.clone(),
        });
    }
    Ok(())
}

#[cfg_attr(not(feature = "atof-streaming"), allow(dead_code))]
enum EndpointMessage {
    Event(String),
    Flush(std_mpsc::Sender<()>),
    Close(std_mpsc::Sender<()>),
}

#[cfg(feature = "atof-streaming")]
enum NdjsonBodyMessage {
    Event(Vec<u8>),
    Flush(std_mpsc::Sender<()>),
}

#[cfg(feature = "atof-streaming")]
impl NdjsonBodyMessage {
    fn acknowledge_if_flush(self) {
        if let Self::Flush(done) = self {
            let _ = done.send(());
        }
    }
}

struct AtofEndpointWorker {
    sender: tokio::sync::mpsc::UnboundedSender<EndpointMessage>,
    timeout: Duration,
}

impl AtofEndpointWorker {
    fn enqueue(&self, raw_json: String) {
        let _ = self.sender.send(EndpointMessage::Event(raw_json));
    }

    fn flush(&self) {
        let (tx, rx) = std_mpsc::channel();
        if self.sender.send(EndpointMessage::Flush(tx)).is_ok()
            && rx.recv_timeout(self.timeout).is_err()
        {
            eprintln!("nemo_relay: timed out flushing ATOF endpoint");
        }
    }

    fn close(&self) {
        let (tx, rx) = std_mpsc::channel();
        if self.sender.send(EndpointMessage::Close(tx)).is_ok()
            && rx.recv_timeout(self.timeout).is_err()
        {
            eprintln!("nemo_relay: timed out closing ATOF endpoint");
        }
    }
}

#[cfg(feature = "atof-streaming")]
fn start_endpoint_workers(configs: &[AtofEndpointConfig]) -> Result<Vec<AtofEndpointWorker>> {
    let mut workers = Vec::with_capacity(configs.len());
    for (index, config) in configs.iter().enumerate() {
        match start_endpoint_worker(index, config.clone()) {
            Ok(worker) => workers.push(worker),
            Err(error) => {
                eprintln!("nemo_relay: invalid ATOF endpoint[{index}]: {error}");
                return Err(AtofExporterError::InvalidEndpoint(format!(
                    "endpoints[{index}]: {error}"
                )));
            }
        }
    }
    Ok(workers)
}

#[cfg(not(feature = "atof-streaming"))]
fn start_endpoint_workers(configs: &[AtofEndpointConfig]) -> Result<Vec<AtofEndpointWorker>> {
    if configs.is_empty() {
        return Ok(Vec::new());
    }
    let message = "ATOF streaming endpoints are not supported in this build".to_string();
    eprintln!("nemo_relay: {message}");
    Err(AtofExporterError::InvalidEndpoint(message))
}

#[cfg(feature = "atof-streaming")]
fn start_endpoint_worker(index: usize, config: AtofEndpointConfig) -> Result<AtofEndpointWorker> {
    validate_endpoint_config(&config)?;
    let timeout = Duration::from_millis(config.timeout_millis);
    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
    std::thread::Builder::new()
        .name(format!("nemo-relay-atof-endpoint-{index}"))
        .spawn(move || run_endpoint_worker(index, config, rx))
        .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
    Ok(AtofEndpointWorker {
        sender: tx,
        timeout,
    })
}

#[cfg(feature = "atof-streaming")]
fn validate_endpoint_config(config: &AtofEndpointConfig) -> Result<()> {
    if config.url.trim().is_empty() {
        return Err(AtofExporterError::InvalidEndpoint(
            "endpoint url must be non-empty".to_string(),
        ));
    }
    if config.timeout_millis == 0 {
        return Err(AtofExporterError::InvalidEndpoint(
            "endpoint timeout_millis must be greater than 0".to_string(),
        ));
    }
    let url = reqwest::Url::parse(&config.url)
        .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
    let valid_scheme = match config.transport {
        AtofEndpointTransport::HttpPost | AtofEndpointTransport::Ndjson => {
            matches!(url.scheme(), "http" | "https")
        }
        AtofEndpointTransport::Websocket => matches!(url.scheme(), "ws" | "wss"),
    };
    if !valid_scheme {
        return Err(AtofExporterError::InvalidEndpoint(format!(
            "endpoint {} transport does not support URL scheme {:?}",
            config.transport.as_str(),
            url.scheme()
        )));
    }
    build_header_map(&config.headers)?;
    Ok(())
}

#[cfg(feature = "atof-streaming")]
fn build_header_map(headers: &HashMap<String, String>) -> Result<reqwest::header::HeaderMap> {
    let mut out = reqwest::header::HeaderMap::new();
    for (key, value) in headers {
        let name = reqwest::header::HeaderName::from_bytes(key.as_bytes())
            .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
        let value = reqwest::header::HeaderValue::from_str(value)
            .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
        out.insert(name, value);
    }
    Ok(out)
}

#[cfg(feature = "atof-streaming")]
fn run_endpoint_worker(
    index: usize,
    config: AtofEndpointConfig,
    rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
) {
    let runtime = match tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
    {
        Ok(runtime) => runtime,
        Err(error) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] runtime failed: {error}");
            return;
        }
    };
    runtime.block_on(async move {
        match config.transport {
            AtofEndpointTransport::HttpPost => run_http_post_endpoint(index, config, rx).await,
            AtofEndpointTransport::Websocket => run_websocket_endpoint(index, config, rx).await,
            AtofEndpointTransport::Ndjson => run_ndjson_endpoint(index, config, rx).await,
        }
    });
}

#[cfg(feature = "atof-streaming")]
async fn run_http_post_endpoint(
    index: usize,
    config: AtofEndpointConfig,
    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
) {
    let client = match reqwest::Client::builder()
        .timeout(Duration::from_millis(config.timeout_millis))
        .default_headers(match build_header_map(&config.headers) {
            Ok(headers) => headers,
            Err(error) => {
                eprintln!("nemo_relay: ATOF endpoint[{index}] disabled: {error}");
                drain_closed(rx).await;
                return;
            }
        })
        .build()
    {
        Ok(client) => client,
        Err(error) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] client build failed: {error}");
            drain_closed(rx).await;
            return;
        }
    };
    while let Some(message) = rx.recv().await {
        match message {
            EndpointMessage::Event(raw_json) => {
                let body = format!("{}\n", endpoint_event_json(&config, raw_json));
                let result = client
                    .post(&config.url)
                    .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson")
                    .body(body)
                    .send()
                    .await;
                match result {
                    Ok(response) if response.status().is_success() => {}
                    Ok(response) => log_http_error(index, "HTTP", response).await,
                    Err(error) => {
                        eprintln!("nemo_relay: ATOF endpoint[{index}] send failed: {error}")
                    }
                }
            }
            EndpointMessage::Flush(done) => {
                let _ = done.send(());
            }
            EndpointMessage::Close(done) => {
                let _ = done.send(());
                return;
            }
        }
    }
}

#[cfg(feature = "atof-streaming")]
async fn run_websocket_endpoint(
    index: usize,
    config: AtofEndpointConfig,
    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
) {
    let mut pending = std::collections::VecDeque::new();
    let mut socket = match connect_websocket(&config).await {
        Ok(socket) => Some(socket),
        Err(error) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] websocket startup failed: {error}");
            None
        }
    };
    while let Some(message) = rx.recv().await {
        match message {
            EndpointMessage::Event(raw_json) => {
                pending.push_back(endpoint_event_json(&config, raw_json));
                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
            }
            EndpointMessage::Flush(done) => {
                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
                let _ = done.send(());
            }
            EndpointMessage::Close(done) => {
                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
                if let Some(mut ws) = socket.take() {
                    let _ = ws.close(None).await;
                }
                let _ = done.send(());
                return;
            }
        }
    }
}

#[cfg(feature = "atof-streaming")]
type AtofWebSocket =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;

#[cfg(feature = "atof-streaming")]
async fn drain_websocket_pending(
    index: usize,
    config: &AtofEndpointConfig,
    socket: &mut Option<AtofWebSocket>,
    pending: &mut std::collections::VecDeque<String>,
) -> bool {
    let timeout = Duration::from_millis(config.timeout_millis);
    match tokio::time::timeout(
        timeout,
        drain_websocket_pending_inner(index, config, socket, pending),
    )
    .await
    {
        Ok(drained) => drained,
        Err(_) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] websocket drain timed out");
            false
        }
    }
}

#[cfg(feature = "atof-streaming")]
async fn drain_websocket_pending_inner(
    index: usize,
    config: &AtofEndpointConfig,
    socket: &mut Option<AtofWebSocket>,
    pending: &mut std::collections::VecDeque<String>,
) -> bool {
    while let Some(raw_json) = pending.front().cloned() {
        if socket.is_none() {
            match connect_websocket(config).await {
                Ok(ws) => *socket = Some(ws),
                Err(error) => {
                    eprintln!(
                        "nemo_relay: ATOF endpoint[{index}] websocket reconnect failed: {error}"
                    );
                    tokio::time::sleep(Duration::from_millis(50)).await;
                    continue;
                }
            }
        }

        let Some(ws) = socket.as_mut() else {
            continue;
        };
        match ws
            .send(tokio_tungstenite::tungstenite::Message::Text(
                raw_json.into(),
            ))
            .await
        {
            Ok(()) => {
                pending.pop_front();
            }
            Err(error) => {
                eprintln!("nemo_relay: ATOF endpoint[{index}] websocket send failed: {error}");
                *socket = None;
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        }
    }
    true
}

#[cfg(feature = "atof-streaming")]
async fn connect_websocket(
    config: &AtofEndpointConfig,
) -> std::result::Result<AtofWebSocket, String> {
    let mut request = config
        .url
        .as_str()
        .into_client_request()
        .map_err(|error| error.to_string())?;
    for (key, value) in &config.headers {
        let name =
            tokio_tungstenite::tungstenite::http::header::HeaderName::from_bytes(key.as_bytes())
                .map_err(|error| error.to_string())?;
        let value = tokio_tungstenite::tungstenite::http::header::HeaderValue::from_str(value)
            .map_err(|error| error.to_string())?;
        request.headers_mut().insert(name, value);
    }
    tokio::time::timeout(
        Duration::from_millis(config.timeout_millis),
        tokio_tungstenite::connect_async(request),
    )
    .await
    .map_err(|error| error.to_string())?
    .map(|(socket, _)| socket)
    .map_err(|error| error.to_string())
}

#[cfg(feature = "atof-streaming")]
async fn run_ndjson_endpoint(
    index: usize,
    config: AtofEndpointConfig,
    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
) {
    let client = match build_ndjson_client(&config) {
        Ok(client) => client,
        Err(error) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] {error}");
            drain_closed(rx).await;
            return;
        }
    };

    let (body_tx, body) = ndjson_body_channel();
    let url = config.url.clone();
    let request = tokio::spawn(async move {
        client
            .post(url)
            .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson")
            .body(body)
            .send()
            .await
    });
    let close_timeout = Duration::from_millis(config.timeout_millis);

    while let Some(message) = rx.recv().await {
        match message {
            EndpointMessage::Event(raw_json) => {
                send_ndjson_event(index, &body_tx, endpoint_event_json(&config, raw_json))
            }
            EndpointMessage::Flush(done) => send_ndjson_flush(index, &body_tx, done),
            EndpointMessage::Close(done) => {
                drop(body_tx);
                finish_ndjson_upload(index, request, close_timeout, done).await;
                return;
            }
        }
    }
}

#[cfg(feature = "atof-streaming")]
fn build_ndjson_client(
    config: &AtofEndpointConfig,
) -> std::result::Result<reqwest::Client, String> {
    let headers =
        build_header_map(&config.headers).map_err(|error| format!("disabled: {error}"))?;
    reqwest::Client::builder()
        .connect_timeout(Duration::from_millis(config.timeout_millis))
        .default_headers(headers)
        .build()
        .map_err(|error| format!("client build failed: {error}"))
}

#[cfg(feature = "atof-streaming")]
fn ndjson_body_channel() -> (
    tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
    reqwest::Body,
) {
    let (body_tx, body_rx) = tokio::sync::mpsc::unbounded_channel::<NdjsonBodyMessage>();
    let body_stream = stream::unfold(body_rx, |mut body_rx| async {
        loop {
            match body_rx.recv().await? {
                NdjsonBodyMessage::Event(bytes) => {
                    return Some((Ok::<_, std::io::Error>(bytes), body_rx));
                }
                NdjsonBodyMessage::Flush(done) => {
                    let _ = done.send(());
                }
            }
        }
    });
    (body_tx, reqwest::Body::wrap_stream(body_stream))
}

#[cfg(feature = "atof-streaming")]
fn send_ndjson_event(
    index: usize,
    body_tx: &tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
    raw_json: String,
) {
    if let Err(error) = body_tx.send(NdjsonBodyMessage::Event(
        format!("{raw_json}\n").into_bytes(),
    )) {
        eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON send failed: {error}");
    }
}

#[cfg(feature = "atof-streaming")]
fn send_ndjson_flush(
    index: usize,
    body_tx: &tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
    done: std_mpsc::Sender<()>,
) {
    if let Err(error) = body_tx.send(NdjsonBodyMessage::Flush(done)) {
        eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON flush failed: {error}");
        error.0.acknowledge_if_flush();
    }
}

#[cfg(feature = "atof-streaming")]
async fn finish_ndjson_upload(
    index: usize,
    request: tokio::task::JoinHandle<reqwest::Result<reqwest::Response>>,
    close_timeout: Duration,
    done: std_mpsc::Sender<()>,
) {
    match tokio::time::timeout(close_timeout, request).await {
        Ok(Ok(Ok(response))) if response.status().is_success() => {}
        Ok(Ok(Ok(response))) => log_http_error(index, "NDJSON HTTP", response).await,
        Ok(Ok(Err(error))) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON upload failed: {error}")
        }
        Ok(Err(error)) => {
            eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON task failed: {error}")
        }
        Err(_) => eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON close timed out"),
    }
    let _ = done.send(());
}

#[cfg(feature = "atof-streaming")]
async fn drain_closed(mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>) {
    while let Some(message) = rx.recv().await {
        match message {
            EndpointMessage::Flush(done) => {
                let _ = done.send(());
            }
            EndpointMessage::Close(done) => {
                let _ = done.send(());
                return;
            }
            EndpointMessage::Event(_) => {}
        }
    }
}

#[cfg(feature = "atof-streaming")]
fn endpoint_event_json(config: &AtofEndpointConfig, raw_json: String) -> String {
    match config.field_name_policy {
        AtofEndpointFieldNamePolicy::Preserve => raw_json,
        AtofEndpointFieldNamePolicy::ReplaceDots => replace_dotted_field_names(&raw_json),
    }
}

#[cfg(feature = "atof-streaming")]
fn replace_dotted_field_names(raw_json: &str) -> String {
    let Ok(mut value) = serde_json::from_str::<Json>(raw_json) else {
        return raw_json.to_string();
    };
    replace_dotted_value_keys(&mut value);
    serde_json::to_string(&value).unwrap_or_else(|_| raw_json.to_string())
}

#[cfg(feature = "atof-streaming")]
fn replace_dotted_value_keys(value: &mut Json) {
    match value {
        Json::Object(object) => replace_dotted_object_keys(object),
        Json::Array(items) => {
            for item in items {
                replace_dotted_value_keys(item);
            }
        }
        _ => {}
    }
}

#[cfg(feature = "atof-streaming")]
fn replace_dotted_object_keys(object: &mut serde_json::Map<String, Json>) {
    let mut old = std::mem::take(object)
        .into_iter()
        .map(|(key, mut value)| {
            replace_dotted_value_keys(&mut value);
            (key, value)
        })
        .collect::<Vec<_>>();
    old.sort_by_key(|(key, _)| !key.contains('.'));

    for (key, value) in old {
        let sanitized_key = key.replace('.', "_");
        let final_key = collision_free_key(object, sanitized_key);
        object.insert(final_key, value);
    }
}

#[cfg(feature = "atof-streaming")]
fn collision_free_key(object: &serde_json::Map<String, Json>, key: String) -> String {
    if !object.contains_key(&key) {
        return key;
    }
    for suffix in 2.. {
        let candidate = format!("{key}_{suffix}");
        if !object.contains_key(&candidate) {
            return candidate;
        }
    }
    unreachable!("unbounded suffix search must find a key")
}

#[cfg(feature = "atof-streaming")]
async fn log_http_error(index: usize, label: &str, response: reqwest::Response) {
    let status = response.status();
    match response.text().await {
        Ok(body) if !body.trim().is_empty() => eprintln!(
            "nemo_relay: ATOF endpoint[{index}] {label} status {status}: {}",
            truncate_log_body(&body)
        ),
        Ok(_) => eprintln!("nemo_relay: ATOF endpoint[{index}] {label} status {status}"),
        Err(error) => eprintln!(
            "nemo_relay: ATOF endpoint[{index}] {label} status {status}; failed to read response body: {error}"
        ),
    }
}

#[cfg(feature = "atof-streaming")]
fn truncate_log_body(body: &str) -> String {
    const LIMIT: usize = 1024;
    let trimmed = body.trim();
    if trimmed.chars().count() <= LIMIT {
        return trimmed.to_string();
    }
    let mut truncated = trimmed.chars().take(LIMIT).collect::<String>();
    truncated.push_str("... <truncated>");
    truncated
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "../../tests/unit/observability/atof_tests.rs"]
mod tests;