datum-mq 0.10.9

use bytes::Bytes;
#[cfg(feature = "rdkafka")]
use rdkafka::message::{BorrowedMessage, Headers, Message, Timestamp};

/// Kafka record header.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KafkaHeader {
    pub key: String,
    pub value: Option<Bytes>,
}

impl KafkaHeader {
    #[must_use]
    pub fn new(key: impl Into<String>, value: impl Into<Option<Bytes>>) -> Self {
        Self {
            key: key.into(),
            value: value.into(),
        }
    }
}

/// Kafka record timestamp.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KafkaTimestamp {
    NotAvailable,
    CreateTime(i64),
    LogAppendTime(i64),
}

#[cfg(feature = "rdkafka")]
impl From<Timestamp> for KafkaTimestamp {
    fn from(timestamp: Timestamp) -> Self {
        match timestamp {
            Timestamp::NotAvailable => Self::NotAvailable,
            Timestamp::CreateTime(value) => Self::CreateTime(value),
            Timestamp::LogAppendTime(value) => Self::LogAppendTime(value),
        }
    }
}

/// Owned Kafka consumer record.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConsumerRecord {
    pub topic: String,
    pub partition: i32,
    pub offset: i64,
    pub timestamp: KafkaTimestamp,
    pub key: Option<Bytes>,
    pub payload: Option<Bytes>,
    pub headers: Vec<KafkaHeader>,
}

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

    #[cfg(feature = "rdkafka")]
    pub(crate) fn from_borrowed(message: &BorrowedMessage<'_>) -> Self {
        let headers = message
            .headers()
            .map(|headers| {
                (0..headers.count())
                    .map(|index| {
                        let header = headers.get(index);
                        KafkaHeader {
                            key: header.key.to_owned(),
                            value: header.value.map(Bytes::copy_from_slice),
                        }
                    })
                    .collect()
            })
            .unwrap_or_default();

        Self {
            topic: message.topic().to_owned(),
            partition: message.partition(),
            offset: message.offset(),
            timestamp: message.timestamp().into(),
            key: message.key().map(Bytes::copy_from_slice),
            payload: message.payload().map(Bytes::copy_from_slice),
            headers,
        }
    }
}

/// Owned Kafka producer record.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProducerRecord {
    pub topic: String,
    pub key: Option<Bytes>,
    pub payload: Option<Bytes>,
    pub partition: Option<i32>,
    pub timestamp: Option<i64>,
    pub headers: Vec<KafkaHeader>,
}

impl ProducerRecord {
    #[must_use]
    pub fn new(topic: impl Into<String>, payload: impl Into<Option<Bytes>>) -> Self {
        Self {
            topic: topic.into(),
            key: None,
            payload: payload.into(),
            partition: None,
            timestamp: None,
            headers: Vec::new(),
        }
    }

    #[must_use]
    pub fn with_key(mut self, key: impl Into<Option<Bytes>>) -> Self {
        self.key = key.into();
        self
    }

    #[must_use]
    pub fn with_partition(mut self, partition: i32) -> Self {
        self.partition = Some(partition);
        self
    }

    #[must_use]
    pub fn with_timestamp(mut self, timestamp: i64) -> Self {
        self.timestamp = Some(timestamp);
        self
    }

    #[must_use]
    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<Option<Bytes>>) -> Self {
        self.headers.push(KafkaHeader::new(key, value));
        self
    }
}