//! Controller statistics used for reporting and flow control.
//!
//! # Design
//!
//! We use a lock-free design to avoid blocking the datapath.  All
//! performance counters are implemented as atomics.  However, updates
//! across multiple counters are not atomic.  More importantly, performance
//! counters can be slightly out of sync with the actual state of the
//! pipeline.  Such discrepancies are bounded and do _not_ accumulate
//! over time.
//!
//! # Example
//!
//! Consider the following interleaving of three threads: transport
//! endpoint 1, transport endpoint 2, and circuit thread:
// The quotes below are needed because of
// https://github.com/rust-lang/rustfmt/issues/4210
//! ```text
//! | Thread        | Action                              | `buffered_input_records` | actual # of buffered records |
//! | ------------- | ----------------------------------- | :----------------------: | :--------------------------: |
//! | 1. endpoint 1 | enqueues 10 records                 |            10            |              10              |
//! | 2. circuit:   | `ControllerStatus::consume_counters`|             0            |              10              |
//! | 3. endpoint 2 | enqueue 10 records                  |            10            |              20              |
//! | 4. circuit    | circuit.step()                      |            10            |              0               |
//! ```
//!
//! Here endpoint 2 buffers additional records after the
//! `ControllerStatus::buffered_input_records` counter is reset to zero.  As a
//! result, all 20 records enqueued by both endpoints are processed
//! by the circuit, but the counter shows that 10 records are still
//! pending.

use super::{EndpointId, InputEndpointConfig, OutputEndpointConfig};
use crate::{
    PipelineState,
    controller::{
        TransactionInfo, TransactionState,
        checkpoint::{CheckpointInputEndpointMetrics, CheckpointOutputEndpointMetrics},
        journal::{InputChecksums, InputLog},
    },
};
use anyhow::Error as AnyError;
use anyhow::anyhow;
use atomic::Atomic;
use base64::{Engine, prelude::BASE64_URL_SAFE_NO_PAD};
use chrono::{DateTime, Utc};
use cpu_time::ProcessTime;
use crossbeam::sync::Unparker;
use dbsp::{samply::SamplyEvent, utils::process_rss_bytes};
use feldera_adapterlib::{
    errors::journal::ControllerError,
    format::{BufferSize, ParseError},
    metrics::ConnectorMetrics,
    transport::{InputReader, Resume, Step, Watermark},
};
use feldera_storage::histogram::SlidingHistogram;
use feldera_types::{
    adapter_stats::{
        CompletedWatermark, ConnectorError, ConnectorHealth, ExternalInputEndpointMetrics,
        ExternalInputEndpointStatus, ExternalOutputEndpointMetrics, ExternalOutputEndpointStatus,
        ShortEndpointConfig,
    },
    config::{FtModel, PipelineConfig},
    coordination::Completion,
    memory_pressure::MemoryPressure,
    suspend::SuspendError,
    time_series::SampleStatistics,
    transaction::CommitProgressSummary,
};
use parking_lot::{RwLock, RwLockReadGuard};
use serde::{Deserialize, Serialize};
use size_of::HumanBytes;
use std::{
    collections::{BTreeMap, BTreeSet, VecDeque},
    fmt::Display,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};
use tokio::sync::{broadcast, watch};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use utoipa::ToSchema;

/// The number of the most recent errors stored for each endpoint as part of the endpoint status.
/// There are separate counters for parse and transport errors and for every tag.
const MAX_CONNECTOR_ERRORS: usize = 100;

/// Completion token.
///
/// A completion token associated with an endpoint identifies a position in the endpoint's
/// input stream and can be used to check whether all inputs up to this position have been
/// fully processed by the pipeline, including output changes getting processed by all
/// output connectors.
#[derive(Debug, ToSchema)]
pub struct CompletionToken {
    /// Pipeline incarnation uuid. Used to reject tokens generated by previous incarnations
    /// of the pipeline (before suspend, crash, or shutdown).
    incarnation: Uuid,

    /// Endpoint id.
    endpoint_id: EndpointId,

    /// Position the input stream that this token is associated with.
    count: u64,
}

#[derive(Serialize, Deserialize)]
struct ApiCompletionToken {
    u: Uuid,
    e: EndpointId,
    c: u64,
}

impl CompletionToken {
    pub fn new(incarnation: Uuid, endpoint_id: EndpointId, count: u64) -> Self {
        Self {
            incarnation,
            endpoint_id,
            count,
        }
    }

    /// Encode token as a URL-safe based64 string.
    pub fn encode(&self) -> String {
        let api = ApiCompletionToken {
            u: self.incarnation,
            e: self.endpoint_id,
            c: self.count,
        };
        let json = serde_json::to_string(&api).unwrap();
        BASE64_URL_SAFE_NO_PAD.encode(json)
    }

    /// Decode token from a URL-safe base64 string.
    pub fn decode(token: &str) -> Result<Self, anyhow::Error> {
        let json = BASE64_URL_SAFE_NO_PAD
            .decode(token)
            .map_err(|e| anyhow!("completion token is not a valid base64 string: {e}"))?;
        match serde_json_path_to_error::from_slice::<ApiCompletionToken>(&json) {
            Ok(api) => Ok(Self {
                incarnation: api.u,
                endpoint_id: api.e,
                count: api.c,
            }),
            Err(e) => Err(anyhow!("error parsing completion token: {e}")),
        }
    }
}

// Keep in sync with feldera_types::ExternalGlobalControllerMetrics
#[derive(Default)]
pub struct GlobalControllerMetrics {
    /// State of the pipeline: running, paused, or terminating.
    state: Atomic<PipelineState>,

    /// The pipeline has been resumed from a checkpoint and is currently bootstrapping
    /// new and modified views.
    bootstrap_in_progress: AtomicBool,

    /// Transaction commit progress, if a transaction is committing.
    pub commit_progress: Mutex<Option<CommitProgressSummary>>,

    /// Time at which the pipeline process started, in seconds since the epoch.
    pub start_time: DateTime<Utc>,

    /// Uniquely identifies the pipeline process that started at `start_time`.
    ///
    /// This is a v7 UUID, meaning that they will be ordered by their creation
    /// time (modulo clock skew).
    pub incarnation_uuid: Uuid,

    /// Time at which the pipeline process from which we resumed started, in
    /// seconds since the epoch.
    ///
    /// If this pipeline process was not started from a checkpoint, this is the
    /// same as `start_time`; otherwise, it is earlier.
    pub initial_start_time: DateTime<Utc>,

    /// Current storage usage in bytes.
    pub storage_bytes: AtomicU64,

    /// Storage usage integrated over time, in megabytes * seconds.
    pub storage_mb_secs: AtomicU64,

    /// Time elapsed while the pipeline is executing a step, multiplied by the
    /// number of foreground and background threads, in milliseconds.
    pub runtime_elapsed_msecs: AtomicU64,

    /// Total number of records currently buffered by all endpoints.
    pub buffered_input_records: AtomicU64,

    /// Total number of bytes currently buffered by all endpoints.
    pub buffered_input_bytes: AtomicU64,

    /// Total number of records received from all endpoints.
    pub total_input_records: AtomicU64,

    /// Total number of bytes received from all endpoints.
    pub total_input_bytes: AtomicU64,

    /// Total number of records pushed to the circuit.
    ///
    /// Equal to `total_input_records - buffered_input_records` in steady state,
    /// but can briefly get out of sync between updates to these two counters.
    /// We track this value for internal use, for when we need an accurate value
    /// and don't want to worry about races, but don't serialize it to avoid
    /// confusing the user with too many related values.
    pub total_circuit_input_records: AtomicU64,

    /// Total number of bytes pushed to the circuit.
    ///
    /// Equal to `total_input_bytes - buffered_input_bytes` in steady state.
    pub total_circuit_input_bytes: AtomicU64,

    /// Total number of input records processed by the DBSP engine.
    ///
    /// Note that some of the outputs produced for these input records
    /// may still be buffered by output connectors.
    /// Use `total_completed_records` for end-to-end progress tracking.
    pub total_processed_records: AtomicU64,

    /// Total bytes of input records processed by the DBSP engine.
    ///
    /// Note that some of the outputs produced for these input records
    /// may still be buffered by output connectors.
    pub total_processed_bytes: AtomicU64,

    /// Total number of input records processed to completion.
    ///
    /// A record is processed to completion if it has been processed by the DBSP engine and
    /// all outputs derived from it have been processed by all output connectors.
    pub total_completed_records: AtomicU64,

    /// If the pipeline is stalled because one or more output connectors' output
    /// buffers are full, this is the time at which the stall began.
    pub output_stall_start: Mutex<Option<Instant>>,

    /// The amount of time the pipeline has stalled due to output connectors,
    /// excluding any current stall in `output_stall_start`.
    pub accumulated_output_stall: Mutex<Duration>,

    /// Number of steps that have been initiated.
    ///
    /// # Interpretation
    ///
    /// This is a count, not a step number.  If `total_initiated_steps` is 0, no
    /// steps have been initiated.  If `total_initiated_steps > 0`, then step
    /// `total_initiated_steps - 1` has been started and all steps previous to
    /// that have been completely processed by the circuit.
    pub total_initiated_steps: Atomic<Step>,

    /// Number of steps whose input records have been processed to completion.
    ///
    /// A record is processed to completion if it has been processed by the DBSP engine and
    /// all outputs derived from it have been processed by all output connectors.
    ///
    /// # Interpretation
    ///
    /// This is a count, not a step number.  If `total_completed_steps` is 0, no
    /// steps have been processed to completion.  If `total_completed_steps >
    /// 0`, then the last step whose input records have been processed to
    /// completion is `total_completed_steps - 1`. A record that was ingested
    /// when `total_initiated_steps` was `n` is fully processed when
    /// `total_completed_steps >= n`.
    pub total_completed_steps: Atomic<Step>,

    /// Forces the controller to perform a step regardless of the state of
    /// input buffers.
    pub step_requested: AtomicBool,
}

impl GlobalControllerMetrics {
    fn new(
        processed_records: u64,
        initial_start_time: Option<DateTime<Utc>>,
        incarnation_uuid: Uuid,
    ) -> Self {
        let start_time = Utc::now();
        let initial_start_time = initial_start_time.unwrap_or(start_time);
        Self {
            state: Atomic::new(PipelineState::Paused),
            bootstrap_in_progress: AtomicBool::new(false),
            commit_progress: Mutex::new(None),
            start_time,
            incarnation_uuid,
            initial_start_time,
            storage_bytes: AtomicU64::new(0),
            storage_mb_secs: AtomicU64::new(0),
            runtime_elapsed_msecs: AtomicU64::new(0),
            buffered_input_records: AtomicU64::new(0),
            buffered_input_bytes: AtomicU64::new(0),
            total_input_records: AtomicU64::new(processed_records),
            total_input_bytes: AtomicU64::new(0),
            total_circuit_input_records: AtomicU64::new(processed_records),
            total_circuit_input_bytes: AtomicU64::new(0),
            total_processed_records: AtomicU64::new(processed_records),
            total_processed_bytes: AtomicU64::new(0),
            total_completed_records: AtomicU64::new(processed_records),
            output_stall_start: Mutex::new(None),
            accumulated_output_stall: Mutex::new(Duration::ZERO),
            step_requested: AtomicBool::new(false),
            total_initiated_steps: Atomic::new(0),
            total_completed_steps: Atomic::new(0),
        }
    }

    pub fn get_state(&self) -> PipelineState {
        self.state.load(Ordering::Acquire)
    }

    fn input_batch(&self, amt: BufferSize) -> u64 {
        self.total_input_bytes
            .fetch_add(amt.bytes as u64, Ordering::Relaxed);
        self.buffered_input_bytes
            .fetch_add(amt.bytes as u64, Ordering::Relaxed);
        self.total_input_records
            .fetch_add(amt.records as u64, Ordering::AcqRel);
        self.buffered_input_records
            .fetch_add(amt.records as u64, Ordering::AcqRel)
    }

    /// `amt` have been pushed to the circuit (but not yet processed, i.e., the
    /// circuit hasn't performed a step yet).
    pub(crate) fn consume_buffered_inputs(&self, amt: BufferSize) {
        self.buffered_input_records
            .fetch_sub(amt.records as u64, Ordering::Release);
        self.total_circuit_input_records
            .fetch_add(amt.records as u64, Ordering::AcqRel);
        self.buffered_input_bytes
            .fetch_sub(amt.bytes as u64, Ordering::Relaxed);
        self.total_circuit_input_bytes
            .fetch_add(amt.bytes as u64, Ordering::Relaxed);
    }

    pub(crate) fn processed_data(&self, amt: BufferSize) -> u64 {
        self.total_processed_bytes
            .fetch_add(amt.bytes as u64, Ordering::AcqRel);
        self.total_processed_records
            .fetch_add(amt.records as u64, Ordering::AcqRel)
            + amt.records as u64
    }

    pub fn num_buffered_input_records(&self) -> u64 {
        self.buffered_input_records.load(Ordering::Acquire)
    }

    pub fn num_buffered_input_bytes(&self) -> u64 {
        self.buffered_input_bytes.load(Ordering::Acquire)
    }

    pub fn num_total_input_records(&self) -> u64 {
        self.total_input_records.load(Ordering::Acquire)
    }

    pub fn num_total_input_bytes(&self) -> u64 {
        self.total_input_bytes.load(Ordering::Acquire)
    }

    pub fn num_total_circuit_input_records(&self) -> u64 {
        self.total_circuit_input_records.load(Ordering::Acquire)
    }

    pub fn num_total_processed_records(&self) -> u64 {
        self.total_processed_records.load(Ordering::Acquire)
    }

    pub fn total_initiated_steps(&self) -> Step {
        self.total_initiated_steps.load(Ordering::Acquire)
    }

    pub fn total_completed_steps(&self) -> Step {
        self.total_completed_steps.load(Ordering::Acquire)
    }

    pub fn num_total_processed_bytes(&self) -> u64 {
        self.total_processed_bytes.load(Ordering::Acquire)
    }

    fn unset_step_requested(&self) -> bool {
        self.step_requested.swap(false, Ordering::Acquire)
    }

    fn bootstrap_in_progress(&self) -> bool {
        self.bootstrap_in_progress.load(Ordering::Acquire)
    }

    fn set_bootstrap_in_progress(&self, bootstrap_in_progress: bool) {
        self.bootstrap_in_progress
            .store(bootstrap_in_progress, Ordering::Release);
    }

    fn set_step_requested(&self) -> bool {
        self.step_requested.swap(true, Ordering::AcqRel)
    }

    pub fn set_commit_progress(&self, commit_progress: Option<CommitProgressSummary>) {
        *self.commit_progress.lock().unwrap() = commit_progress;
    }

    pub fn update_output_stall_start(&self, stalled: bool) {
        let mut output_stall_start = self.output_stall_start.lock().unwrap();
        if stalled != output_stall_start.is_some() {
            if let Some(start) = &*output_stall_start {
                *self.accumulated_output_stall.lock().unwrap() += start.elapsed();
            }
            *output_stall_start = stalled.then(Instant::now);
        }
    }

    pub fn current_output_stall_duration(&self) -> Duration {
        self.output_stall_start
            .lock()
            .unwrap()
            .map(|instant| instant.elapsed())
            .unwrap_or_default()
    }

    pub fn total_output_stall_duration(&self) -> Duration {
        self.current_output_stall_duration() + *self.accumulated_output_stall.lock().unwrap()
    }
}

// `ShardedLock` is a read/write lock optimized for fast reads.
// Write access is only required when adding or removing an endpoint.
// Regular stats updates only require a read lock thanks to the use of
// atomics.
type InputsStatus = RwLock<BTreeMap<EndpointId, InputEndpointStatus>>;
type OutputsStatus = RwLock<BTreeMap<EndpointId, OutputEndpointStatus>>;

/// Controller statistics.
///
/// Keep in sync with feldera_types::ExternalControllerStatus
pub struct ControllerStatus {
    /// Global controller configuration.
    pub pipeline_config: PipelineConfig,

    /// Global controller metrics.
    pub global_metrics: GlobalControllerMetrics,

    /// Statistics time series.
    pub time_series: Mutex<VecDeque<SampleStatistics>>,

    /// Broadcast channel for notifying subscribers about new time series data.
    pub time_series_notifier: broadcast::Sender<SampleStatistics>,

    /// Watch channel for notifying subscribers about [Completion] updates.
    pub completion_notifier: watch::Sender<Completion>,

    /// Input endpoint configs and metrics.
    pub(crate) inputs: InputsStatus,

    /// Output endpoint configs and metrics.
    outputs: OutputsStatus,
}

impl ControllerStatus {
    pub fn new(
        pipeline_config: PipelineConfig,
        processed_records: u64,
        initial_start_time: Option<DateTime<Utc>>,
        incarnation_uuid: Uuid,
    ) -> Self {
        let (time_series_notifier, _) = broadcast::channel(1024); // Buffer for up to 1024 time series updates
        let (completion_notifier, _) = watch::channel(Completion::default());
        Self {
            pipeline_config,
            global_metrics: GlobalControllerMetrics::new(
                processed_records,
                initial_start_time,
                incarnation_uuid,
            ),
            time_series: Mutex::new(VecDeque::with_capacity(60)),
            time_series_notifier,
            completion_notifier,
            inputs: RwLock::new(BTreeMap::new()),
            outputs: RwLock::new(BTreeMap::new()),
        }
    }

    pub fn state(&self) -> PipelineState {
        self.global_metrics.state.load(Ordering::Relaxed)
    }

    /// Set the state to `desired`.
    ///
    /// Setting the state to [PipelineState::Terminated] is permanent; the state
    /// can't be re-set to any other state after that.
    pub fn set_state(&self, desired: PipelineState) {
        if desired == PipelineState::Terminated {
            self.global_metrics.state.store(desired, Ordering::Relaxed);
        } else {
            let mut current = self.state();
            while current != desired && current != PipelineState::Terminated {
                match self.global_metrics.state.compare_exchange(
                    current,
                    desired,
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => break,
                    Err(changed) => current = changed,
                }
            }
        }
    }

    /// Lookup input endpoint by name.
    pub fn input_endpoint_id_by_name(
        &self,
        endpoint_name: &str,
    ) -> Result<EndpointId, ControllerError> {
        let inputs = self.input_status();

        for (endpoint_id, descr) in inputs.iter() {
            if descr.endpoint_name == endpoint_name {
                return Ok(*endpoint_id);
            }
        }

        Err(ControllerError::unknown_input_endpoint(endpoint_name))
    }

    /// Lookup output endpoint by name.
    pub fn output_endpoint_id_by_name(
        &self,
        endpoint_name: &str,
    ) -> Result<EndpointId, ControllerError> {
        let outputs = self.output_status();

        for (endpoint_id, descr) in outputs.iter() {
            if descr.endpoint_name == endpoint_name {
                return Ok(*endpoint_id);
            }
        }

        Err(ControllerError::unknown_output_endpoint(endpoint_name))
    }

    /// Sets `endpoint`'s paused state to `paused` and returns whether it was
    /// previously paused.
    pub fn set_input_endpoint_paused(&self, endpoint: &EndpointId, paused: bool) -> Option<bool> {
        Some(
            self.input_status()
                .get(endpoint)?
                .set_paused_by_user(paused),
        )
    }

    /// Pauses `endpoint` and returns whether the endpoint was previously
    /// paused.
    pub fn pause_input_endpoint(&self, endpoint: &EndpointId) -> Option<bool> {
        self.set_input_endpoint_paused(endpoint, true)
    }

    /// Unpauses `endpoint` and returns whether the endpoint was previously
    /// paused.
    pub fn start_input_endpoint(&self, endpoint: &EndpointId) -> Option<bool> {
        self.set_input_endpoint_paused(endpoint, false)
    }

    pub fn is_input_endpoint_paused(&self, endpoint: &EndpointId) -> Option<bool> {
        Some(self.input_status().get(endpoint)?.is_paused_by_user())
    }

    /// Invoked when one of the input endpoints has finished processing
    /// to check if any paused endpoints have all their 'start_after'
    /// dependencies satisfied and can be started.
    pub fn start_dependencies(&self, backpressure_thread_unparker: &Unparker) {
        let mut endpoints_to_start: Vec<EndpointId> = Vec::new();

        let inputs = self.input_status();

        let mut unfinished_labels = BTreeSet::new();
        for input in inputs.values() {
            if !input.finished() {
                unfinished_labels.extend(input.config.connector_config.labels.iter());
            }
        }

        for (endpoint_id, input) in inputs.iter() {
            if !input.is_paused_by_user() {
                continue;
            }
            if let Some(start_after) = &input.config.connector_config.start_after
                && start_after
                    .iter()
                    .all(|label| !unfinished_labels.contains(label))
            {
                info!(
                    "starting endpoint {}, whose 'start_after' dependencies have been satisfied",
                    input.endpoint_name
                );
                endpoints_to_start.push(*endpoint_id);
            }
        }

        drop(inputs);

        for endpoint_id in endpoints_to_start.iter() {
            self.start_input_endpoint(endpoint_id);
        }

        if !endpoints_to_start.is_empty() {
            backpressure_thread_unparker.unpark();
        }
    }

    pub fn remove_input(&self, endpoint_id: &EndpointId) -> Option<InputEndpointStatus> {
        self.inputs.write().remove(endpoint_id)
    }

    pub fn remove_output(&self, endpoint_id: &EndpointId) {
        self.outputs.write().remove(endpoint_id);
    }

    /// Initialize stats for a new output endpoint.
    pub fn add_output(
        &self,
        endpoint_id: &EndpointId,
        endpoint_name: &str,
        config: &OutputEndpointConfig,
        initial_statistics: Option<&CheckpointOutputEndpointMetrics>,
    ) {
        // Initialize the `total_processed_input_records` counter on the new endpoint to `total_processed_records`:
        // logically the new endpoint is up to speed with the outputs produced by the pipeline so far and only needs to
        // process any future outputs.
        let total_processed_records = self
            .global_metrics
            .total_processed_records
            .load(Ordering::Acquire);
        self.outputs.write().insert(
            *endpoint_id,
            OutputEndpointStatus::new(
                endpoint_name,
                config,
                total_processed_records,
                initial_statistics,
            ),
        );
    }

    /// Total number of records currently buffered by all input endpoints.
    pub fn num_buffered_input_records(&self) -> u64 {
        self.global_metrics.num_buffered_input_records()
    }

    /// Total number of bytes currently buffered by all input endpoints.
    pub fn num_buffered_input_bytes(&self) -> u64 {
        self.global_metrics.num_buffered_input_bytes()
    }

    /// Total number of records received from all input endpoints.
    pub fn num_total_input_records(&self) -> u64 {
        self.global_metrics.num_total_input_records()
    }

    /// Total number of bytes received from all input endpoints.
    pub fn num_total_input_bytes(&self) -> u64 {
        self.global_metrics.num_total_input_bytes()
    }

    pub fn num_total_circuit_input_records(&self) -> u64 {
        self.global_metrics.num_total_circuit_input_records()
    }

    pub fn num_total_processed_records(&self) -> u64 {
        self.global_metrics.num_total_processed_records()
    }

    pub fn num_total_processed_bytes(&self) -> u64 {
        self.global_metrics.num_total_processed_bytes()
    }

    pub fn processed_data(&self, amt: BufferSize) -> u64 {
        let new_processed_records = self.global_metrics.processed_data(amt);
        let ts = Utc::now();
        self.inputs.read().values().for_each(|input| {
            input
                .completed_frontier
                .update_processed(new_processed_records, ts);
        });
        new_processed_records
    }

    pub fn unset_step_requested(&self) -> bool {
        self.global_metrics.unset_step_requested()
    }

    pub fn bootstrap_in_progress(&self) -> bool {
        self.global_metrics.bootstrap_in_progress()
    }

    pub fn set_bootstrap_in_progress(&self, bootstrap_in_progress: bool) {
        self.global_metrics
            .set_bootstrap_in_progress(bootstrap_in_progress);
    }

    pub fn request_step(&self, circuit_thread_unparker: &Unparker) {
        let old = self.global_metrics.set_step_requested();
        if !old {
            circuit_thread_unparker.unpark();
        }
    }

    /// Input endpoint stats.
    pub fn input_status(&self) -> RwLockReadGuard<'_, BTreeMap<EndpointId, InputEndpointStatus>> {
        self.inputs.read_recursive()
    }

    /// Register connector-specific metrics for an input endpoint.
    pub fn set_custom_metrics(&self, endpoint_id: EndpointId, metrics: Arc<dyn ConnectorMetrics>) {
        if let Some(status) = self.inputs.write().get_mut(&endpoint_id) {
            status.custom_metrics = Some(metrics);
        }
    }

    /// Output endpoint stats.
    pub fn output_status(&self) -> RwLockReadGuard<'_, BTreeMap<EndpointId, OutputEndpointStatus>> {
        self.outputs.read_recursive()
    }

    /// Number of records buffered by the endpoint or 0 if the endpoint
    /// doesn't exist (the latter is possible if the endpoint is being
    /// destroyed).
    pub fn num_input_endpoint_buffered_records(&self, endpoint_id: &EndpointId) -> u64 {
        match &self.input_status().get(endpoint_id) {
            None => 0,
            Some(endpoint_stats) => endpoint_stats
                .metrics
                .buffered_records
                .load(Ordering::Acquire),
        }
    }

    pub fn num_total_completed_records(&self) -> u64 {
        self.global_metrics
            .total_completed_records
            .load(Ordering::Acquire)
    }

    /// Get information about outstanding output records at each output endpoint.
    ///
    /// Returns a map of output endpoint names to the number of queued records, queued batches, and buffered records.
    pub fn pending_output_records(&self) -> BTreeMap<String, (u64, u64, u64)> {
        let mut pending_output_records = BTreeMap::new();
        for endpoint_stats in self.output_status().values() {
            pending_output_records.insert(
                endpoint_stats.endpoint_name.clone(),
                (
                    endpoint_stats
                        .metrics
                        .queued_records
                        .load(Ordering::Acquire),
                    endpoint_stats
                        .metrics
                        .queued_batches
                        .load(Ordering::Acquire),
                    endpoint_stats
                        .metrics
                        .buffered_records
                        .load(Ordering::Acquire),
                ),
            );
        }

        pending_output_records
    }

    /// Update `global_metrics.total_completed_records` and `global_metrics.total_completed_steps`.
    ///
    /// Must be invoked any time this metric can change, i.e., after every output
    /// produced by an output connector as well as after each step.
    ///
    /// Computes `total_completed_records` as the minimum `total_processed_input_records` across
    /// all output connectors. In case there are no output connectors attached to the pipeline,
    /// returns `total_processed_records`.
    ///
    /// Also updates watermark trackers in input endpoints.
    pub fn update_total_completed_records(&self) {
        // Compute new `total_completed_records` atomically, protected by the output status lock.
        // However we don't want to call `watermarks_update_completed` while holding the lock, as that
        // will take the input status lock. This is not necessarily a bug, but it will complicate
        // reasoning about potential deadlocks.
        //
        // So we call `watermarks_update_completed` outside the atomic section, but use compare_exchange to
        // ensure that the biggest value wins among multiple writers (`total_completed_records` increases
        // monotonically).

        let output_status = self.output_status();

        let mut total_completed_records = self.global_metrics.num_total_processed_records();
        let mut total_completed_steps = self.global_metrics.total_initiated_steps();
        for output_ep in output_status.values() {
            total_completed_records =
                total_completed_records.min(output_ep.num_total_processed_input_records());
            total_completed_steps =
                total_completed_steps.min(output_ep.num_total_processed_steps());
        }

        drop(output_status);

        // Use `fetch_max` to ensure that the biggest value wins among multiple writers.
        let old_records = self
            .global_metrics
            .total_completed_records
            .fetch_max(total_completed_records, Ordering::SeqCst);
        if total_completed_records >= old_records {
            self.watermarks_update_completed(total_completed_records);
        }

        let old_steps = self
            .global_metrics
            .total_completed_steps
            .fetch_max(total_completed_steps, Ordering::SeqCst);
        if total_completed_steps >= old_steps {
            self.completion_notifier.send_if_modified(|state| {
                if state.total_completed_steps != total_completed_steps {
                    state.total_completed_steps = total_completed_steps;
                    true
                } else {
                    false
                }
            });
        }
    }

    /// Notify all watermark trackers about a new value for `total_completed_records`.
    pub fn watermarks_update_completed(&self, total_completed_records: u64) {
        let ts = Utc::now();
        self.inputs.read().values().for_each(|input| {
            input
                .completed_frontier
                .update_completed(total_completed_records, ts, &input.metrics);
        })
    }

    /// Generate completion token for an endpoint based on the current number of records ingested by this endpoint.
    pub fn completion_token(
        &self,
        endpoint_name: &str,
        endpoint_id: EndpointId,
    ) -> Result<CompletionToken, ControllerError> {
        if let Some(endpoint_stats) = self.input_status().get(&endpoint_id) {
            let num_input_records = endpoint_stats.completion_token(
                self.global_metrics.total_initiated_steps(),
                self.global_metrics.total_completed_steps(),
            );
            Ok(CompletionToken::new(
                self.global_metrics.incarnation_uuid,
                endpoint_id,
                num_input_records,
            ))
        } else {
            Err(ControllerError::unknown_input_endpoint(endpoint_name))
        }
    }

    /// Check completion status of `completion_token`.
    ///
    /// Returns `None` if the token is definitely incomplete, `Some(step)` if it
    /// is complete once `GlobalControllerMetrics::total_completed_steps`
    /// reaches `step`, or an error if the token is invalid.
    pub fn completion_status(
        &self,
        completion_token: &CompletionToken,
    ) -> Result<Option<Step>, ControllerError> {
        if completion_token.incarnation != self.global_metrics.incarnation_uuid {
            return Err(ControllerError::pipeline_restarted(&format!(
                "Completion token was created by a previous incarnation of the pipeline (incarnation uuid: {}) and is not valid for the current incarnation ({}). This indicates that the pipeline was suspended and resumed from a checkpoint or restarted after a failure.",
                &completion_token.incarnation, &self.global_metrics.incarnation_uuid
            )));
        }

        if let Some(endpoint_stats) = self.input_status().get(&completion_token.endpoint_id) {
            Ok(endpoint_stats.completion_status(completion_token.count))
        } else {
            Err(ControllerError::unknown_endpoint_in_completion_token(
                completion_token.endpoint_id,
            ))
        }
    }

    /// Update the global counters after receiving a new input batch.
    ///
    /// This method is used for inserts that don't belong to an endpoint, e.g.,
    /// happen by executing an ad-hoc INSERT query.
    pub(super) fn input_batch_global(&self, amt: BufferSize, circuit_thread_unparker: &Unparker) {
        let num_records = amt.records as u64;
        // Increment buffered_records; unpark circuit thread once
        // `min_batch_size_records` is exceeded.
        let old = self.global_metrics.input_batch(amt);

        if old == 0
            || (old <= self.pipeline_config.global.min_batch_size_records
                && old + num_records > self.pipeline_config.global.min_batch_size_records)
        {
            circuit_thread_unparker.unpark();
        }
    }

    /// Update counters after receiving a new input batch.
    ///
    /// # Arguments
    ///
    /// * `endpoint_id` - id of the input endpoint.
    /// * `num_bytes` - number of bytes received.
    /// * `num_records` - number of records in the deserialized batch.
    /// * `global_config` - global controller config.
    /// * `circuit_thread_unparker` - unparker used to wake up the circuit
    ///   thread if the total number of buffered records exceeds
    ///   `min_batch_size_records`.
    /// * `backpressure_thread_unparker` - unparker used to wake up the
    ///   backpressure thread if the endpoint is full.
    pub(super) fn input_batch_from_endpoint(
        &self,
        endpoint_id: EndpointId,
        amt: BufferSize,
        backpressure_thread_unparker: &Unparker,
    ) {
        // There is a potential race condition if the endpoint is currently
        // being removed. In this case, it's safe to ignore this operation.
        if !amt.is_empty() {
            let inputs = self.input_status();
            if let Some(endpoint_stats) = inputs.get(&endpoint_id) {
                let old = endpoint_stats.add_buffered(amt);
                let threshold = endpoint_stats.config.connector_config.max_queued_records;
                if old < threshold && old + amt.records as u64 >= threshold {
                    backpressure_thread_unparker.unpark();
                }
            }
        }
    }

    /// Update counters after receiving an end-of-input event on an input
    /// endpoint.
    ///
    /// # Arguments
    ///
    /// * `endpoint_id` - id of the input endpoint.
    /// * `num_records` - number of records returned by `Parser::eoi`.
    /// * `circuit_thread_unparker` - unparker used to wake up the circuit
    ///   thread if the total number of buffered records exceeds
    ///   `min_batch_size_records`.
    pub fn eoi(
        &self,
        endpoint_id: EndpointId,
        circuit_thread_unparker: &Unparker,
        backpressure_thread_unparker: &Unparker,
    ) {
        if self.global_metrics.num_buffered_input_records() == 0 {
            circuit_thread_unparker.unpark();
        }

        let mut finished = false;

        // Update endpoint counters, no need to wake up the backpressure thread since we
        // won't see any more inputs from this endpoint.
        let inputs = self.input_status();
        if let Some(endpoint_stats) = inputs.get(&endpoint_id) {
            endpoint_stats.eoi();
            finished = endpoint_stats.finished();
        };

        drop(inputs);

        // Check if we can start any paused endpoints waiting for this endpoint to finish.
        // It's possible that the endpoint is at the end of input, but not all of its updates
        // have been processed yet. In this case, `finished` will be `false`, and we will call
        // `start_dependencies` when the endpoint is finished in the `completed()` method below.
        if finished {
            self.start_dependencies(backpressure_thread_unparker);
        }
    }

    /// Endpoint pushed additional records to the circuit.
    pub fn extended(
        &self,
        endpoint_id: EndpointId,
        step_results: StepResults,
        watermarks: Vec<Watermark>,
        backpressure_thread_unparker: &Unparker,
    ) {
        let inputs = self.input_status();
        self.global_metrics
            .consume_buffered_inputs(step_results.amt);

        let mut finished = false;

        if let Some(endpoint_stats) = inputs.get(&endpoint_id) {
            endpoint_stats.extended(
                step_results,
                watermarks,
                self.num_total_circuit_input_records(),
                self.global_metrics.total_initiated_steps(),
                self.global_metrics.total_completed_steps(),
            );
            finished = endpoint_stats.finished();
        };

        drop(inputs);

        // Check if we can start any paused endpoints waiting for this endpoint to finish.
        if finished {
            self.start_dependencies(backpressure_thread_unparker);
        }
    }

    pub fn enqueue_batch(&self, endpoint_id: EndpointId, num_records: usize) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.enqueue_batch(num_records);
        }
    }

    pub fn buffer_batch(
        &self,
        endpoint_id: EndpointId,
        num_records: usize,
        circuit_thread_unparker: &Unparker,
    ) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            let old = endpoint_stats.buffer_batch(num_records);
            if old - (num_records as u64)
                <= endpoint_stats.config.connector_config.max_queued_records
                && old >= endpoint_stats.config.connector_config.max_queued_records
            {
                circuit_thread_unparker.unpark();
            }
        };
    }

    /// `processed` is statistics about records processed by the circuit before
    /// this output batch was produced or `None` if the circuit is executing a
    /// transaction.
    pub fn output_batch(
        &self,
        endpoint_id: EndpointId,
        processed: Option<ProcessedRecords>,
        num_records: usize,
        circuit_thread_unparker: &Unparker,
    ) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            let threshold = endpoint_stats.config.connector_config.max_queued_records;

            let old = endpoint_stats.output_batch(processed, num_records);

            let new = old - (num_records as u64);
            if (old >= threshold && new <= threshold) || new == 0 {
                circuit_thread_unparker.unpark();
            }
        };
        self.update_total_completed_records();
    }

    pub fn update_output_memory(&self, endpoint_id: EndpointId, memory: usize) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats
                .metrics
                .memory
                .store(memory as u64, Ordering::Release);
        }
    }

    pub fn output_buffered_batches(&self, endpoint_id: EndpointId, processed: ProcessedRecords) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.output_buffered_batches(processed);
        }
        self.update_total_completed_records();
    }

    pub fn output_buffer(&self, endpoint_id: EndpointId, num_bytes: usize, num_records: usize) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.output_buffer(num_bytes, num_records);
        };
    }

    pub fn output_buffers_full(&self) -> bool {
        self.output_status()
            .values()
            .any(|endpoint_stats| endpoint_stats.is_buffer_full())
    }

    /// Returns the names of output connectors whose buffers are currently full.
    pub fn output_buffers_full_names(&self) -> Vec<String> {
        self.output_status()
            .values()
            .filter_map(|endpoint_stats| {
                if endpoint_stats.is_buffer_full() {
                    Some(endpoint_stats.endpoint_name.clone())
                } else {
                    None
                }
            })
            .collect()
    }

    pub fn parse_error(&self, endpoint_id: EndpointId, tag: Option<&str>, error: &ParseError) {
        if let Some(endpoint_stats) = self.input_status().get(&endpoint_id) {
            endpoint_stats.parse_error(tag, error);
        }
    }

    pub fn encode_error(&self, endpoint_id: EndpointId, tag: Option<&str>, error: &AnyError) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.encode_error(tag, error);
        }
    }

    pub fn input_transport_error(
        &self,
        endpoint_id: EndpointId,
        fatal: bool,
        tag: Option<&str>,
        error: &AnyError,
    ) {
        if let Some(endpoint_stats) = self.input_status().get(&endpoint_id) {
            endpoint_stats.transport_error(fatal, tag, error);
        }
    }

    pub fn output_transport_error(
        &self,
        endpoint_id: EndpointId,
        fatal: bool,
        tag: Option<&str>,
        error: &AnyError,
    ) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.transport_error(fatal, tag, error);
        }
    }

    pub fn update_input_connector_health(&self, endpoint_id: EndpointId, health: ConnectorHealth) {
        if let Some(endpoint_stats) = self.input_status().get(&endpoint_id) {
            endpoint_stats.update_health(health);
        }
    }

    pub fn update_output_connector_health(&self, endpoint_id: EndpointId, health: ConnectorHealth) {
        if let Some(endpoint_stats) = self.output_status().get(&endpoint_id) {
            endpoint_stats.update_health(health);
        }
    }

    /// True if the pipeline has processed all inputs to completion.
    pub fn pipeline_complete(&self) -> bool {
        // All input endpoints (if any) are at end of input.
        if !self
            .input_status()
            .values()
            .filter(|endpoint_stats| !endpoint_stats.endpoint_name.contains(".api-ingress-"))
            .all(|endpoint_stats| endpoint_stats.is_eoi())
        {
            return false;
        }

        // All received records have been processed by the circuit.
        let total_input_records = self.num_total_input_records();

        if self.num_total_processed_records() != total_input_records {
            return false;
        }

        // Outputs have been pushed to their respective transport endpoints.
        if !self.output_status().values().all(|endpoint_stats| {
            endpoint_stats.num_total_processed_input_records() == total_input_records
        }) {
            return false;
        }

        true
    }

    /// Convert from internal `ControllerStatus` to the public API type in `feldera_types`.
    ///
    /// This function ensures that the "duplicate" types in `feldera_types::adapter_stats`
    /// stay correlated with the original types and won't go out of sync.
    pub fn to_api_type(
        &self,
        suspend_error: Result<(), SuspendError>,
        pipeline_complete: bool,
        transaction_info: TransactionInfo,
        memory_pressure: MemoryPressure,
        memory_pressure_epoch: u64,
    ) -> feldera_types::adapter_stats::ExternalControllerStatus {
        use feldera_types::adapter_stats;

        let total_processed_records = self
            .global_metrics
            .total_processed_records
            .load(Ordering::Acquire);

        #[allow(clippy::manual_unwrap_or_default)]
        #[allow(clippy::manual_unwrap_or)]
        let global_metrics = adapter_stats::ExternalGlobalControllerMetrics {
            state: match self.global_metrics.state.load(Ordering::Acquire) {
                PipelineState::Paused => adapter_stats::PipelineState::Paused,
                PipelineState::Running => adapter_stats::PipelineState::Running,
                PipelineState::Terminated => adapter_stats::PipelineState::Terminated,
            },
            bootstrap_in_progress: self
                .global_metrics
                .bootstrap_in_progress
                .load(Ordering::Acquire),
            transaction_status: match transaction_info.transaction_state {
                TransactionState::None => adapter_stats::TransactionStatus::NoTransaction,
                TransactionState::Started { .. } => {
                    adapter_stats::TransactionStatus::TransactionInProgress
                }
                TransactionState::Committing { .. } => {
                    adapter_stats::TransactionStatus::CommitInProgress
                }
            },
            transaction_msecs: transaction_info
                .transaction_state
                .start_time()
                .and_then(|start| start.elapsed().as_millis().try_into().ok()),
            transaction_records: transaction_info
                .transaction_state
                .processed_records()
                .map(|n| total_processed_records - n),
            transaction_id: if let Some(tid) = transaction_info.transaction_state.tid() {
                // The current transaction ID.
                tid
            } else if let Some(tid) = transaction_info.initiators.transaction_id {
                // The transaction that will start when we execute a step.
                tid
            } else {
                // No transaction
                0
            },
            commit_progress: self.global_metrics.commit_progress.lock().unwrap().clone(),
            transaction_initiators: transaction_info.initiators.to_api_type(),
            start_time: self.global_metrics.start_time,
            incarnation_uuid: self.global_metrics.incarnation_uuid,
            initial_start_time: self.global_metrics.initial_start_time,
            storage_bytes: self.global_metrics.storage_bytes.load(Ordering::Relaxed),
            storage_mb_secs: self.global_metrics.storage_mb_secs.load(Ordering::Relaxed),
            uptime_msecs: (Utc::now() - self.global_metrics.start_time)
                .num_milliseconds()
                .try_into()
                .unwrap_or(0),
            rss_bytes: process_rss_bytes().unwrap_or_else(|| {
                error!("Failed to fetch process RSS");
                0
            }),
            memory_pressure,
            memory_pressure_epoch,
            cpu_msecs: match ProcessTime::try_now() {
                Ok(time) => time.as_duration().as_millis() as u64,
                Err(e) => {
                    error!("Failed to fetch process times: {e}");
                    0
                }
            },
            runtime_elapsed_msecs: self
                .global_metrics
                .runtime_elapsed_msecs
                .load(Ordering::Relaxed),
            buffered_input_records: self
                .global_metrics
                .buffered_input_records
                .load(Ordering::Acquire),
            buffered_input_bytes: self
                .global_metrics
                .buffered_input_bytes
                .load(Ordering::Acquire),
            total_input_records: self
                .global_metrics
                .total_input_records
                .load(Ordering::Acquire),
            total_input_bytes: self
                .global_metrics
                .total_input_bytes
                .load(Ordering::Acquire),
            total_processed_records,
            total_processed_bytes: self
                .global_metrics
                .total_processed_bytes
                .load(Ordering::Acquire),
            total_completed_records: self
                .global_metrics
                .total_completed_records
                .load(Ordering::Acquire),
            output_stall_msecs: self
                .global_metrics
                .current_output_stall_duration()
                .as_millis()
                .try_into()
                .unwrap_or(u64::MAX),
            total_initiated_steps: self.global_metrics.total_initiated_steps(),
            total_completed_steps: self.global_metrics.total_completed_steps(),
            pipeline_complete,
        };

        // Convert input endpoints and sort by endpoint_name to match serialize_inputs behavior
        let mut inputs: Vec<_> = self
            .input_status()
            .values()
            .map(|input| input.to_api_type(false))
            .collect();
        inputs.sort_by(|ep1, ep2| ep1.endpoint_name.cmp(&ep2.endpoint_name));

        // Convert output endpoints and sort by endpoint_name to match serialize_outputs behavior
        let mut outputs: Vec<_> = self
            .output_status()
            .values()
            .map(|output| output.to_api_type(false))
            .collect();
        outputs.sort_by(|ep1, ep2| ep1.endpoint_name.cmp(&ep2.endpoint_name));

        adapter_stats::ExternalControllerStatus {
            global_metrics,
            suspend_error: suspend_error.err(),
            inputs,
            outputs,
        }
    }
}

/// Keep in sync with feldera_types::ExternalInputEndpointMetrics
pub struct InputEndpointMetrics {
    /// Total bytes pushed to the endpoint since it was created.
    pub total_bytes: AtomicU64,

    /// Total records pushed to the endpoint since it was created.
    pub total_records: AtomicU64,

    /// Number of records currently buffered by the endpoint
    /// (not yet consumed by the circuit).
    pub buffered_records: AtomicU64,

    /// Number of bytes currently buffered by the endpoint
    /// (not yet consumed by the circuit).
    pub buffered_bytes: AtomicU64,

    /// Number of records pushed from the endpoint's queue to the circuit.
    ///
    /// Equal to `total_records - buffered_records` in steady state.
    pub circuit_input_records: AtomicU64,

    /// Number of bytes pushed from the endpoint's queue to the circuit.
    ///
    /// Equal to `total_bytes - buffered_bytes` in steady state.
    pub circuit_input_bytes: AtomicU64,

    pub num_transport_errors: AtomicU64,

    pub num_parse_errors: AtomicU64,

    pub end_of_input: AtomicBool,

    pub processing_latency_micros_histogram: Mutex<SlidingHistogram>,

    pub completion_latency_micros_histogram: Mutex<SlidingHistogram>,
}

impl Default for InputEndpointMetrics {
    fn default() -> Self {
        Self {
            total_bytes: AtomicU64::new(0),
            total_records: AtomicU64::new(0),
            buffered_records: AtomicU64::new(0),
            buffered_bytes: AtomicU64::new(0),
            circuit_input_records: AtomicU64::new(0),
            circuit_input_bytes: AtomicU64::new(0),
            num_transport_errors: AtomicU64::new(0),
            num_parse_errors: AtomicU64::new(0),
            end_of_input: AtomicBool::new(false),
            processing_latency_micros_histogram: Mutex::new(Self::processing_latency_histogram()),
            completion_latency_micros_histogram: Mutex::new(Self::completion_latency_histogram()),
        }
    }
}

impl InputEndpointMetrics {
    pub fn to_api_type(&self) -> ExternalInputEndpointMetrics {
        ExternalInputEndpointMetrics {
            total_bytes: self.total_bytes.load(Ordering::Relaxed),
            total_records: self.total_records.load(Ordering::Relaxed),
            buffered_records: self.buffered_records.load(Ordering::Relaxed),
            buffered_bytes: self.buffered_bytes.load(Ordering::Relaxed),
            num_transport_errors: self.num_transport_errors.load(Ordering::Relaxed),
            num_parse_errors: self.num_parse_errors.load(Ordering::Relaxed),
            end_of_input: self.end_of_input.load(Ordering::Relaxed),
        }
    }
}

// Latency histogram creation functions.
impl InputEndpointMetrics {
    pub fn processing_latency_histogram() -> SlidingHistogram {
        SlidingHistogram::new(10_000, Duration::from_secs(600))
    }

    pub fn completion_latency_histogram() -> SlidingHistogram {
        SlidingHistogram::new(10_000, Duration::from_secs(600))
    }
}

#[derive(Debug)]
pub struct StepResults {
    pub amt: BufferSize,
    pub resume: Option<Resume>,
}

#[derive(Debug)]
pub struct MissingReplay;

impl TryFrom<StepResults> for InputLog {
    type Error = MissingReplay;

    fn try_from(results: StepResults) -> Result<Self, Self::Error> {
        match results.resume {
            Some(Resume::Replay { seek, replay, hash }) => Ok(InputLog {
                data: replay,
                metadata: seek,
                checksums: InputChecksums {
                    num_records: results.amt.records as u64,
                    hash,
                },
            }),
            _ => Err(MissingReplay),
        }
    }
}

impl StepResults {
    pub(super) fn checksums(&self) -> Option<InputChecksums> {
        match self.resume {
            Some(Resume::Replay { hash, .. }) => Some(InputChecksums {
                hash,
                num_records: self.amt.records as u64,
            }),
            _ => None,
        }
    }
}

const MAX_TOKENLIST_LEN: usize = 1_000_000;

/// Tracks the set of tokens issued by the connector and the state needed to report
/// completion status for these tokens.
///
/// # Processing
///
/// Processing a token has two phases:
///
/// 1. Initialize the token consists as just the number of records ingested *by
///    the connector* at the time it was created.  After the circuit ingests all
///    of the records associated with the token, we additionally label the token
///    with the step in which the last record was ingested.
///
/// 2. Wait until all of the records output for the labeled step have been sent
///    to the output connector, by watching the `total_completed_steps` counter.
///
/// # Garbage collection
///
/// A new token is generated on each /completion_token API call, as well as for each
/// HTTP /ingress request. We therefore need to be careful to bound the number of tokens
/// in the list.
///
/// We do this by removing all except one token whose global offset is `<= total_completed_records`
/// from the list.  This bounds the number of tokens to the size of the connectors
/// input queue + the number of inputs currently in the circuit of buffered in the output
/// connectors, which can be quite large. Therefore we additionally bound the list to `MAX_TOKENLIST_LEN`
/// tokens, dropping the oldest token when queue is full. This is safe (i.e., we can never
/// report an operation that hasn't been completed as completed), but it can delay the
/// completion notification. In the worst case, if there are always `> MAX_TOKENLIST_LEN` incomplete
/// tokens in the queue, we will never issue any completion notifications, but that's
/// probably ok, as this means that no one's waiting for these tokens.
#[repr(transparent)]
struct TokenList(Mutex<TokenListInner>);

struct TokenListInner {
    /// List of (token, step) mappings.
    token_list: VecDeque<(u64, Option<Step>)>,
}

impl TokenListInner {
    fn new() -> Self {
        Self {
            token_list: VecDeque::new(),
        }
    }

    /// Perform queue maintenance.
    ///
    /// Invoked when a new token is pushed to the queue or when additional records are pushed from
    /// the connector to the circuit.
    ///
    /// * Set `total_initiated_steps` as the step for all tokens that are <= connector_input_records
    ///   and that don't have a step yet.
    /// * GC all but one token whose step is `< total_completed_steps`.
    /// * Remove tokens that exceed queue capacity.
    ///
    /// # Arguments
    ///
    /// * `connector_circuit_input_records` - total number of inputs **from this connector** ingested by the circuit.
    /// * `total_initiated_steps` - Number of steps that have been initiated.
    /// * `total_completed_steps` - Number of steps whose output has been delivered to connectors
    fn update(
        &mut self,
        connector_circuit_input_records: u64,
        total_initiated_steps: Step,
        total_completed_steps: Step,
    ) {
        // Label all unlabeled tokens <= `connector_circuit_input_records` with `total_initiated_steps`.
        let first_unlabeled = self
            .token_list
            .partition_point(|(_token, label)| label.is_some());

        for (token, label) in self.token_list.range_mut(first_unlabeled..) {
            if *token <= connector_circuit_input_records {
                *label = Some(total_initiated_steps)
            } else {
                break;
            }
        }

        // Delete all but one token whose global offset is `<= total_completed_steps`.
        let first_completed = self.token_list.partition_point(|(_token, label)| {
            if let Some(label) = label {
                *label <= total_completed_steps
            } else {
                false
            }
        });
        if first_completed > 0 {
            self.token_list.drain(0..first_completed - 1);
        }

        // If we've overflowed, delete the newest ones so that the old ones have
        // a chance to complete.
        self.token_list.truncate(MAX_TOKENLIST_LEN);
    }

    /// Create a token to track records ingested by the connector up to now.
    ///
    /// # Arguments
    ///
    /// * `token_input_records` - the number of records queued by the connector so far
    ///   (not all of these record have been pushed to the circuit).
    fn add_token(&mut self, token_input_records: u64) {
        // Don't track more than one token for the same offset.
        if let Some((token, _label)) = self.token_list.back() {
            // The number of ingested records grows monotonically.
            debug_assert!(*token <= token_input_records);
            if *token >= token_input_records {
                return;
            }
        }

        self.token_list.push_back((token_input_records, None));
    }

    /// Check completion status of `token`, an offset in the connector's input stream.
    ///
    /// Returns `None` if the token is definitely incomplete, or `Some(step)` if
    /// it is complete once [GlobalControllerMetrics::total_completed_steps]
    /// reaches `step`.
    fn completion_status(&self, token: u64) -> Option<Step> {
        // The first queued token `>= token`.
        let nearest_index = self
            .token_list
            .partition_point(|(queue_token, _label)| *queue_token < token);

        if let Some((_token, label)) = self.token_list.get(nearest_index) {
            *label
        } else {
            warn!("Client requested completion status of unknown token {token}");
            None
        }
    }
}

impl TokenList {
    fn new() -> Self {
        Self(Mutex::new(TokenListInner::new()))
    }
    fn update(
        &self,
        connector_circuit_input_records: u64,
        total_initiated_steps: Step,
        total_completed_steps: Step,
    ) {
        self.0.lock().unwrap().update(
            connector_circuit_input_records,
            total_initiated_steps,
            total_completed_steps,
        );
    }

    fn add_token(&self, token_input_records: u64) {
        self.0.lock().unwrap().add_token(token_input_records);
    }

    /// Check completion status of `token`, an offset in the connector's input stream.
    ///
    /// Returns `None` if the token is definitely incomplete, or `Some(step)` if
    /// it is complete once [GlobalControllerMetrics::total_completed_steps]
    /// reaches `step`.
    fn completion_status(&self, token: u64) -> Option<Step> {
        self.0.lock().unwrap().completion_status(token)
    }
}

/// Tracks how records ingested by the endpoint are processed by the pipeline.
///
/// ## Watermarks
///
/// A watermark is a marker associated with a batch of records ingested by the endpoint.
/// It consists of the timestamp when the connector started receiving data from the wire
/// and optional transport-specific metadata that identifies the position in the input
/// data stream, e.g., Kafka offset.
///
/// ## Processing latency and completion latency
///
/// **Processing latency** is the time it takes the pipeline to process a batch of records from
/// when they are ingested from the wire to when the circuit completes processing these records.
/// This includes: parsing, queuing, and processing by the circuit.
///
/// **Completion latency** includes processing latency plus the time it takes the pipeline to push
/// outputs to all output endpoints (this includes output buffering delay if
/// output buffers are enabled for some endpoints).
///
/// ```text
///             ┌──────────────────────────────────────────────────────────────────────────────┐
///             │                               Feldera                                        │
///             │                                                                              │
///          ┌──│──►input connector──────┐                              ┌────►output connector─│───┐
///          │  │                        │                              │                      │   │
///          │  │                        │                              │                      │   │
///          │  │                        │                              │                      │   │
///          │  │                        ▼                              │                      │   ▼
/// data source │                   input queue ──────►circuit─────►output queue               │ data sink
///             │                                                                              │
///             │                                                                              │
///             └──────────────────────────────────────────────────────────────────────────────┘
///
///                        1. end-to-end latency
/// ────────────────────────────────────────────────────────────────────────────────────────────►
///
///                        2. processing latency
///               ──────────────────────────────────────────────────────►
///
///                        3. completion latency
///               ──────────────────────────────────────────────────────────────────────────────►
/// ```
///
/// ## Tracking
///
/// Watermark tracking starts when the connector pushes a batch of records to the circuit before
/// the circuit performs a step. It notifies the circuit about the number of ingested records
/// via `extended()` call, and provides a list of watermarks, which indicate what input data will
/// be processed when the step completes, e.g.,:
/// * 100 records ingested at time t1, that corresponds to Kafka offsets 12000
/// * 100 records ingested at time t2, that corresponds to Kafka offsets 12100
/// * ...
///
/// When the step completes, the controller marks all watermarks belonging to the data processed in this step
/// as processed and adds the processed_at timestamp to them.
///
/// Later, when the outputs produced by this step have been pushed to all output endpoints,
/// the controller marks these watermarks as completed and updates two histograms: processing latency
/// and completion latency.
///
/// In addition, the controller keeps track of the latest watermark whose metadata is not `None`
/// and reports it via the `/status` API. This allows clients to track how far their input data
/// has been processed by the pipeline.
#[repr(transparent)]
pub(crate) struct WatermarkTracker(Mutex<WatermarkTrackerInner>);

impl WatermarkTracker {
    fn new() -> Self {
        Self(Mutex::new(WatermarkTrackerInner::new()))
    }

    fn append(&self, watermarks: Vec<Watermark>, global_offset: u64) {
        self.0.lock().unwrap().append(watermarks, global_offset);
    }

    /// Called when `total_processed_records` changes. Updates processed_at timestamps.
    fn update_processed(&self, processed_records: u64, ts: DateTime<Utc>) {
        self.0
            .lock()
            .unwrap()
            .update_processed(processed_records, ts);
    }

    /// Called when `total_completed_records` changes. Updates histograms, `completed_watermark`,
    /// and removes completed watermarks from the list.
    fn update_completed(
        &self,
        completed_records: u64,
        ts: DateTime<Utc>,
        metrics: &InputEndpointMetrics,
    ) {
        self.0
            .lock()
            .unwrap()
            .update_completed(completed_records, ts, metrics);
    }

    #[allow(dead_code)]
    pub(crate) fn debug(&self) -> String {
        format!("{:?}", self.0.lock().unwrap())
    }

    pub fn completed_watermark(&self) -> Option<CompletedWatermark> {
        self.0.lock().unwrap().completed_watermark.clone()
    }
}

/// Bound the number of tracked watermarks to avoid unbounded memory growth.
const MAX_TRACKED_WATERMARKS: usize = 10_000;

#[derive(Debug)]
struct WatermarkTrackerInner {
    /// The last `total_completed_records` value reported by the controller.
    total_completed_records: u64,

    /// The last `total_processed_records` value reported by the controller.
    total_processed_records: u64,

    /// Not-yet-completed watermarks with metadata=None, in the order they were added.
    watermark_list: VecDeque<WatermarkListEntry>,

    /// Not-yet-completed watermarks with metadata, in the order they were added.
    /// We keep these watermarks separate so that a large number of watermarks without
    /// metadata doesn't cause us to lose all watermarks with metadata, which are necessary
    /// to implement frontier tracking.
    watermark_with_metadata_list: VecDeque<WatermarkListEntry>,

    /// The latest completed watermark with non-empty metadata, if any.
    completed_watermark: Option<CompletedWatermark>,
}

impl WatermarkTrackerInner {
    fn new() -> Self {
        Self {
            total_completed_records: 0,
            total_processed_records: 0,
            watermark_list: VecDeque::new(),
            watermark_with_metadata_list: VecDeque::new(),
            completed_watermark: None,
        }
    }

    fn append(&mut self, watermarks: Vec<Watermark>, global_offset: u64) {
        for watermark in watermarks {
            if watermark.metadata.is_some() {
                // Keep the queue bounded. This gives older watermarks
                // a chance to complete; otherwise we risk always evicting all watermarks before the complete
                // and never reporting any completed watermarks.
                if self.watermark_with_metadata_list.len() >= MAX_TRACKED_WATERMARKS {
                    continue;
                }

                self.watermark_with_metadata_list
                    .push_back(WatermarkListEntry {
                        watermark,
                        global_offset,
                        processed_at: None,
                    });
            } else {
                if self.watermark_list.len() >= MAX_TRACKED_WATERMARKS {
                    continue;
                }
                self.watermark_list.push_back(WatermarkListEntry {
                    watermark,
                    global_offset,
                    processed_at: None,
                });
            }
        }
    }

    fn update_processed(&mut self, processed_records: u64, ts: DateTime<Utc>) {
        self.total_processed_records = processed_records;

        for entry in self.watermark_list.iter_mut() {
            if entry.global_offset <= processed_records && entry.processed_at.is_none() {
                entry.processed_at = Some(ts);
            } else {
                break;
            }
        }
        for entry in self.watermark_with_metadata_list.iter_mut() {
            if entry.global_offset <= processed_records && entry.processed_at.is_none() {
                entry.processed_at = Some(ts);
            } else {
                break;
            }
        }
    }

    fn update_completed(
        &mut self,
        completed_records: u64,
        ts: DateTime<Utc>,
        metrics: &InputEndpointMetrics,
    ) {
        self.total_completed_records = completed_records;

        while let Some(entry) = self.watermark_list.front() {
            if entry.global_offset <= completed_records {
                let entry = self.watermark_list.pop_front().unwrap();

                let processed_at = entry.processed_at.unwrap_or(ts);

                if let Ok(mut histogram) = metrics.processing_latency_micros_histogram.lock() {
                    histogram.record(
                        processed_at.timestamp_micros()
                            - entry.watermark.timestamp.timestamp_micros(),
                    )
                };

                if let Ok(mut histogram) = metrics.completion_latency_micros_histogram.lock() {
                    histogram.record(
                        ts.timestamp_micros() - entry.watermark.timestamp.timestamp_micros(),
                    )
                };
            } else {
                break;
            }
        }

        while let Some(entry) = self.watermark_with_metadata_list.front() {
            if entry.global_offset <= completed_records {
                let entry = self.watermark_with_metadata_list.pop_front().unwrap();

                let processed_at = entry.processed_at.unwrap_or(ts);

                if let Ok(mut histogram) = metrics.processing_latency_micros_histogram.lock() {
                    histogram.record(
                        processed_at.timestamp_micros()
                            - entry.watermark.timestamp.timestamp_micros(),
                    )
                };

                if let Ok(mut histogram) = metrics.completion_latency_micros_histogram.lock() {
                    histogram.record(
                        ts.timestamp_micros() - entry.watermark.timestamp.timestamp_micros(),
                    )
                };

                if let Some(metadata) = entry.watermark.metadata {
                    self.completed_watermark = Some(CompletedWatermark {
                        metadata,
                        ingested_at: entry.watermark.timestamp,
                        processed_at,
                        completed_at: ts,
                    })
                };
            } else {
                break;
            }
        }
    }
}

#[derive(Debug)]
struct WatermarkListEntry {
    watermark: Watermark,
    global_offset: u64,
    processed_at: Option<DateTime<Utc>>,
}

/// Input endpoint status information.
///
/// Keep in sync with [feldera_types::adapter_stats::ExternalInputEndpointStatus]
pub struct InputEndpointStatus {
    pub endpoint_name: String,

    /// Endpoint configuration (doesn't change).
    pub config: InputEndpointConfig,

    /// Performance metrics.
    pub metrics: InputEndpointMetrics,

    /// The first fatal error that occurred at the endpoint.
    pub fatal_error: Mutex<Option<String>>,

    /// Recent transport errors.
    pub transport_errors: Mutex<ConnectorErrorList>,

    /// Recent parse errors.
    pub parse_errors: Mutex<ConnectorErrorList>,

    /// Health status of the connector.
    pub health: Mutex<Option<ConnectorHealth>>,

    /// Progress within the latest step.
    pub progress: Mutex<Option<StepResults>>,

    /// May be None during endpoint initialization.
    pub reader: Option<Arc<dyn InputReader>>,

    /// Endpoint support for fault tolerance.
    pub fault_tolerance: Option<FtModel>,

    /// Endpoint has been paused by the user.
    ///
    /// When `true`, the endpoint doesn't produce any data even when the pipeline
    /// is running.
    ///
    /// This flag is set to `true` on startup if the `paused` flag in the
    /// endpoint configuration is `true`. At runtime, the value of the flag is
    /// controlled via the `/tables/<table_name>/connectors/<connector_name>/start` and
    /// `/tables/<table_name>/connectors/<connector_name>/pause` endpoints.
    pub paused: AtomicBool,

    /// Endpoint is currently a barrier checkpointing and suspend.
    ///
    /// An endpoint blocks checkpoint and suspend if its current input position
    /// is not one where it can resume.  In such a case, the user can still
    /// request checkpoint and suspend, which will be carried out as soon as the
    /// endpoint or endpoints advance beyond the barriers.
    pub barrier: AtomicBool,

    /// Connector-specific metrics for Prometheus export, registered via
    /// [`InputConsumer::set_custom_metrics`].
    pub custom_metrics: Option<Arc<dyn ConnectorMetrics>>,

    /// Completion tokens associated with the endpoint.
    completion_tokens: TokenList,

    pub(crate) completed_frontier: WatermarkTracker,
}

impl InputEndpointStatus {
    /// Convert the endpoint status to the API type.
    ///
    /// If `full` is true, generates a full version of the endpoint status,
    /// including the list of errors. This status can be returned via the
    /// `/tables/<table_name>/connectors/<connector_name>/status` endpoint.
    ///
    /// If `full` is false, generates a trimmed version of the endpoint status,
    /// not including the list of errors. This status can be returned via the
    /// `/stats` endpoint.
    pub fn to_api_type(&self, full: bool) -> ExternalInputEndpointStatus {
        ExternalInputEndpointStatus {
            endpoint_name: self.endpoint_name.clone(),
            config: feldera_types::adapter_stats::ShortEndpointConfig {
                stream: self.config.stream.clone().into_owned(),
            },
            metrics: self.metrics.to_api_type(),
            fatal_error: self.fatal_error(),
            parse_errors: if full {
                Some(self.parse_errors.lock().unwrap().to_api_type())
            } else {
                None
            },
            transport_errors: if full {
                Some(self.transport_errors.lock().unwrap().to_api_type())
            } else {
                None
            },
            health: self.health.lock().unwrap().clone(),
            paused: self.is_paused_by_user(),
            barrier: self.is_barrier(),
            completed_frontier: self.completed_frontier.completed_watermark(),
        }
    }
}

impl InputEndpointStatus {
    pub fn new(
        endpoint_name: &str,
        config: InputEndpointConfig,
        fault_tolerance: Option<FtModel>,
        initial_statistics: Option<&CheckpointInputEndpointMetrics>,
    ) -> Self {
        let paused_by_user =
            config.connector_config.paused || config.connector_config.start_after.is_some();

        Self {
            endpoint_name: endpoint_name.to_string(),
            config,
            metrics: initial_statistics
                .map_or(InputEndpointMetrics::default(), |initial_statistics| {
                    initial_statistics.into()
                }),
            fatal_error: Mutex::new(None),
            transport_errors: Mutex::new(ConnectorErrorList::new()),
            parse_errors: Mutex::new(ConnectorErrorList::new()),
            health: Mutex::new(None),
            progress: Mutex::new(None),
            paused: AtomicBool::new(paused_by_user),
            barrier: AtomicBool::new(false),
            reader: None,
            fault_tolerance,
            custom_metrics: None,
            completion_tokens: TokenList::new(),
            completed_frontier: WatermarkTracker::new(),
        }
    }

    /// Increment the number of buffered bytes and records; return
    /// the previous number of buffered records.
    fn add_buffered(&self, amt: BufferSize) -> u64 {
        // We are only updating statistics here, so no need to pay for
        // strong consistency.
        self.metrics
            .total_bytes
            .fetch_add(amt.bytes as u64, Ordering::Relaxed);
        self.metrics
            .total_records
            .fetch_add(amt.records as u64, Ordering::Relaxed);
        self.metrics
            .buffered_bytes
            .fetch_add(amt.bytes as u64, Ordering::Relaxed);
        self.metrics
            .buffered_records
            .fetch_add(amt.records as u64, Ordering::AcqRel)
    }

    fn eoi(&self) {
        debug!("endpoint {} has reached end of input", self.endpoint_name);
        self.metrics.end_of_input.store(true, Ordering::Release);
    }

    fn is_eoi(&self) -> bool {
        self.metrics.end_of_input.load(Ordering::Acquire)
    }

    /// True if the endpoint has reached end of input and doesn't have
    /// any more buffered records.
    fn finished(&self) -> bool {
        self.is_eoi() && self.metrics.buffered_records.load(Ordering::Acquire) == 0
    }

    /// Increment parser error counter.
    fn parse_error(&self, tag: Option<&str>, error: &ParseError) {
        let last_error = self.metrics.num_parse_errors.fetch_add(1, Ordering::AcqRel);

        self.parse_errors
            .lock()
            .unwrap()
            .add_error(tag, error, last_error + 1);
    }

    /// Increment transport error counter.  If this is the first fatal error,
    /// save it in `self.fatal_error`.
    pub fn transport_error(&self, fatal: bool, tag: Option<&str>, error: &AnyError) {
        let last_error = self
            .metrics
            .num_transport_errors
            .fetch_add(1, Ordering::AcqRel);

        self.transport_errors
            .lock()
            .unwrap()
            .add_error(tag, error, last_error + 1);

        if fatal {
            let mut fatal_error = self.fatal_error.lock().unwrap();
            if fatal_error.is_none() {
                *fatal_error = Some(error.to_string());
            }
        }
    }

    pub fn update_health(&self, health: ConnectorHealth) {
        *self.health.lock().unwrap() = Some(health);
    }

    /// True if the endpoint's `paused_by_user` flag is set to `true`.
    pub fn is_paused_by_user(&self) -> bool {
        self.paused.load(Ordering::Acquire)
    }

    /// Sets the paused-by-user flag to `paused` and returns the previous value.
    pub fn set_paused_by_user(&self, paused: bool) -> bool {
        self.paused.swap(paused, Ordering::Release)
    }

    /// True if the number of records buffered by the endpoint exceeds
    /// its `max_queued_records` config parameter.
    pub fn is_full(&self) -> bool {
        let buffered_records = self.metrics.buffered_records.load(Ordering::Acquire);
        let max_queued_records = self.config.connector_config.max_queued_records;
        buffered_records >= max_queued_records
    }

    /// Endpoint pushed additional records to the circuit.
    ///
    /// Updates:
    /// * `buffered_records`
    /// * `circuit_input_records`
    /// * `completion_tokens`
    ///
    /// # Arguments
    ///
    /// `total_circuit_input_records`, `total_initiated_steps`,
    /// `total_completed_steps` - current pipeline-level metrics (already
    /// updated with `step_results`) used to update `completion_tokens`.
    fn extended(
        &self,
        step_results: StepResults,
        watermarks: Vec<Watermark>,
        total_circuit_input_records: u64,
        total_initiated_steps: Step,
        total_completed_steps: Step,
    ) {
        let num_records = step_results.amt.records as u64;
        let num_bytes = step_results.amt.bytes as u64;
        SamplyEvent::new("input")
            .with_category("Step")
            .with_tooltip(|| {
                format!(
                    "{} submitted {num_records} records ({} bytes) for step {total_initiated_steps}",
                    &self.endpoint_name,
                    HumanBytes::from(num_bytes)
                )
            })
            .record();
        *self.progress.lock().unwrap() = Some(step_results);
        self.metrics
            .buffered_records
            .fetch_sub(num_records, Ordering::Relaxed);
        self.metrics
            .buffered_bytes
            .fetch_sub(num_bytes, Ordering::Relaxed);
        self.metrics
            .circuit_input_bytes
            .fetch_add(num_bytes, Ordering::Relaxed);
        let old_circuit_input_records = self
            .metrics
            .circuit_input_records
            .fetch_add(num_records, Ordering::AcqRel);
        self.completion_tokens.update(
            old_circuit_input_records + num_records,
            total_initiated_steps,
            total_completed_steps,
        );
        self.completed_frontier
            .append(watermarks, total_circuit_input_records);
    }

    /// Push a new completion token that corresponds to the current input offset to the
    /// `completion_tokens` queue.
    ///
    /// Returns the input offset that can be used to build `CompletionToken`.
    fn completion_token(&self, total_initiated_steps: Step, total_completed_steps: Step) -> u64 {
        let token_input_records = self.metrics.total_records.load(Ordering::Acquire);

        // To avoid a race, we add a token _before_ reading the circuit_input_records
        // metric and using it to update the token list.
        self.completion_tokens.add_token(token_input_records);

        let connector_circuit_input_records =
            self.metrics.circuit_input_records.load(Ordering::Acquire);

        self.completion_tokens.update(
            connector_circuit_input_records,
            total_initiated_steps,
            total_completed_steps,
        );
        token_input_records
    }

    /// Check completion status of `token`, an offset in the connector's input stream.
    ///
    /// Returns `None` if the token is definitely incomplete, or `Some(step)` if
    /// it is complete once [GlobalControllerMetrics::total_completed_steps]
    /// reaches `step`.
    fn completion_status(&self, token: u64) -> Option<Step> {
        self.completion_tokens.completion_status(token)
    }

    pub fn is_barrier(&self) -> bool {
        self.barrier.load(Ordering::Relaxed)
    }

    pub fn set_barrier(&self, barrier: bool) {
        self.barrier.store(barrier, Ordering::Relaxed);
    }

    pub fn fatal_error(&self) -> Option<String> {
        self.fatal_error.lock().unwrap().clone()
    }
}

// Keep in sync with feldera_types::ExternalOutputEndpointMetrics
#[derive(Default)]
pub struct OutputEndpointMetrics {
    /// Records and bytes sent on the underlying transport (HTTP, Kafka, etc.)
    /// to the endpoint.
    ///
    /// These only increase.
    pub transmitted_records: AtomicU64,
    pub transmitted_bytes: AtomicU64,

    /// Number of queued records.
    ///
    /// These are the records sent by the main circuit thread to the endpoint thread.
    /// Upon dequeuing the record, it gets buffered or sent directly to the output
    /// transport.
    ///
    /// These increase as records are queued.  They decrease as records are
    /// dequeued either to be buffered (see
    /// [buffered_records](Self::buffered_records)) or to be output directly.
    pub queued_records: AtomicU64,
    pub queued_batches: AtomicU64,

    /// Number of records pushed to the output buffer.
    ///
    /// Note that this may not be equal to the current size of the
    /// buffer, since the buffer consolidates records, e.g., inserts
    /// and deletes can cancel out over time.
    ///
    /// These increase as records are moved from queues to buffers.  They then
    /// fall abruptly to 0 because the buffers are always flushed as a whole.
    pub buffered_records: AtomicU64,
    pub buffered_batches: AtomicU64,

    pub num_encode_errors: AtomicU64,
    pub num_transport_errors: AtomicU64,

    /// The number of input records processed by the circuit.
    ///
    /// This metric tracks the end-to-end progress of the pipeline: the output
    /// of this endpoint is equal to the output of the circuit after
    /// processing `total_processed_input_records` records.
    ///
    /// In a multihost pipeline, this count reflects only the input records
    /// processed on the same host as the output endpoint, which is not usually
    /// meaningful.
    pub total_processed_input_records: AtomicU64,

    /// The number of steps whose input records have been processed by the
    /// endpoint.
    ///
    /// This is meaningful in a multihost pipeline because steps are
    /// synchronized across all of the hosts.
    ///
    /// # Interpretation
    ///
    /// This is a count, not a step number.  If `total_processed_steps` is 0, no
    /// steps have been processed to completion.  If `total_processed_steps >
    /// 0`, then the last step whose input records have been processed to
    /// completion is `total_processed_steps - 1`. A record that was ingested in
    /// step `n` is fully processed when `total_processed_steps > n`.
    pub total_processed_steps: Atomic<Step>,

    /// Extra memory in use beyond that used for queuing records.  Not all
    /// output connectors report this.
    pub memory: AtomicU64,
}

impl OutputEndpointMetrics {
    fn new(
        total_processed_input_records: u64,
        initial_statistics: Option<&CheckpointOutputEndpointMetrics>,
    ) -> Self {
        let initial_statistics = initial_statistics.cloned().unwrap_or_default();
        Self {
            transmitted_records: AtomicU64::new(initial_statistics.transmitted_records),
            transmitted_bytes: AtomicU64::new(initial_statistics.transmitted_bytes),
            queued_records: AtomicU64::new(0),
            queued_batches: AtomicU64::new(0),
            buffered_records: AtomicU64::new(0),
            buffered_batches: AtomicU64::new(0),
            num_encode_errors: AtomicU64::new(initial_statistics.num_encode_errors),
            num_transport_errors: AtomicU64::new(initial_statistics.num_transport_errors),
            total_processed_input_records: AtomicU64::new(total_processed_input_records),
            total_processed_steps: Atomic::new(0),
            memory: AtomicU64::new(0),
        }
    }

    /// Obtains a consistent snapshot of these statistics.
    pub fn snapshot(&self) -> ExternalOutputEndpointMetrics {
        // Fetch the statistics repeatedly until they don't change from one read
        // to the next.  It's pretty unlikely that they'll change at all, since
        // we're racing against I/O.
        let mut snapshot1 = self.to_api_type();
        loop {
            let snapshot2 = self.to_api_type();
            if snapshot1 == snapshot2 {
                return snapshot1;
            }
            snapshot1 = self.to_api_type();
            if snapshot1 == snapshot2 {
                return snapshot1;
            }
        }
    }

    pub fn to_api_type(&self) -> ExternalOutputEndpointMetrics {
        ExternalOutputEndpointMetrics {
            transmitted_records: self.transmitted_records.load(Ordering::Acquire),
            transmitted_bytes: self.transmitted_bytes.load(Ordering::Relaxed),
            queued_records: self.queued_records.load(Ordering::Relaxed),
            queued_batches: self.queued_batches.load(Ordering::Relaxed),
            buffered_records: self.buffered_records.load(Ordering::Relaxed),
            buffered_batches: self.buffered_batches.load(Ordering::Relaxed),
            num_encode_errors: self.num_encode_errors.load(Ordering::Relaxed),
            num_transport_errors: self.num_transport_errors.load(Ordering::Relaxed),
            total_processed_input_records: self
                .total_processed_input_records
                .load(Ordering::Relaxed),
            total_processed_steps: self.total_processed_steps.load(Ordering::Relaxed),
            memory: self.memory.load(Ordering::Relaxed),
        }
    }
}

/// Statistics for a collection of records being processed by an output
/// endpoint.
#[derive(Copy, Clone, Default)]
pub struct ProcessedRecords {
    /// Total number of input records ingested by the circuit across all
    /// endpoints corresponding to these records.
    ///
    /// When the records represented by these statistics are fully processed by
    /// the output endpoint, this value becomes the output endpoint's new
    /// [OutputEndpointMetrics::total_processed_input_records].
    pub total_processed_input_records: u64,

    /// The number of steps whose input records have been processed by the
    /// endpoint, once the endpoint has fully processed these records.
    ///
    /// When the records represented by these statistics are fully processed by
    /// the output endpoint, this value becomes the output endpoint's new
    /// [OutputEndpointMetrics::total_processed_steps].
    pub total_processed_steps: Step,
}

/// Recent connector errors.
///
/// Stores up to MAX_CONNECTOR_ERRORS most recent errors for each tag.
/// When the number of errors for a tag exceeds MAX_CONNECTOR_ERRORS, the oldest error is removed.
#[derive(Debug)]
pub struct ConnectorErrorList {
    errors: BTreeMap<Option<String>, VecDeque<ConnectorError>>,
}

impl Default for ConnectorErrorList {
    fn default() -> Self {
        Self::new()
    }
}

impl ConnectorErrorList {
    pub fn new() -> Self {
        Self {
            errors: BTreeMap::new(),
        }
    }

    pub fn add_error<E: Display>(&mut self, tag: Option<&str>, error: E, index: u64) {
        let tag = tag.map(|tag| tag.to_string());
        let entry: &mut VecDeque<ConnectorError> = self.errors.entry(tag.clone()).or_default();

        entry.push_back(ConnectorError {
            timestamp: Utc::now(),
            tag: tag.map(|tag| tag.to_string()),
            index,
            message: error.to_string(),
        });
        if entry.len() > MAX_CONNECTOR_ERRORS {
            entry.pop_front();
        }
    }

    pub fn to_api_type(&self) -> Vec<ConnectorError> {
        let mut errors = self
            .errors
            .iter()
            .flat_map(|(_tag, errors)| errors.iter().cloned())
            .collect::<Vec<_>>();
        errors.sort_by_key(|error| error.index);
        errors
    }
}

/// Output endpoint status information.
/// // Keep in sync with feldera_types::ExternalOutputEndpointStatus
pub struct OutputEndpointStatus {
    pub endpoint_name: String,

    /// Endpoint configuration (doesn't change).
    pub config: OutputEndpointConfig,

    /// Performance metrics.
    pub metrics: OutputEndpointMetrics,

    /// The first fatal error that occurred at the endpoint.
    pub fatal_error: Mutex<Option<String>>,

    /// Recent encode errors.
    pub encode_errors: Mutex<ConnectorErrorList>,

    /// Recent transport errors.
    pub transport_errors: Mutex<ConnectorErrorList>,

    pub health: Mutex<Option<ConnectorHealth>>,
}

impl OutputEndpointStatus {
    /// Convert the endpoint status to the API type.
    ///
    /// If `full` is true, generates a full version of the endpoint status,
    /// including the list of errors. This status can be returned via the
    /// `/tables/<table_name>/connectors/<connector_name>/status` endpoint.
    ///
    /// If `full` is false, generates a trimmed version of the endpoint status,
    /// not including the list of errors. This status can be returned via the
    /// `/stats` endpoint.
    pub fn to_api_type(&self, full: bool) -> ExternalOutputEndpointStatus {
        ExternalOutputEndpointStatus {
            endpoint_name: self.endpoint_name.clone(),
            config: ShortEndpointConfig {
                stream: self.config.stream.clone().into_owned(),
            },
            metrics: self.metrics.to_api_type(),
            fatal_error: self.fatal_error(),
            encode_errors: if full {
                Some(self.encode_errors.lock().unwrap().to_api_type())
            } else {
                None
            },
            transport_errors: if full {
                Some(self.transport_errors.lock().unwrap().to_api_type())
            } else {
                None
            },
            health: self.health.lock().unwrap().clone(),
        }
    }

    pub fn is_buffer_full(&self) -> bool {
        self.metrics.queued_records.load(Ordering::Acquire)
            >= self.config.connector_config.max_queued_records
    }

    pub fn is_busy(&self) -> bool {
        self.metrics.buffered_records.load(Ordering::Relaxed) != 0
            || self.metrics.queued_records.load(Ordering::Relaxed) != 0
    }

    pub fn fatal_error(&self) -> Option<String> {
        self.fatal_error.lock().unwrap().clone()
    }
}

/// Public read API.
impl OutputEndpointStatus {
    pub fn transmitted_records(&self) -> u64 {
        self.metrics.transmitted_records.load(Ordering::Acquire)
    }
}

impl OutputEndpointStatus {
    fn new(
        endpoint_name: &str,
        config: &OutputEndpointConfig,
        total_processed_records: u64,
        initial_statistics: Option<&CheckpointOutputEndpointMetrics>,
    ) -> Self {
        Self {
            endpoint_name: endpoint_name.to_string(),
            config: config.clone(),
            metrics: OutputEndpointMetrics::new(total_processed_records, initial_statistics),
            fatal_error: Mutex::new(None),
            encode_errors: Mutex::new(ConnectorErrorList::new()),
            transport_errors: Mutex::new(ConnectorErrorList::new()),
            health: Mutex::new(None),
        }
    }

    fn enqueue_batch(&self, num_records: usize) {
        self.metrics
            .queued_records
            .fetch_add(num_records as u64, Ordering::AcqRel);
        self.metrics.queued_batches.fetch_add(1, Ordering::AcqRel);
    }

    /// A batch has been pushed to the output transport directly from the queue,
    /// bypassing the buffer.
    fn output_batch(&self, processed: Option<ProcessedRecords>, num_records: usize) -> u64 {
        if let Some(processed) = processed {
            self.metrics
                .total_processed_input_records
                .store(processed.total_processed_input_records, Ordering::Release);
            self.metrics
                .total_processed_steps
                .store(processed.total_processed_steps, Ordering::Release);
        }

        let old = self
            .metrics
            .queued_records
            .fetch_sub(num_records as u64, Ordering::AcqRel);
        self.metrics.queued_batches.fetch_sub(1, Ordering::AcqRel);

        old
    }

    /// A batch has been read from a queue into the buffer.
    fn buffer_batch(&self, num_records: usize) -> u64 {
        self.metrics
            .buffered_records
            .fetch_add(num_records as u64, Ordering::AcqRel);
        // Don't count empty batches.
        if num_records > 0 {
            self.metrics.buffered_batches.fetch_add(1, Ordering::AcqRel);
        }

        let old = self
            .metrics
            .queued_records
            .fetch_sub(num_records as u64, Ordering::AcqRel);
        self.metrics.queued_batches.fetch_sub(1, Ordering::AcqRel);

        old
    }

    /// The content of the buffer has been pushed to the transport endpoint.
    ///
    /// # Arguments
    ///
    /// * `processed` - the output of the endpoint is now up to speed with the
    ///   output of the circuit after processing these records.
    fn output_buffered_batches(&self, processed: ProcessedRecords) {
        self.metrics.buffered_records.store(0, Ordering::Release);
        self.metrics.buffered_batches.store(0, Ordering::Release);

        self.metrics
            .total_processed_input_records
            .store(processed.total_processed_input_records, Ordering::Release);
        self.metrics
            .total_processed_steps
            .store(processed.total_processed_steps, Ordering::Release);
    }

    fn output_buffer(&self, num_bytes: usize, num_records: usize) {
        self.metrics
            .transmitted_bytes
            .fetch_add(num_bytes as u64, Ordering::Relaxed);
        self.metrics
            .transmitted_records
            .fetch_add(num_records as u64, Ordering::Relaxed);
    }

    /// Increment encoder error counter.
    fn encode_error(&self, tag: Option<&str>, error: &AnyError) {
        let last_error = self
            .metrics
            .num_encode_errors
            .fetch_add(1, Ordering::AcqRel);

        self.encode_errors
            .lock()
            .unwrap()
            .add_error(tag, error, last_error + 1);
    }

    /// Increment error counter.  If this is the first fatal error,
    /// save it in `self.fatal_error`.
    fn transport_error(&self, fatal: bool, tag: Option<&str>, error: &AnyError) {
        let last_error = self
            .metrics
            .num_transport_errors
            .fetch_add(1, Ordering::AcqRel);

        self.transport_errors
            .lock()
            .unwrap()
            .add_error(tag, error, last_error + 1);

        if fatal {
            let mut fatal_error = self.fatal_error.lock().unwrap();
            if fatal_error.is_none() {
                *fatal_error = Some(error.to_string());
            }
        }
    }

    fn update_health(&self, health: ConnectorHealth) {
        *self.health.lock().unwrap() = Some(health);
    }

    fn num_total_processed_input_records(&self) -> u64 {
        self.metrics
            .total_processed_input_records
            .load(Ordering::Acquire)
    }

    fn num_total_processed_steps(&self) -> u64 {
        self.metrics.total_processed_steps.load(Ordering::Acquire)
    }
}
