datum-mq 0.11.1

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 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()
    }
}

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

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

    #[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));
    }
}