datum-mq 0.10.5

use std::{collections::BTreeMap, time::Duration};

#[cfg(feature = "rdkafka")]
use rdkafka::ClientConfig;
#[cfg(feature = "rdkafka")]
pub use rdkafka::Offset;

/// Kafka assignment start offset used when the `rdkafka` feature is disabled.
#[cfg(not(feature = "rdkafka"))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Offset {
    Beginning,
    End,
    Stored,
    Offset(i64),
    OffsetTail(i64),
    Invalid,
}

use crate::{MqError, MqResult, TopicPartition};

/// Raw Kafka property map with typed helpers for common Datum defaults.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KafkaConfig {
    properties: BTreeMap<String, String>,
}

impl KafkaConfig {
    #[must_use]
    pub fn new(bootstrap_servers: impl Into<String>) -> Self {
        let mut properties = BTreeMap::new();
        properties.insert("bootstrap.servers".to_owned(), bootstrap_servers.into());
        Self { properties }
    }

    #[must_use]
    pub fn empty() -> Self {
        Self {
            properties: BTreeMap::new(),
        }
    }

    #[must_use]
    pub fn get(&self, key: &str) -> Option<&str> {
        self.properties.get(key).map(String::as_str)
    }

    #[must_use]
    pub fn properties(&self) -> &BTreeMap<String, String> {
        &self.properties
    }

    pub fn set(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
        self.properties.insert(key.into(), value.into());
        self
    }

    #[must_use]
    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.set(key, value);
        self
    }

    /// Sets the common TLS properties understood by both backends.
    #[must_use]
    pub fn with_tls(
        mut self,
        ca_location: impl Into<String>,
        certificate_location: impl Into<String>,
        key_location: impl Into<String>,
    ) -> Self {
        self.set("security.protocol", "ssl");
        self.set("ssl.ca.location", ca_location);
        self.set("ssl.certificate.location", certificate_location);
        self.set("ssl.key.location", key_location);
        self
    }

    /// Sets SASL/SSL PLAIN properties.
    ///
    /// This is a property helper only. The selected backend must support the
    /// requested SASL mechanism.
    #[must_use]
    pub fn with_sasl_plain(
        mut self,
        username: impl Into<String>,
        password: impl Into<String>,
    ) -> Self {
        self.set("security.protocol", "sasl_ssl");
        self.set("sasl.mechanism", "PLAIN");
        self.set("sasl.username", username);
        self.set("sasl.password", password);
        self
    }

    #[cfg(feature = "rdkafka")]
    pub(crate) fn to_client_config(&self) -> ClientConfig {
        let mut config = ClientConfig::new();
        for (key, value) in &self.properties {
            config.set(key, value);
        }
        config
    }
}

/// Offset commit mode for `KafkaSource`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommitPolicy {
    /// Datum-managed at-least-once commits. This is the default.
    Manual,
    /// Use Kafka auto-commit. This is an escape hatch and may lose records.
    AutoKafka,
    /// Downstream stores offsets externally; Kafka commits are optional.
    External,
}

/// Consumer implementation used by `KafkaSource`.
///
/// `Native` is the default. `Rdkafka` remains selectable when the default-on
/// `rdkafka` Cargo feature is enabled. The producer has a separate
/// [`KafkaProducerBackend`] selector.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum KafkaConsumerBackend {
    #[default]
    Native,
    #[cfg(feature = "rdkafka")]
    Rdkafka,
}

/// Producer implementation used by [`crate::KafkaSink`].
///
/// `Rdkafka` remains the default when the default-on `rdkafka` Cargo feature is
/// enabled. Without that feature, `Native` is the only available backend and
/// therefore the default. Native production provides at-least-once delivery
/// with bounded retries, but does not yet implement Kafka idempotent production.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum KafkaProducerBackend {
    #[cfg(feature = "rdkafka")]
    #[default]
    Rdkafka,
    #[cfg_attr(not(feature = "rdkafka"), default)]
    Native,
}

/// Topic subscription mode.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Subscription {
    Topics(Vec<String>),
    Pattern(String),
    Assignment(Vec<TopicPartitionOffset>),
}

impl Subscription {
    #[must_use]
    pub fn topics<I, S>(topics: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        Self::Topics(topics.into_iter().map(Into::into).collect())
    }

    #[must_use]
    pub fn pattern(pattern: impl Into<String>) -> Self {
        let pattern = pattern.into();
        if pattern.starts_with('^') {
            Self::Pattern(pattern)
        } else {
            Self::Pattern(format!("^{pattern}"))
        }
    }

    #[must_use]
    pub fn assignment<I>(partitions: I) -> Self
    where
        I: IntoIterator<Item = TopicPartitionOffset>,
    {
        Self::Assignment(partitions.into_iter().collect())
    }

    pub(crate) fn validate(&self) -> MqResult<()> {
        match self {
            Subscription::Topics(topics) if topics.is_empty() => Err(MqError::InvalidConfig(
                "Kafka subscription requires at least one topic".to_owned(),
            )),
            Subscription::Pattern(pattern) if pattern.is_empty() || pattern == "^" => {
                Err(MqError::InvalidConfig(
                    "Kafka pattern subscription requires a non-empty pattern".to_owned(),
                ))
            }
            Subscription::Assignment(partitions) if partitions.is_empty() => {
                Err(MqError::InvalidConfig(
                    "Kafka assignment requires at least one partition".to_owned(),
                ))
            }
            _ => Ok(()),
        }
    }
}

/// A concrete topic/partition plus the starting offset for assignment subscriptions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TopicPartitionOffset {
    pub topic: String,
    pub partition: i32,
    pub offset: Offset,
}

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

    #[must_use]
    pub fn beginning(topic: impl Into<String>, partition: i32) -> Self {
        Self::new(topic, partition, Offset::Beginning)
    }

    #[must_use]
    pub fn end(topic: impl Into<String>, partition: i32) -> Self {
        Self::new(topic, partition, Offset::End)
    }

    #[must_use]
    pub fn absolute(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
        Self::new(topic, partition, Offset::Offset(offset))
    }

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

/// Consumer settings for `KafkaSource`.
#[derive(Debug, Clone)]
pub struct KafkaConsumerSettings {
    pub config: KafkaConfig,
    pub group_id: String,
    pub consumer_backend: KafkaConsumerBackend,
    pub commit_policy: CommitPolicy,
    pub high_watermark: usize,
    pub low_watermark: usize,
    pub drain_timeout: Duration,
    pub poll_timeout: Duration,
    pub poll_batch_size: usize,
    pub commit_batch_size: usize,
    pub commit_interval: Duration,
    pub low_latency: bool,
    pub commit_sync: bool,
    pub include_payload_timestamps: bool,
}

impl KafkaConsumerSettings {
    #[must_use]
    pub fn new(bootstrap_servers: impl Into<String>, group_id: impl Into<String>) -> Self {
        Self {
            config: KafkaConfig::new(bootstrap_servers),
            group_id: group_id.into(),
            consumer_backend: KafkaConsumerBackend::Native,
            commit_policy: CommitPolicy::Manual,
            high_watermark: 4096,
            low_watermark: 2048,
            drain_timeout: Duration::from_secs(30),
            poll_timeout: Duration::from_millis(100),
            poll_batch_size: 256,
            commit_batch_size: 10_000,
            commit_interval: Duration::from_millis(100),
            low_latency: false,
            commit_sync: false,
            include_payload_timestamps: true,
        }
    }

    #[must_use]
    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.config.set(key, value);
        self
    }

    #[must_use]
    pub fn with_commit_policy(mut self, policy: CommitPolicy) -> Self {
        self.commit_policy = policy;
        self
    }

    #[must_use]
    pub fn with_consumer_backend(mut self, backend: KafkaConsumerBackend) -> Self {
        self.consumer_backend = backend;
        self
    }

    #[must_use]
    pub fn with_backpressure(mut self, low_watermark: usize, high_watermark: usize) -> Self {
        assert!(
            low_watermark < high_watermark,
            "low watermark must be below high watermark"
        );
        self.low_watermark = low_watermark;
        self.high_watermark = high_watermark;
        self
    }

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

    #[must_use]
    pub fn with_poll_batch_size(mut self, batch_size: usize) -> Self {
        assert!(batch_size > 0, "poll batch size must be greater than zero");
        self.poll_batch_size = batch_size;
        self
    }

    #[must_use]
    pub fn with_commit_batch_size(mut self, batch_size: usize) -> Self {
        assert!(
            batch_size > 0,
            "commit batch size must be greater than zero"
        );
        self.commit_batch_size = batch_size;
        self
    }

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

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

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

    #[cfg(feature = "rdkafka")]
    pub(crate) fn to_client_config(&self) -> ClientConfig {
        let mut config = self.config.to_client_config();
        config.set("group.id", &self.group_id);
        match self.commit_policy {
            CommitPolicy::AutoKafka => {
                config.set("enable.auto.commit", "true");
                config.set("enable.auto.offset.store", "true");
            }
            CommitPolicy::Manual | CommitPolicy::External => {
                config.set("enable.auto.commit", "false");
                config.set("enable.auto.offset.store", "false");
            }
        }
        config.set("enable.partition.eof", "false");
        config
    }
}

/// Producer settings for `KafkaSink`.
#[derive(Debug, Clone)]
pub struct KafkaProducerSettings {
    pub config: KafkaConfig,
    pub producer_backend: KafkaProducerBackend,
    pub in_flight_limit: usize,
    pub queue_timeout: Duration,
    pub flush_timeout: Duration,
    pub drain_timeout: Duration,
}

impl KafkaProducerSettings {
    #[must_use]
    pub fn new(bootstrap_servers: impl Into<String>) -> Self {
        let settings = Self {
            config: KafkaConfig::new(bootstrap_servers)
                .with("enable.idempotence", "true")
                .with("acks", "all")
                .with("retries", "2147483647")
                .with("max.in.flight.requests.per.connection", "5")
                .with("linger.ms", "5")
                .with("batch.num.messages", "10000")
                .with("batch.size", "131072")
                .with("queue.buffering.max.messages", "1000000")
                .with("queue.buffering.max.kbytes", "1048576"),
            producer_backend: KafkaProducerBackend::default(),
            in_flight_limit: 65_536,
            queue_timeout: Duration::from_secs(30),
            flush_timeout: Duration::from_secs(30),
            drain_timeout: Duration::from_secs(30),
        };
        if settings.producer_backend == KafkaProducerBackend::Native {
            settings.with_producer_backend(KafkaProducerBackend::Native)
        } else {
            settings
        }
    }

    #[must_use]
    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.config.set(key, value);
        self
    }

    /// Selects the producer backend.
    ///
    /// Selecting the native backend replaces the rdkafka idempotence defaults
    /// with the native producer's bounded at-least-once defaults. Setting
    /// `enable.idempotence=true` after this call is rejected at materialization
    /// because native idempotent produce is not implemented yet. Native
    /// Produce requests default to one in flight per broker connection;
    /// requests for the same topic-partition are always serialized.
    #[must_use]
    pub fn with_producer_backend(mut self, backend: KafkaProducerBackend) -> Self {
        self.producer_backend = backend;
        if backend == KafkaProducerBackend::Native {
            self.config.set("enable.idempotence", "false");
            self.config.set("retries", "3");
            self.config
                .set("max.in.flight.requests.per.connection", "1");
        }
        self
    }

    /// Sets the maximum number of native Produce requests in flight on each
    /// broker connection.
    ///
    /// The native backend preserves per-partition order by never overlapping
    /// two requests that touch the same topic-partition. On a retriable
    /// failure it halts that connection's pipeline, settles every issued
    /// response, and retries failed requests in their original order. A lost
    /// response can still duplicate an accepted batch because native
    /// idempotent produce is not implemented. The value must be non-zero.
    #[must_use]
    pub fn with_max_in_flight_requests_per_connection(mut self, limit: usize) -> Self {
        assert!(
            limit > 0,
            "max in-flight requests must be greater than zero"
        );
        self.config
            .set("max.in.flight.requests.per.connection", limit.to_string());
        self
    }

    #[must_use]
    pub fn with_in_flight_limit(mut self, in_flight_limit: usize) -> Self {
        assert!(
            in_flight_limit > 0,
            "in-flight limit must be greater than zero"
        );
        self.in_flight_limit = in_flight_limit;
        self
    }

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

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

    #[cfg(feature = "rdkafka")]
    pub(crate) fn to_client_config(&self) -> ClientConfig {
        self.config.to_client_config()
    }
}

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

    #[test]
    fn raw_properties_are_preserved() {
        let config = KafkaConfig::new("127.0.0.1:9092")
            .with("client.id", "datum-test")
            .with("linger.ms", "5");

        assert_eq!(config.get("bootstrap.servers"), Some("127.0.0.1:9092"));
        assert_eq!(config.get("client.id"), Some("datum-test"));
        assert_eq!(config.get("linger.ms"), Some("5"));
    }

    #[test]
    #[cfg(feature = "rdkafka")]
    fn consumer_manual_commit_disables_auto_commit_and_store() {
        let settings = KafkaConsumerSettings::new("localhost:9092", "group");
        let config = settings.to_client_config();
        assert_eq!(config.get("enable.auto.commit"), Some("false"));
        assert_eq!(config.get("enable.auto.offset.store"), Some("false"));
        assert_eq!(config.get("group.id"), Some("group"));
    }

    #[test]
    fn consumer_backend_defaults_to_native() {
        let settings = KafkaConsumerSettings::new("localhost:9092", "group");
        assert_eq!(settings.consumer_backend, KafkaConsumerBackend::Native);
    }

    #[test]
    fn consumer_backend_can_be_selected_explicitly() {
        let settings = KafkaConsumerSettings::new("localhost:9092", "group")
            .with_consumer_backend(KafkaConsumerBackend::Native);
        assert_eq!(settings.consumer_backend, KafkaConsumerBackend::Native);
    }

    #[test]
    fn native_producer_backend_uses_bounded_non_idempotent_defaults() {
        let settings = KafkaProducerSettings::new("localhost:9092")
            .with_producer_backend(KafkaProducerBackend::Native);
        assert_eq!(settings.producer_backend, KafkaProducerBackend::Native);
        assert_eq!(settings.config.get("enable.idempotence"), Some("false"));
        assert_eq!(settings.config.get("retries"), Some("3"));
        assert_eq!(
            settings.config.get("max.in.flight.requests.per.connection"),
            Some("1")
        );
    }

    #[test]
    fn producer_backend_default_matches_compiled_backends() {
        let settings = KafkaProducerSettings::new("localhost:9092");
        #[cfg(feature = "rdkafka")]
        assert_eq!(settings.producer_backend, KafkaProducerBackend::Rdkafka);
        #[cfg(not(feature = "rdkafka"))]
        assert_eq!(settings.producer_backend, KafkaProducerBackend::Native);
    }

    #[test]
    fn tls_and_sasl_helpers_set_librdkafka_properties() {
        let tls = KafkaConfig::new("broker").with_tls("/ca.pem", "/cert.pem", "/key.pem");
        assert_eq!(tls.get("security.protocol"), Some("ssl"));
        assert_eq!(tls.get("ssl.ca.location"), Some("/ca.pem"));

        let sasl = KafkaConfig::new("broker").with_sasl_plain("user", "pass");
        assert_eq!(sasl.get("security.protocol"), Some("sasl_ssl"));
        assert_eq!(sasl.get("sasl.mechanism"), Some("PLAIN"));
        assert_eq!(sasl.get("sasl.username"), Some("user"));
        assert_eq!(sasl.get("sasl.password"), Some("pass"));
    }

    #[test]
    fn subscription_validation_rejects_empty_inputs() {
        assert!(
            Subscription::topics(Vec::<String>::new())
                .validate()
                .is_err()
        );
        assert!(Subscription::pattern("").validate().is_err());
        assert!(Subscription::assignment(Vec::new()).validate().is_err());
    }
}