datum-mq 0.11.2

use std::sync::{Arc, Mutex};

use datum::{Flow, Keep, NotUsed, Sink, StreamCompletion, StreamError};

use crate::{
    KafkaMetrics, KafkaProducerSettings, MqError, MqResult, ProducerRecord,
    native::{NativeKafkaProducerControl, NativeKafkaProducerHandle},
};

const PRODUCER_BATCH_SIZE: usize = 256;

/// Kafka producer sink entry points.
pub struct KafkaSink;

impl KafkaSink {
    /// Creates a producer sink.
    ///
    /// The input record carries its target topic. The sink completes each element
    /// by enqueuing bounded 256-record groups into the native producer owner.
    /// Delivery failures fail the stream; call
    /// [`KafkaProducerControl::drain_and_shutdown`] at shutdown to wait for final
    /// acknowledgements before closing. Production uses record-batch v2 plus
    /// bounded, idempotent retries by default, preserving the no-duplicate
    /// guarantee. Its bounded per-broker pipeline serializes each topic-partition
    /// and halts and settles issued responses before retrying.
    #[must_use]
    pub fn plain(settings: KafkaProducerSettings) -> Sink<ProducerRecord, KafkaProducerControl> {
        Flow::identity()
            .grouped(PRODUCER_BATCH_SIZE)
            .to_mat(Self::batched(settings), Keep::right)
    }

    fn batched(settings: KafkaProducerSettings) -> Sink<Vec<ProducerRecord>, KafkaProducerControl> {
        Sink::setup(move |_materializer, _attributes| {
            let metrics = KafkaMetrics::default();
            let (native, handle) = NativeKafkaProducerControl::start(&settings, metrics.clone())
                .expect("native Kafka producer configuration must be valid at materialization");
            let control = KafkaProducerControl::new(native, metrics);
            Sink::foreach_result(move |records| {
                let handle: NativeKafkaProducerHandle = handle.clone();
                handle.send(records).map_err(StreamError::from)
            })
            .map_materialized_value(move |completion| {
                control.attach_completion(completion);
                control.clone()
            })
        })
    }
}

/// Materialized control for a Kafka producer sink.
#[derive(Clone)]
pub struct KafkaProducerControl {
    state: Arc<ProducerControlState>,
}

struct ProducerControlState {
    native: NativeKafkaProducerControl,
    metrics: KafkaMetrics,
    completion: Mutex<Option<StreamCompletion<NotUsed>>>,
}

impl KafkaProducerControl {
    fn new(native: NativeKafkaProducerControl, metrics: KafkaMetrics) -> Self {
        Self {
            state: Arc::new(ProducerControlState {
                native,
                metrics,
                completion: Mutex::new(None),
            }),
        }
    }

    fn attach_completion(&self, completion: StreamCompletion<NotUsed>) {
        if let Ok(mut slot) = self.state.completion.lock() {
            *slot = Some(completion);
        }
    }

    #[must_use]
    pub fn metrics(&self) -> KafkaMetrics {
        self.state.metrics.clone()
    }

    /// Waits for accepted input records to receive delivery acknowledgements, then flushes.
    pub fn drain_and_shutdown(&self) -> MqResult<()> {
        let completion = self
            .state
            .completion
            .lock()
            .map_err(|_| MqError::Failed("Kafka producer completion lock poisoned".to_owned()))?
            .take();
        let completion_result = completion.map_or(Ok(()), |completion| {
            completion
                .wait()
                .map(|_| ())
                .map_err(|error| MqError::Failed(error.to_string()))
        });
        completion_result.and(self.state.native.finish())
    }

    pub fn flush(&self) -> MqResult<()> {
        self.state.native.flush()
    }
}