datum-mq 0.10.4

use std::{collections::BTreeMap, fmt, sync::Arc};

use crate::{MqError, MqResult};

/// Kafka topic/partition identifier.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct TopicPartition {
    pub topic: String,
    pub partition: i32,
}

impl TopicPartition {
    #[must_use]
    pub fn new(topic: impl Into<String>, partition: i32) -> Self {
        Self {
            topic: topic.into(),
            partition,
        }
    }
}

/// A Kafka commit position. `offset` is the next offset to consume.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OffsetCommit {
    pub topic: String,
    pub partition: i32,
    pub offset: i64,
}

impl OffsetCommit {
    #[must_use]
    pub fn topic_partition(&self) -> TopicPartition {
        TopicPartition::new(self.topic.clone(), self.partition)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) struct OffsetPosition {
    pub(crate) topic: String,
    pub(crate) partition: i32,
    pub(crate) offset: i64,
}

impl OffsetPosition {
    pub(crate) fn new(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
        Self {
            topic: topic.into(),
            partition,
            offset,
        }
    }

    pub(crate) fn topic_partition(&self) -> TopicPartition {
        TopicPartition::new(self.topic.clone(), self.partition)
    }
}

pub(crate) trait OffsetCommitter: Send + Sync {
    fn stage_offset(&self, position: OffsetPosition) -> MqResult<()>;
    fn commit_offset(&self, position: OffsetPosition) -> MqResult<()>;
}

/// Context value emitted by `KafkaSource::committable`.
#[derive(Clone)]
pub struct KafkaOffset {
    position: OffsetPosition,
    committer: Option<Arc<dyn OffsetCommitter>>,
}

impl KafkaOffset {
    pub(crate) fn new(position: OffsetPosition, committer: Arc<dyn OffsetCommitter>) -> Self {
        Self {
            position,
            committer: Some(committer),
        }
    }

    #[must_use]
    pub fn detached(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
        Self {
            position: OffsetPosition::new(topic, partition, offset),
            committer: None,
        }
    }

    #[must_use]
    pub fn topic(&self) -> &str {
        &self.position.topic
    }

    #[must_use]
    pub fn partition(&self) -> i32 {
        self.position.partition
    }

    #[must_use]
    pub fn offset(&self) -> i64 {
        self.position.offset
    }

    #[must_use]
    pub fn next_offset(&self) -> i64 {
        self.position.offset + 1
    }

    #[must_use]
    pub fn topic_partition(&self) -> TopicPartition {
        self.position.topic_partition()
    }

    /// Marks this offset as completed downstream and commits the highest contiguous
    /// completed position for the partition.
    pub fn commit(&self) -> MqResult<()> {
        let Some(committer) = &self.committer else {
            return Err(MqError::ControlUnavailable);
        };
        committer.commit_offset(self.position.clone())
    }

    /// Marks this offset as completed downstream without sending a Kafka commit.
    ///
    /// Use this for batched commits: stage earlier offsets in a contiguous batch,
    /// then call [`KafkaOffset::commit`] on the last offset after the downstream
    /// checkpoint succeeds.
    pub fn mark_processed(&self) -> MqResult<()> {
        let Some(committer) = &self.committer else {
            return Err(MqError::ControlUnavailable);
        };
        committer.stage_offset(self.position.clone())
    }
}

impl fmt::Debug for KafkaOffset {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KafkaOffset")
            .field("topic", &self.position.topic)
            .field("partition", &self.position.partition)
            .field("offset", &self.position.offset)
            .finish_non_exhaustive()
    }
}

/// Batch of offsets for external checkpointing or batched downstream commit sinks.
#[derive(Debug, Clone, Default)]
pub struct KafkaOffsetBatch {
    offsets: Vec<KafkaOffset>,
}

impl KafkaOffsetBatch {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    #[must_use]
    pub fn from_offsets<I>(offsets: I) -> Self
    where
        I: IntoIterator<Item = KafkaOffset>,
    {
        Self {
            offsets: offsets.into_iter().collect(),
        }
    }

    pub fn push(&mut self, offset: KafkaOffset) {
        self.offsets.push(offset);
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.offsets.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.offsets.is_empty()
    }

    pub fn commit(&self) -> MqResult<()> {
        let mut by_partition = BTreeMap::<TopicPartition, &KafkaOffset>::new();
        for offset in &self.offsets {
            by_partition
                .entry(offset.topic_partition())
                .and_modify(|current| {
                    if offset.offset() > current.offset() {
                        *current = offset;
                    }
                })
                .or_insert(offset);
        }
        for offset in by_partition.values() {
            offset.commit()?;
        }
        Ok(())
    }
}

/// Committable watermark for a batch of Kafka records.
///
/// `KafkaBatchOffset` is emitted by batch-oriented sources. It advances one
/// offset watermark per topic/partition after the downstream batch checkpoint
/// succeeds, avoiding one offset handle per record while preserving
/// at-least-once commit ordering.
#[derive(Clone)]
pub struct KafkaBatchOffset {
    positions: Vec<OffsetPosition>,
    committer: KafkaBatchCommitter,
}

#[derive(Clone)]
enum KafkaBatchCommitter {
    Offset(Option<Arc<dyn OffsetCommitter>>),
    Custom(Arc<dyn Fn() -> MqResult<()> + Send + Sync>),
}

impl KafkaBatchOffset {
    pub(crate) fn new(positions: Vec<OffsetPosition>, committer: Arc<dyn OffsetCommitter>) -> Self {
        Self {
            positions,
            committer: KafkaBatchCommitter::Offset(Some(committer)),
        }
    }
    pub(crate) fn custom<F>(positions: Vec<OffsetPosition>, commit: F) -> Self
    where
        F: Fn() -> MqResult<()> + Send + Sync + 'static,
    {
        Self {
            positions,
            committer: KafkaBatchCommitter::Custom(Arc::new(commit)),
        }
    }

    #[must_use]
    pub fn detached<I>(positions: I) -> Self
    where
        I: IntoIterator<Item = OffsetCommit>,
    {
        Self {
            positions: positions
                .into_iter()
                .map(|position| {
                    OffsetPosition::new(position.topic, position.partition, position.offset - 1)
                })
                .collect(),
            committer: KafkaBatchCommitter::Offset(None),
        }
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.positions.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.positions.is_empty()
    }

    /// Marks the batch as completed downstream and commits eligible Kafka
    /// watermarks according to the source committer settings.
    pub fn commit(&self) -> MqResult<()> {
        match &self.committer {
            KafkaBatchCommitter::Offset(Some(committer)) => {
                for position in &self.positions {
                    committer.commit_offset(position.clone())?;
                }
                Ok(())
            }
            KafkaBatchCommitter::Offset(None) => Err(MqError::ControlUnavailable),
            KafkaBatchCommitter::Custom(commit) => commit(),
        }
    }
}

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

#[derive(Debug, Clone, Default)]
pub(crate) struct OffsetTracker {
    partitions: BTreeMap<TopicPartition, PartitionOffsetState>,
}

#[derive(Debug, Clone, Default)]
struct PartitionOffsetState {
    emitted_next: Option<i64>,
    processed_next: Option<i64>,
    last_committed: Option<i64>,
}

impl OffsetTracker {
    pub(crate) fn observe(&mut self, position: &OffsetPosition) -> u64 {
        let state = self
            .partitions
            .entry(position.topic_partition())
            .or_default();
        let next = position.offset + 1;
        state.emitted_next = Some(state.emitted_next.map_or(next, |current| current.max(next)));
        state
            .processed_next
            .get_or_insert_with(|| position.offset.min(next));
        state.last_committed.get_or_insert(position.offset);
        self.outstanding()
    }

    pub(crate) fn process(&mut self, position: &OffsetPosition) -> (u64, u64) {
        let state = self
            .partitions
            .entry(position.topic_partition())
            .or_default();
        let next = position.offset + 1;
        state.emitted_next = Some(state.emitted_next.map_or(next, |current| current.max(next)));
        state.processed_next = Some(
            state
                .processed_next
                .map_or(next, |current| current.max(next)),
        );
        state.last_committed.get_or_insert(position.offset);
        (self.outstanding(), self.uncommitted())
    }

    pub(crate) fn due_commits(&self) -> Vec<OffsetCommit> {
        self.partitions
            .iter()
            .filter_map(|(partition, state)| {
                let processed = state.processed_next?;
                if state
                    .last_committed
                    .is_some_and(|committed| committed >= processed)
                {
                    return None;
                }
                Some(OffsetCommit {
                    topic: partition.topic.clone(),
                    partition: partition.partition,
                    offset: processed,
                })
            })
            .collect()
    }

    pub(crate) fn mark_committed(&mut self, commits: &[OffsetCommit]) -> (u64, i64) {
        for commit in commits {
            if let Some(state) = self.partitions.get_mut(&commit.topic_partition()) {
                state.last_committed = Some(
                    state
                        .last_committed
                        .map_or(commit.offset, |current| current.max(commit.offset)),
                );
            }
        }
        (self.outstanding(), self.last_committed_watermark())
    }

    pub(crate) fn outstanding(&self) -> u64 {
        self.partitions
            .values()
            .map(|state| {
                let emitted = state.emitted_next.unwrap_or(0);
                let processed = state.processed_next.unwrap_or(emitted);
                emitted.saturating_sub(processed) as u64
            })
            .sum()
    }

    pub(crate) fn uncommitted(&self) -> u64 {
        self.partitions
            .values()
            .map(|state| {
                let processed = state.processed_next.unwrap_or(0);
                let committed = state.last_committed.unwrap_or(processed);
                processed.saturating_sub(committed) as u64
            })
            .sum()
    }

    pub(crate) fn last_committed_watermark(&self) -> i64 {
        self.partitions
            .values()
            .filter_map(|state| state.last_committed)
            .max()
            .unwrap_or(-1)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;

    fn pos(offset: i64) -> OffsetPosition {
        OffsetPosition::new("topic", 0, offset)
    }

    #[test]
    fn commits_contiguous_offsets_in_order() {
        let mut tracker = OffsetTracker::default();
        tracker.observe(&pos(5));
        tracker.observe(&pos(6));

        let (outstanding, uncommitted) = tracker.process(&pos(5));
        assert_eq!(outstanding, 1);
        assert_eq!(uncommitted, 1);
        assert_eq!(tracker.due_commits()[0].offset, 6);

        let commits = tracker.due_commits();
        tracker.mark_committed(&commits);
        let (outstanding, uncommitted) = tracker.process(&pos(6));
        assert_eq!(outstanding, 0);
        assert_eq!(uncommitted, 1);
        assert_eq!(tracker.due_commits()[0].offset, 7);
    }

    #[test]
    fn watermark_processing_advances_monotonically() {
        let mut tracker = OffsetTracker::default();
        tracker.observe(&pos(10));
        tracker.observe(&pos(11));
        tracker.observe(&pos(12));

        let (outstanding, uncommitted) = tracker.process(&pos(11));
        assert_eq!(outstanding, 1);
        assert_eq!(uncommitted, 2);
        assert_eq!(tracker.due_commits()[0].offset, 12);

        let (outstanding, uncommitted) = tracker.process(&pos(10));
        assert_eq!(outstanding, 1);
        assert_eq!(uncommitted, 2);

        let (outstanding, uncommitted) = tracker.process(&pos(12));
        assert_eq!(outstanding, 0);
        assert_eq!(uncommitted, 3);
        assert_eq!(tracker.due_commits()[0].offset, 13);
    }

    #[derive(Debug, Default)]
    struct FakeCommitter {
        staged: Mutex<Vec<OffsetPosition>>,
        committed: Mutex<Vec<OffsetPosition>>,
    }

    impl OffsetCommitter for FakeCommitter {
        fn stage_offset(&self, position: OffsetPosition) -> MqResult<()> {
            self.staged.lock().expect("staged lock").push(position);
            Ok(())
        }

        fn commit_offset(&self, position: OffsetPosition) -> MqResult<()> {
            self.committed
                .lock()
                .expect("committed lock")
                .push(position);
            Ok(())
        }
    }

    #[test]
    fn offset_batch_commits_last_offset_per_partition() {
        let committer = Arc::new(FakeCommitter::default());
        let offsets = [
            KafkaOffset::new(pos(0), committer.clone()),
            KafkaOffset::new(OffsetPosition::new("topic", 1, 0), committer.clone()),
            KafkaOffset::new(pos(1), committer.clone()),
            KafkaOffset::new(OffsetPosition::new("topic", 1, 1), committer.clone()),
        ];

        KafkaOffsetBatch::from_offsets(offsets)
            .commit()
            .expect("batch commits");

        let staged = committer.staged.lock().expect("staged lock");
        let committed = committer.committed.lock().expect("committed lock");
        assert_eq!(staged.len(), 0);
        assert_eq!(committed.len(), 2);
        assert!(committed.iter().all(|position| position.offset == 1));
    }

    #[test]
    fn kafka_batch_offset_commits_partition_watermarks() {
        let committer = Arc::new(FakeCommitter::default());
        let batch = KafkaBatchOffset::new(
            vec![pos(11), OffsetPosition::new("topic", 1, 21)],
            committer.clone(),
        );

        batch.commit().expect("batch watermark commits");

        let staged = committer.staged.lock().expect("staged lock");
        let committed = committer.committed.lock().expect("committed lock");
        assert_eq!(staged.len(), 0);
        assert_eq!(
            committed.as_slice(),
            &[pos(11), OffsetPosition::new("topic", 1, 21)]
        );
    }
}