datum-cdc 0.10.1

use std::{
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

use datum::{Channel, Source, SourceWithContext};
use pgwire_replication::{ReplicationClient, ReplicationEvent};
use tokio::sync::{mpsc, oneshot};

use crate::{
    CdcCheckpointHandle, CdcCheckpointStore, CdcError, CdcLag, CdcOffset, CdcResult, CdcStart,
    ChangeEvent, FeedbackCommand, FeedbackCoordinator, FeedbackMode, PgLsn, PostgresCdcConfig,
    ReconnectSettings, SlotLifecycle, SourceMetadata, TransactionMeta, admin,
    config::SourceSettings,
    pgoutput::{DecodedChange, PgOutputDecoder},
    runtime,
};

/// Entry point for CDC source builders.
pub struct CdcSource;

impl CdcSource {
    #[must_use]
    pub fn postgres() -> PostgresCdcBuilder {
        PostgresCdcBuilder::default()
    }
}

/// Builder for a PostgreSQL pgoutput source.
#[derive(Clone)]
pub struct PostgresCdcBuilder {
    postgres: Option<PostgresCdcConfig>,
    slot: Option<String>,
    publication: Option<String>,
    start: CdcStart,
    feedback: FeedbackMode,
    slot_lifecycle: SlotLifecycle,
    buffer_capacity: usize,
    replication_buffer_events: usize,
    status_interval: Duration,
    idle_wakeup_interval: Duration,
    validate_publication: bool,
    checkpoint_store: Option<Arc<dyn CdcCheckpointStore>>,
    reconnect: Option<ReconnectSettings>,
}

impl Default for PostgresCdcBuilder {
    fn default() -> Self {
        Self {
            postgres: None,
            slot: None,
            publication: None,
            start: CdcStart::default(),
            feedback: FeedbackMode::default(),
            slot_lifecycle: SlotLifecycle::default(),
            buffer_capacity: 8192,
            replication_buffer_events: 8192,
            status_interval: Duration::from_secs(1),
            idle_wakeup_interval: Duration::from_secs(10),
            validate_publication: true,
            checkpoint_store: None,
            reconnect: Some(ReconnectSettings::default()),
        }
    }
}

impl PostgresCdcBuilder {
    #[must_use]
    pub fn connect(mut self, config: PostgresCdcConfig) -> Self {
        self.postgres = Some(config);
        self
    }

    pub fn connect_url(self, url: impl AsRef<str>) -> CdcResult<Self> {
        Ok(self.connect(PostgresCdcConfig::from_url(url)?))
    }

    #[must_use]
    pub fn slot(mut self, slot: impl Into<String>) -> Self {
        self.slot = Some(slot.into());
        self
    }

    #[must_use]
    pub fn publication(mut self, publication: impl Into<String>) -> Self {
        self.publication = Some(publication.into());
        self
    }

    #[must_use]
    pub fn start_from(mut self, start: CdcStart) -> Self {
        self.start = start;
        self
    }

    #[must_use]
    pub fn feedback(mut self, feedback: FeedbackMode) -> Self {
        self.feedback = feedback;
        self
    }

    #[must_use]
    pub fn slot_lifecycle(mut self, lifecycle: SlotLifecycle) -> Self {
        self.slot_lifecycle = lifecycle;
        self
    }

    #[must_use]
    pub fn buffer_capacity(mut self, capacity: usize) -> Self {
        self.buffer_capacity = capacity;
        self
    }

    #[must_use]
    pub fn replication_buffer_events(mut self, capacity: usize) -> Self {
        self.replication_buffer_events = capacity;
        self
    }

    #[must_use]
    pub fn status_interval(mut self, interval: Duration) -> Self {
        self.status_interval = interval;
        self
    }

    #[must_use]
    pub fn idle_wakeup_interval(mut self, interval: Duration) -> Self {
        self.idle_wakeup_interval = interval;
        self
    }

    #[must_use]
    pub fn validate_publication(mut self, enabled: bool) -> Self {
        self.validate_publication = enabled;
        self
    }

    #[must_use]
    pub fn checkpoint_store<S>(mut self, store: S) -> Self
    where
        S: CdcCheckpointStore,
    {
        self.checkpoint_store = Some(Arc::new(store));
        self
    }

    #[must_use]
    pub fn checkpoint_store_arc(mut self, store: Arc<dyn CdcCheckpointStore>) -> Self {
        self.checkpoint_store = Some(store);
        self
    }

    #[must_use]
    pub fn reconnect(mut self, settings: ReconnectSettings) -> Self {
        self.reconnect = Some(settings);
        self
    }

    #[must_use]
    pub fn disable_reconnect(mut self) -> Self {
        self.reconnect = None;
        self
    }

    pub fn build(self) -> CdcResult<Source<ChangeEvent, CdcHandle>> {
        let settings = Arc::new(self.finalize()?);
        Ok(source_from_settings(settings))
    }

    pub fn build_with_context(
        self,
    ) -> CdcResult<SourceWithContext<ChangeEvent, CdcOffset, CdcHandle>> {
        Ok(self
            .build()?
            .as_source_with_context(|event| event.lsn.clone()))
    }

    pub fn restart_factory(
        self,
    ) -> CdcResult<impl Fn() -> Source<ChangeEvent, CdcHandle> + Clone + Send + Sync + 'static>
    {
        let settings = Arc::new(self.finalize()?);
        Ok(move || source_from_settings(Arc::clone(&settings)))
    }

    fn finalize(self) -> CdcResult<SourceSettings> {
        if self.buffer_capacity == 0 {
            return Err(CdcError::Config(
                "CDC source buffer capacity must be greater than zero".into(),
            ));
        }
        if self.replication_buffer_events == 0 {
            return Err(CdcError::Config(
                "replication event buffer capacity must be greater than zero".into(),
            ));
        }
        let postgres = self
            .postgres
            .ok_or_else(|| CdcError::Config("PostgreSQL connection is required".into()))?;
        if !matches!(postgres.tls, crate::CdcTlsConfig::Disable) {
            return Err(CdcError::Config(
                "TLS replication is not enabled in this workspace build because pgwire-replication 0.3.2 forces rustls/ring and conflicts with Datum's aws-lc-rs rustls provider; use sslmode=disable/prefer for this MVP"
                    .into(),
            ));
        }
        let slot = self
            .slot
            .ok_or_else(|| CdcError::Config("replication slot is required".into()))?;
        let publication = self
            .publication
            .ok_or_else(|| CdcError::Config("publication is required".into()))?;
        Ok(SourceSettings {
            source: SourceMetadata {
                database: postgres.database.clone(),
                slot: slot.clone(),
                publication: publication.clone(),
            },
            postgres,
            slot,
            publication,
            start: self.start,
            feedback: self.feedback,
            slot_lifecycle: self.slot_lifecycle,
            buffer_capacity: self.buffer_capacity,
            replication_buffer_events: self.replication_buffer_events,
            status_interval: self.status_interval,
            idle_wakeup_interval: self.idle_wakeup_interval,
            validate_publication: self.validate_publication,
            checkpoint_store: self.checkpoint_store,
            reconnect: self.reconnect,
        })
    }
}

fn source_from_settings(settings: Arc<SourceSettings>) -> Source<ChangeEvent, CdcHandle> {
    Source::channel(settings.buffer_capacity).map_materialized_value(move |channel| {
        let (commands, command_rx) = mpsc::unbounded_channel();
        let health = Arc::new(HealthState::default());
        let checkpoint = CdcCheckpointHandle::new(
            settings.slot.clone(),
            settings.checkpoint_store.clone(),
            commands.clone(),
        );
        let task_settings = Arc::clone(&settings);
        let task_health = Arc::clone(&health);
        runtime::spawn(async move {
            run_carrier(task_settings, channel, command_rx, task_health).await;
        });
        CdcHandle {
            settings: Arc::clone(&settings),
            checkpoint,
            commands,
            health,
        }
    })
}

/// Materialized management handle for a CDC source.
pub struct CdcHandle {
    settings: Arc<SourceSettings>,
    checkpoint: CdcCheckpointHandle,
    commands: mpsc::UnboundedSender<FeedbackCommand>,
    health: Arc<HealthState>,
}

impl CdcHandle {
    #[must_use]
    pub fn checkpoint_handle(&self) -> CdcCheckpointHandle {
        self.checkpoint.clone()
    }

    #[must_use]
    pub fn health(&self) -> CdcHealth {
        self.health.snapshot()
    }

    pub fn stop(&self) -> CdcResult<()> {
        self.commands
            .send(FeedbackCommand::Stop)
            .map_err(|_| CdcError::ChannelClosed)
    }

    pub async fn drop_slot(&self) -> CdcResult<()> {
        self.drop_slot_inner(false).await
    }

    pub async fn force_drop_slot(&self) -> CdcResult<()> {
        self.drop_slot_inner(true).await
    }

    pub async fn lag(&self) -> CdcResult<CdcLag> {
        admin::sample_lag(&self.settings.postgres, &self.settings.slot).await
    }

    async fn drop_slot_inner(&self, force: bool) -> CdcResult<()> {
        let (reply, wait) = oneshot::channel();
        if self
            .commands
            .send(FeedbackCommand::DropSlot { force, reply })
            .is_ok()
        {
            return wait
                .await
                .map_err(|err| CdcError::Runtime(format!("drop-slot reply lost: {err}")))?;
        }
        admin::drop_slot(
            &self.settings.postgres,
            &self.settings.slot,
            self.settings.slot_lifecycle,
            force,
        )
        .await
    }
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CdcHealth {
    pub running: bool,
    pub last_error: Option<String>,
    pub last_received_lsn: PgLsn,
    pub latest_server_wal_end: PgLsn,
    pub last_feedback_lsn: PgLsn,
    pub emitted_events: u64,
    pub reconnects: u64,
}

#[derive(Default)]
struct HealthState {
    snapshot: Mutex<CdcHealth>,
    emitted_events: AtomicU64,
    reconnects: AtomicU64,
}

impl HealthState {
    fn snapshot(&self) -> CdcHealth {
        let mut snapshot = self
            .snapshot
            .lock()
            .expect("CDC health state poisoned")
            .clone();
        snapshot.emitted_events = self.emitted_events.load(Ordering::Acquire);
        snapshot.reconnects = self.reconnects.load(Ordering::Acquire);
        snapshot
    }

    fn set_running(&self, running: bool) {
        self.update(|health| health.running = running);
    }

    fn set_error(&self, error: impl ToString) {
        let error = error.to_string();
        self.update(|health| health.last_error = Some(error));
    }

    fn clear_error(&self) {
        self.update(|health| health.last_error = None);
    }

    fn record_received(&self, lsn: PgLsn) {
        self.update(|health| health.last_received_lsn = lsn);
    }

    fn record_server_wal_end(&self, lsn: PgLsn) {
        self.update(|health| health.latest_server_wal_end = lsn);
    }

    fn record_feedback(&self, lsn: PgLsn) {
        self.update(|health| health.last_feedback_lsn = lsn);
    }

    fn record_emitted(&self) {
        self.emitted_events.fetch_add(1, Ordering::AcqRel);
    }

    fn record_reconnect(&self) {
        self.reconnects.fetch_add(1, Ordering::AcqRel);
    }

    fn update(&self, f: impl FnOnce(&mut CdcHealth)) {
        let mut health = self.snapshot.lock().expect("CDC health state poisoned");
        f(&mut health);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CarrierStop {
    SourceClosed,
    Stopped,
    DroppedSlot,
    ReplicationEnded,
}

async fn run_carrier(
    settings: Arc<SourceSettings>,
    channel: Channel<ChangeEvent>,
    mut commands: mpsc::UnboundedReceiver<FeedbackCommand>,
    health: Arc<HealthState>,
) {
    health.set_running(true);
    let mut next_backoff = settings
        .reconnect
        .as_ref()
        .map(|reconnect| reconnect.min_backoff);
    let deadline = settings.start_deadline();
    loop {
        match run_once(&settings, &channel, &mut commands, &health).await {
            Ok(CarrierStop::SourceClosed | CarrierStop::Stopped | CarrierStop::DroppedSlot) => {
                break;
            }
            Ok(CarrierStop::ReplicationEnded) => {
                if settings.reconnect.is_none() {
                    break;
                }
            }
            Err(error) => {
                health.set_error(&error);
                if !should_reconnect(&error, settings.reconnect.as_ref(), deadline) {
                    break;
                }
            }
        }

        let Some(reconnect) = settings.reconnect.as_ref() else {
            break;
        };
        let delay = next_backoff.unwrap_or(reconnect.min_backoff);
        tokio::time::sleep(delay).await;
        health.record_reconnect();
        next_backoff = Some(delay.saturating_mul(2).min(reconnect.max_backoff));
    }
    channel.close();
    health.set_running(false);
}

fn should_reconnect(
    error: &CdcError,
    reconnect: Option<&ReconnectSettings>,
    deadline: Option<Instant>,
) -> bool {
    if reconnect.is_none() || deadline.is_some_and(|deadline| Instant::now() >= deadline) {
        return false;
    }
    matches!(
        error,
        CdcError::Replication(_) | CdcError::Postgres(_) | CdcError::Io(_)
    )
}

async fn run_once(
    settings: &SourceSettings,
    channel: &Channel<ChangeEvent>,
    commands: &mut mpsc::UnboundedReceiver<FeedbackCommand>,
    health: &HealthState,
) -> CdcResult<CarrierStop> {
    let start_lsn = admin::prepare_slot_and_start_lsn(settings).await?;
    let mut feedback = FeedbackCoordinator::new(settings.slot.clone(), start_lsn);
    let replication_config = settings.postgres.replication_config(
        &settings.slot,
        &settings.publication,
        start_lsn,
        settings.status_interval,
        settings.idle_wakeup_interval,
        settings.replication_buffer_events,
    );
    let mut client = ReplicationClient::connect(replication_config).await?;
    health.clear_error();
    let mut decoder = PgOutputDecoder::default();
    let mut tx = None;

    loop {
        tokio::select! {
            Some(command) = commands.recv() => {
                if let Some(stop) = handle_command(command, settings, &mut client, &mut feedback, health).await? {
                    return Ok(stop);
                }
            }
            event = client.recv() => {
                let Some(event) = event? else {
                    return Ok(CarrierStop::ReplicationEnded);
                };
                let mut context = EventContext {
                    settings,
                    channel,
                    client: &mut client,
                    decoder: &mut decoder,
                    feedback: &mut feedback,
                    tx: &mut tx,
                    health,
                };
                if let Some(stop) = handle_event(event, &mut context).await? {
                    return Ok(stop);
                }
            }
        }
    }
}

async fn handle_command(
    command: FeedbackCommand,
    settings: &SourceSettings,
    client: &mut ReplicationClient,
    feedback: &mut FeedbackCoordinator,
    health: &HealthState,
) -> CdcResult<Option<CarrierStop>> {
    match command {
        FeedbackCommand::Checkpoint(offset) => {
            match settings.feedback {
                FeedbackMode::AfterCheckpoint => {
                    if let Some(lsn) = feedback.checkpoint(&offset)? {
                        client.update_applied_lsn(lsn.into());
                        health.record_feedback(lsn);
                    }
                }
            }
            Ok(None)
        }
        FeedbackCommand::Stop => {
            client.stop();
            Ok(Some(CarrierStop::Stopped))
        }
        FeedbackCommand::DropSlot { force, reply } => {
            client.stop();
            let _ = client.shutdown().await;
            let result = admin::drop_slot(
                &settings.postgres,
                &settings.slot,
                settings.slot_lifecycle,
                force,
            )
            .await;
            let stop = if result.is_ok() {
                CarrierStop::DroppedSlot
            } else {
                CarrierStop::Stopped
            };
            let _ = reply.send(result);
            Ok(Some(stop))
        }
    }
}

struct EventContext<'a> {
    settings: &'a SourceSettings,
    channel: &'a Channel<ChangeEvent>,
    client: &'a mut ReplicationClient,
    decoder: &'a mut PgOutputDecoder,
    feedback: &'a mut FeedbackCoordinator,
    tx: &'a mut Option<PendingTransaction>,
    health: &'a HealthState,
}

async fn handle_event(
    event: ReplicationEvent,
    context: &mut EventContext<'_>,
) -> CdcResult<Option<CarrierStop>> {
    match event {
        ReplicationEvent::KeepAlive { wal_end, .. } => {
            context.health.record_server_wal_end(wal_end.into());
        }
        ReplicationEvent::Begin {
            final_lsn,
            xid,
            commit_time_micros,
        } => {
            *context.tx = Some(PendingTransaction {
                xid,
                final_lsn: final_lsn.into(),
                commit_time_micros,
                changes: Vec::new(),
            });
        }
        ReplicationEvent::XLogData {
            wal_start,
            wal_end,
            data,
            ..
        } => {
            context.health.record_received(wal_end.into());
            let changes = context.decoder.decode(&data)?;
            if !changes.is_empty() {
                let pending = context.tx.as_mut().ok_or_else(|| {
                    CdcError::Parse(format!(
                        "pgoutput row data at WAL {wal_start} arrived outside a transaction"
                    ))
                })?;
                pending.changes.extend(changes);
            }
        }
        ReplicationEvent::Commit {
            lsn,
            end_lsn,
            commit_time_micros,
        } => {
            let Some(mut pending) = context.tx.take() else {
                return Ok(None);
            };
            let _begin_final_lsn = pending.final_lsn;
            pending.commit_time_micros = commit_time_micros;
            let commit_lsn = PgLsn::from(lsn);
            let tx_end_lsn = PgLsn::from(end_lsn);
            let event_count = u32::try_from(pending.changes.len()).map_err(|_| {
                CdcError::Parse("transaction emitted more than u32::MAX events".into())
            })?;
            if event_count == 0 {
                // An empty transaction has nothing to checkpoint downstream, so
                // route it through the feedback coordinator: it may only advance
                // PostgreSQL feedback once every earlier transaction is
                // confirmed, otherwise the slot could skip an unconfirmed
                // transaction on restart or reconnect (at-least-once violation).
                if let Some(lsn) = context.feedback.note_empty_tx(tx_end_lsn)? {
                    context.client.update_applied_lsn(lsn.into());
                    context.health.record_feedback(lsn);
                }
                return Ok(None);
            }
            context.feedback.register_tx(tx_end_lsn, event_count)?;
            for (index, change) in pending.changes.into_iter().enumerate() {
                let event = build_change_event(
                    context.settings,
                    change,
                    CommittedEventMeta {
                        xid: pending.xid,
                        commit_time_micros: pending.commit_time_micros,
                        commit_lsn,
                        tx_end_lsn,
                        event_index: u32::try_from(index)
                            .expect("event index bounded by event_count"),
                        event_count,
                    },
                );
                if context.channel.send(event).await.is_err() {
                    return Ok(Some(CarrierStop::SourceClosed));
                }
                context.health.record_emitted();
            }
        }
        ReplicationEvent::Message { .. } | ReplicationEvent::StoppedAt { .. } => {}
    }
    Ok(None)
}

#[derive(Debug, Clone, Copy)]
struct CommittedEventMeta {
    xid: u32,
    commit_time_micros: i64,
    commit_lsn: PgLsn,
    tx_end_lsn: PgLsn,
    event_index: u32,
    event_count: u32,
}

fn build_change_event(
    settings: &SourceSettings,
    change: DecodedChange,
    meta: CommittedEventMeta,
) -> ChangeEvent {
    let offset = CdcOffset {
        slot: settings.slot.clone(),
        tx_end_lsn: meta.tx_end_lsn,
        commit_lsn: meta.commit_lsn,
        xid: meta.xid,
        event_index: meta.event_index,
        event_count: meta.event_count,
    };
    let tx = TransactionMeta {
        xid: meta.xid,
        commit_time_micros: meta.commit_time_micros,
        event_index: meta.event_index,
        event_count: meta.event_count,
    };
    ChangeEvent {
        source: settings.source.clone(),
        schema: change.relation.schema.clone(),
        table: change.relation.table.clone(),
        op: change.op,
        before: change.before,
        after: change.after,
        truncate: change.truncate,
        lsn: offset,
        tx,
        relation: change.relation,
    }
}

#[derive(Debug)]
struct PendingTransaction {
    xid: u32,
    final_lsn: PgLsn,
    commit_time_micros: i64,
    changes: Vec<DecodedChange>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::MemoryCheckpointStore;

    #[test]
    fn builder_requires_connection_slot_and_publication() {
        let err = match CdcSource::postgres().build() {
            Ok(_) => panic!("builder unexpectedly succeeded"),
            Err(err) => err,
        };
        assert!(matches!(err, CdcError::Config(_)));
    }

    #[test]
    fn builder_creates_context_source() {
        let config = PostgresCdcConfig::from_url("postgresql://datum@127.0.0.1:5432/db").unwrap();
        let source = CdcSource::postgres()
            .connect(config)
            .slot("slot")
            .publication("pub")
            .checkpoint_store(MemoryCheckpointStore::new())
            .build_with_context();
        assert!(source.is_ok());
    }
}