datum-sql 0.10.3

use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use arrow::array::{
    Array, RecordBatch, TimestampMicrosecondArray, TimestampMillisecondArray,
    TimestampNanosecondArray, TimestampSecondArray,
};
use arrow::datatypes::{DataType, Schema, TimeUnit};
use datafusion::common::{DataFusionError, Result};
use datum::{NotUsed, Source, StreamCompletion, StreamError, StreamResult, UniqueKillSwitch};

use crate::{ChangelogBatch, stream_error};

/// Event-time progress in nanoseconds since the Unix epoch.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Watermark {
    timestamp_ns: i64,
}

impl Watermark {
    /// Creates a watermark from nanoseconds since the Unix epoch.
    #[must_use]
    pub const fn new(timestamp_ns: i64) -> Self {
        Self { timestamp_ns }
    }

    /// Returns the watermark timestamp in nanoseconds since the Unix epoch.
    #[must_use]
    pub const fn timestamp_ns(self) -> i64 {
        self.timestamp_ns
    }
}

/// Future checkpoint/alignment marker carried outside SQL row data.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SqlBarrier {
    epoch: u64,
}

impl SqlBarrier {
    /// Creates a barrier marker for a checkpoint epoch.
    #[must_use]
    pub const fn new(epoch: u64) -> Self {
        Self { epoch }
    }

    /// Returns the checkpoint epoch.
    #[must_use]
    pub const fn epoch(self) -> u64 {
        self.epoch
    }
}

/// In-band SQL stream envelope for data and control records.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SqlEvent<T> {
    /// User data payload.
    Data(T),
    /// Event-time progress marker.
    Watermark(Watermark),
    /// Checkpoint or alignment marker.
    Barrier(SqlBarrier),
}

impl<T> SqlEvent<T> {
    /// Maps only the data branch and passes watermarks and barriers through.
    pub fn map_data_result<U, F>(self, f: F) -> StreamResult<SqlEvent<U>>
    where
        F: FnOnce(T) -> StreamResult<U>,
    {
        match self {
            Self::Data(data) => f(data).map(SqlEvent::Data),
            Self::Watermark(watermark) => Ok(SqlEvent::Watermark(watermark)),
            Self::Barrier(barrier) => Ok(SqlEvent::Barrier(barrier)),
        }
    }

    /// Returns the contained watermark when this event is a watermark.
    #[must_use]
    pub const fn as_watermark(&self) -> Option<Watermark> {
        match self {
            Self::Watermark(watermark) => Some(*watermark),
            Self::Data(_) | Self::Barrier(_) => None,
        }
    }
}

/// Strategy for producing watermarks from an event-time column.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WatermarkStrategy {
    /// Emit `max_observed_event_time - max_out_of_orderness`, monotonically.
    BoundedOutOfOrderness {
        /// Maximum event-time delay tolerated before a row is considered late.
        max_out_of_orderness: Duration,
    },
}

/// Event-time declaration for a registered SQL table.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EventTimeConfig {
    column: Arc<str>,
    strategy: WatermarkStrategy,
}

impl EventTimeConfig {
    /// Declares event time with a bounded-out-of-orderness watermark strategy.
    #[must_use]
    pub fn bounded_out_of_orderness(
        column: impl Into<Arc<str>>,
        max_out_of_orderness: Duration,
    ) -> Self {
        Self {
            column: column.into(),
            strategy: WatermarkStrategy::BoundedOutOfOrderness {
                max_out_of_orderness,
            },
        }
    }

    /// Returns the Arrow timestamp column used for event time.
    #[must_use]
    pub fn column(&self) -> &str {
        &self.column
    }

    /// Returns the watermark strategy.
    #[must_use]
    pub const fn strategy(&self) -> &WatermarkStrategy {
        &self.strategy
    }

    pub(crate) fn resolve(&self, schema: &Schema) -> Result<ResolvedEventTimeConfig> {
        let column_index = schema.index_of(&self.column).map_err(|_| {
            DataFusionError::Plan(format!(
                "event-time column '{}' does not exist in table schema",
                self.column
            ))
        })?;
        let field = schema.field(column_index);
        if !matches!(field.data_type(), DataType::Timestamp(_, _)) {
            return Err(DataFusionError::Plan(format!(
                "event-time column '{}' must be an Arrow Timestamp, found {:?}",
                self.column,
                field.data_type()
            )));
        }
        if field.is_nullable() {
            return Err(DataFusionError::Plan(format!(
                "event-time column '{}' must be non-nullable",
                self.column
            )));
        }

        let max_out_of_orderness_ns = match &self.strategy {
            WatermarkStrategy::BoundedOutOfOrderness {
                max_out_of_orderness,
            } => duration_to_ns(*max_out_of_orderness)?,
        };

        Ok(ResolvedEventTimeConfig {
            column_index,
            max_out_of_orderness_ns,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ResolvedEventTimeConfig {
    column_index: usize,
    max_out_of_orderness_ns: i64,
}

/// Payloads that expose the Arrow batch used to read event-time columns.
pub trait SqlEventPayload {
    /// Returns the Arrow batch containing the configured event-time column.
    fn event_time_batch(&self) -> &RecordBatch;

    /// Partition identity for one row within this source.
    ///
    /// The identity is local to a single source. Sources that do not expose
    /// partition metadata return `None` and are treated as one logical
    /// partition.
    fn event_time_partition(&self, _row: usize) -> Option<i64> {
        None
    }

    /// Partitions that can still produce rows for this source.
    ///
    /// Returning assigned-but-idle partitions lets the watermark generator hold
    /// progress until every active partition has produced event-time progress.
    fn event_time_active_partitions(&self) -> Option<Vec<i64>> {
        None
    }
}

impl SqlEventPayload for RecordBatch {
    fn event_time_batch(&self) -> &RecordBatch {
        self
    }
}

impl SqlEventPayload for ChangelogBatch {
    fn event_time_batch(&self) -> &RecordBatch {
        self.batch()
    }
}

#[must_use]
/// Wraps every payload from a source as `SqlEvent::Data`.
pub fn data_events<T, Mat>(source: Source<T, Mat>) -> Source<SqlEvent<T>, Mat>
where
    T: Send + 'static,
    Mat: Send + 'static,
{
    source.map(SqlEvent::Data)
}

/// Assigns event-time watermarks to a source.
///
/// The returned source emits each data payload followed by a watermark whenever
/// the configured bounded-out-of-orderness strategy advances.
pub fn assign_event_time_watermarks<T, Mat>(
    source: Source<T, Mat>,
    schema: &Schema,
    config: EventTimeConfig,
) -> Result<Source<SqlEvent<T>, Mat>>
where
    T: SqlEventPayload + Send + 'static,
    Mat: Send + 'static,
{
    let resolved = config.resolve(schema)?;
    Ok(assign_resolved_event_time_watermarks(source, resolved))
}

pub(crate) fn assign_resolved_event_time_watermarks<T, Mat>(
    source: Source<T, Mat>,
    config: ResolvedEventTimeConfig,
) -> Source<SqlEvent<T>, Mat>
where
    T: SqlEventPayload + Send + 'static,
    Mat: Send + 'static,
{
    source.try_stateful_map_concat(WatermarkGenerator::new(config), |generator, payload| {
        generator.apply(payload)
    })
}

/// Applies a function to the data branch of a SQL event stream.
pub fn map_sql_event_data<T, U, Mat, F>(
    source: Source<SqlEvent<T>, Mat>,
    f: F,
) -> Source<SqlEvent<U>, Mat>
where
    T: Send + 'static,
    U: Send + 'static,
    Mat: Send + 'static,
    F: Fn(T) -> StreamResult<U> + Send + Sync + 'static,
{
    source.try_map(move |event| event.map_data_result(&f))
}

#[derive(Clone)]
struct WatermarkGenerator {
    config: ResolvedEventTimeConfig,
    max_seen_event_time_ns_by_partition: BTreeMap<i64, Option<i64>>,
    last_emitted_watermark_ns: Option<i64>,
}

impl WatermarkGenerator {
    fn new(config: ResolvedEventTimeConfig) -> Self {
        Self {
            config,
            max_seen_event_time_ns_by_partition: BTreeMap::new(),
            last_emitted_watermark_ns: None,
        }
    }

    fn apply<T>(&mut self, payload: T) -> StreamResult<Vec<SqlEvent<T>>>
    where
        T: SqlEventPayload,
    {
        if let Some(active_partitions) = payload.event_time_active_partitions() {
            for partition in active_partitions {
                self.max_seen_event_time_ns_by_partition
                    .entry(partition)
                    .or_insert(None);
            }
        }
        let batch_max_by_partition =
            max_event_time_ns_by_partition(&payload, self.config.column_index)?;
        let mut out = vec![SqlEvent::Data(payload)];
        for (partition, batch_max) in batch_max_by_partition {
            let entry = self
                .max_seen_event_time_ns_by_partition
                .entry(partition)
                .or_insert(None);
            *entry = Some(entry.map_or(batch_max, |previous| previous.max(batch_max)));
        }

        if let Some(min_seen) = self.min_seen_event_time_ns() {
            let watermark_ns = min_seen
                .checked_sub(self.config.max_out_of_orderness_ns)
                .ok_or_else(|| stream_error("event-time watermark underflowed i64 nanoseconds"))?;
            if self
                .last_emitted_watermark_ns
                .is_none_or(|previous| watermark_ns > previous)
            {
                self.last_emitted_watermark_ns = Some(watermark_ns);
                out.push(SqlEvent::Watermark(Watermark::new(watermark_ns)));
            }
        }

        Ok(out)
    }

    fn min_seen_event_time_ns(&self) -> Option<i64> {
        let mut min_seen = None;
        for value in self.max_seen_event_time_ns_by_partition.values() {
            let value = (*value)?;
            min_seen = Some(min_seen.map_or(value, |current: i64| current.min(value)));
        }
        min_seen
    }
}

const DEFAULT_EVENT_TIME_PARTITION: i64 = 0;

fn max_event_time_ns_by_partition<T>(
    payload: &T,
    column_index: usize,
) -> StreamResult<BTreeMap<i64, i64>>
where
    T: SqlEventPayload,
{
    let batch = payload.event_time_batch();
    if batch.num_rows() == 0 {
        return Ok(BTreeMap::new());
    }
    if column_index >= batch.num_columns() {
        return Err(stream_error(format!(
            "event-time column index {column_index} is out of range for {} columns",
            batch.num_columns()
        )));
    }

    let column = batch.column(column_index);
    match batch.schema().field(column_index).data_type() {
        DataType::Timestamp(TimeUnit::Second, _) => {
            timestamp_max_ns_by_partition::<TimestampSecondArray, T>(
                column.as_any().downcast_ref::<TimestampSecondArray>(),
                TimeUnit::Second,
                payload,
            )
        }
        DataType::Timestamp(TimeUnit::Millisecond, _) => {
            timestamp_max_ns_by_partition::<TimestampMillisecondArray, T>(
                column.as_any().downcast_ref::<TimestampMillisecondArray>(),
                TimeUnit::Millisecond,
                payload,
            )
        }
        DataType::Timestamp(TimeUnit::Microsecond, _) => {
            timestamp_max_ns_by_partition::<TimestampMicrosecondArray, T>(
                column.as_any().downcast_ref::<TimestampMicrosecondArray>(),
                TimeUnit::Microsecond,
                payload,
            )
        }
        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
            timestamp_max_ns_by_partition::<TimestampNanosecondArray, T>(
                column.as_any().downcast_ref::<TimestampNanosecondArray>(),
                TimeUnit::Nanosecond,
                payload,
            )
        }
        other => Err(stream_error(format!(
            "event-time column must be an Arrow Timestamp, found {other:?}"
        ))),
    }
}

fn timestamp_max_ns_by_partition<A, T>(
    array: Option<&A>,
    unit: TimeUnit,
    payload: &T,
) -> StreamResult<BTreeMap<i64, i64>>
where
    A: Array + TimestampValues,
    T: SqlEventPayload,
{
    let array = array.ok_or_else(|| stream_error("event-time column array type mismatch"))?;
    let multiplier = timestamp_unit_multiplier(unit);
    let mut max_by_partition = BTreeMap::new();
    for row in 0..array.len() {
        if array.is_null(row) {
            return Err(stream_error(format!(
                "event-time column contains null at row {row}"
            )));
        }
        let value = array.value_at(row);
        let ns = value
            .checked_mul(multiplier)
            .ok_or_else(|| stream_error("event-time value overflowed i64 nanoseconds"))?;
        let partition = payload
            .event_time_partition(row)
            .unwrap_or(DEFAULT_EVENT_TIME_PARTITION);
        max_by_partition
            .entry(partition)
            .and_modify(|previous: &mut i64| *previous = (*previous).max(ns))
            .or_insert(ns);
    }
    Ok(max_by_partition)
}

trait TimestampValues {
    fn value_at(&self, row: usize) -> i64;
}

impl TimestampValues for TimestampSecondArray {
    fn value_at(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampValues for TimestampMillisecondArray {
    fn value_at(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampValues for TimestampMicrosecondArray {
    fn value_at(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampValues for TimestampNanosecondArray {
    fn value_at(&self, row: usize) -> i64 {
        self.value(row)
    }
}

const fn timestamp_unit_multiplier(unit: TimeUnit) -> i64 {
    match unit {
        TimeUnit::Second => 1_000_000_000,
        TimeUnit::Millisecond => 1_000_000,
        TimeUnit::Microsecond => 1_000,
        TimeUnit::Nanosecond => 1,
    }
}

fn duration_to_ns(duration: Duration) -> Result<i64> {
    i64::try_from(duration.as_nanos()).map_err(|_| {
        DataFusionError::Plan(format!(
            "event-time watermark delay {duration:?} exceeds i64 nanoseconds"
        ))
    })
}

/// Materialized handle for a streaming SQL query.
pub struct ContinuousQueryHandle {
    kill_switch: UniqueKillSwitch,
    completion: StreamCompletion<NotUsed>,
}

impl ContinuousQueryHandle {
    pub(crate) fn new(
        kill_switch: UniqueKillSwitch,
        completion: StreamCompletion<NotUsed>,
    ) -> Self {
        Self {
            kill_switch,
            completion,
        }
    }

    /// Requests graceful shutdown of the continuous query.
    pub fn cancel(&self) {
        self.kill_switch.shutdown();
    }

    /// Aborts the continuous query with a stream error.
    pub fn abort(&self, error: StreamError) {
        self.kill_switch.abort(error);
    }

    /// Waits for query completion.
    pub fn wait(self) -> StreamResult<NotUsed> {
        self.completion.wait()
    }

    /// Polls completion without blocking.
    #[must_use]
    pub fn try_wait(&mut self) -> Option<StreamResult<NotUsed>> {
        self.completion.try_wait()
    }
}

impl fmt::Debug for ContinuousQueryHandle {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ContinuousQueryHandle")
            .finish_non_exhaustive()
    }
}