datum-sql 0.10.2

use std::collections::BTreeMap;
use std::sync::Arc;

use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use arrow_json::{LineDelimitedWriter, ReaderBuilder};
use datafusion::common::{DataFusionError, Result};
use datum::{NotUsed, Source};
use datum_mq::{
    KafkaConsumerSettings, KafkaControl, KafkaPayloadBatch, KafkaPayloadRecord, KafkaSource,
    Subscription, TopicPartition,
};

use super::{BatchEnvelope, EnvelopedRecordBatch};
use crate::{CommittableRecordBatch, SourceCommit, SqlSourcePosition};

/// One Kafka partition range covered by an emitted Arrow batch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KafkaPartitionOffset {
    /// Kafka topic name.
    pub topic: String,
    /// Kafka partition number.
    pub partition: i32,
    /// First offset included in the batch.
    pub first_offset: i64,
    /// Last offset included in the batch.
    pub last_offset: i64,
}

impl KafkaPartitionOffset {
    /// Returns the next offset after this batch range.
    #[must_use]
    pub fn next_offset(&self) -> i64 {
        self.last_offset + 1
    }
}

/// Kafka topic/partition assigned to the SQL source.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct KafkaTopicPartition {
    pub topic: String,
    pub partition: i32,
}

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

/// Source position for one Kafka-backed SQL batch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KafkaSourcePosition {
    offsets: Vec<KafkaPartitionOffset>,
    row_partitions: Vec<i32>,
    active_partitions: Vec<KafkaTopicPartition>,
}

impl KafkaSourcePosition {
    /// Creates a source position from partition offset ranges.
    #[must_use]
    pub fn from_offsets<I>(offsets: I) -> Self
    where
        I: IntoIterator<Item = KafkaPartitionOffset>,
    {
        let mut offsets = offsets.into_iter().collect::<Vec<_>>();
        offsets.sort_by(|left, right| {
            left.topic
                .cmp(&right.topic)
                .then(left.partition.cmp(&right.partition))
        });
        let active_partitions = offsets
            .iter()
            .map(|offset| KafkaTopicPartition::new(offset.topic.clone(), offset.partition))
            .collect();
        Self {
            offsets,
            row_partitions: Vec::new(),
            active_partitions,
        }
    }

    #[must_use]
    pub fn from_offsets_with_row_partitions<I, R, A>(
        offsets: I,
        row_partitions: R,
        active_partitions: A,
    ) -> Self
    where
        I: IntoIterator<Item = KafkaPartitionOffset>,
        R: IntoIterator<Item = i32>,
        A: IntoIterator<Item = KafkaTopicPartition>,
    {
        let mut position = Self::from_offsets(offsets);
        position.row_partitions = row_partitions.into_iter().collect();
        position.active_partitions = active_partitions.into_iter().collect();
        position.active_partitions.sort();
        position.active_partitions.dedup();
        position
    }

    /// Builds a source position from a `datum-mq` payload batch record list.
    #[must_use]
    pub fn from_payload_records(topic: &str, records: &[KafkaPayloadRecord]) -> Self {
        Self::from_payload_records_and_partitions(topic, records, Vec::new())
    }

    #[must_use]
    pub fn from_payload_batch(topic: &str, batch: &KafkaPayloadBatch) -> Self {
        Self::from_payload_records_and_partitions(
            topic,
            batch.records(),
            batch.active_partitions().to_vec(),
        )
    }

    fn from_payload_records_and_partitions(
        topic: &str,
        records: &[KafkaPayloadRecord],
        active_partitions: Vec<TopicPartition>,
    ) -> Self {
        let mut by_partition = BTreeMap::<i32, (i64, i64)>::new();
        for record in records {
            by_partition
                .entry(record.partition)
                .and_modify(|range| {
                    range.0 = range.0.min(record.offset);
                    range.1 = range.1.max(record.offset);
                })
                .or_insert((record.offset, record.offset));
        }
        let position = Self::from_offsets(by_partition.into_iter().map(|(partition, range)| {
            KafkaPartitionOffset {
                topic: topic.to_owned(),
                partition,
                first_offset: range.0,
                last_offset: range.1,
            }
        }));
        let row_partitions = records
            .iter()
            .map(|record| record.partition)
            .collect::<Vec<_>>();
        let mut active_partitions = active_partitions
            .into_iter()
            .filter(|partition| partition.topic == topic)
            .map(|partition| KafkaTopicPartition::new(partition.topic, partition.partition))
            .collect::<Vec<_>>();
        if active_partitions.is_empty() {
            active_partitions = position
                .offsets
                .iter()
                .map(|offset| KafkaTopicPartition::new(offset.topic.clone(), offset.partition))
                .collect();
        }
        active_partitions.sort();
        active_partitions.dedup();
        Self {
            offsets: position.offsets,
            row_partitions,
            active_partitions,
        }
    }

    /// Returns the partition ranges covered by the batch.
    #[must_use]
    pub fn offsets(&self) -> &[KafkaPartitionOffset] {
        &self.offsets
    }

    /// Returns true when the batch contains no offsets.
    #[must_use]
    pub fn row_partitions(&self) -> &[i32] {
        &self.row_partitions
    }

    #[must_use]
    pub fn active_partitions(&self) -> &[KafkaTopicPartition] {
        &self.active_partitions
    }

    #[must_use]
    pub fn has_event_time_partition_mapping(&self) -> bool {
        !self.row_partitions.is_empty() || self.offsets.len() == 1
    }

    #[must_use]
    pub fn partition_for_row(&self, row: usize) -> Option<i32> {
        if let Some(partition) = self.row_partitions.get(row) {
            return Some(*partition);
        }
        match self.offsets.as_slice() {
            [offset] => Some(offset.partition),
            _ => None,
        }
    }

    #[must_use]
    pub fn slice_rows(&self, offset: usize, len: usize) -> Self {
        let row_partitions = if self.row_partitions.is_empty() {
            Vec::new()
        } else {
            self.row_partitions
                .iter()
                .skip(offset)
                .take(len)
                .copied()
                .collect()
        };
        Self {
            offsets: self.offsets.clone(),
            row_partitions,
            active_partitions: self.active_partitions.clone(),
        }
    }

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

/// Format seam for MQ payload batches.
pub trait MqPayloadFormat: Clone + Send + Sync + 'static {
    /// Returns the Arrow schema produced by this format.
    fn schema(&self) -> SchemaRef;
    /// Returns the schema revision attached to emitted batches.
    fn schema_revision(&self) -> u64;
    /// Decodes raw payload bytes into an Arrow batch.
    fn decode_payloads(&self, payloads: &[&[u8]]) -> Result<RecordBatch>;
}

/// JSON object payloads decoded into a user-supplied Arrow schema.
#[derive(Debug, Clone)]
pub struct JsonRowFormat {
    schema: SchemaRef,
    schema_revision: u64,
    strict_mode: bool,
}

impl JsonRowFormat {
    /// Creates a JSON-row format with schema revision zero and non-strict mode.
    #[must_use]
    pub fn new(schema: SchemaRef) -> Self {
        Self {
            schema,
            schema_revision: 0,
            strict_mode: false,
        }
    }

    /// Sets the schema revision attached to decoded batches.
    #[must_use]
    pub fn with_schema_revision(mut self, schema_revision: u64) -> Self {
        self.schema_revision = schema_revision;
        self
    }

    /// Enables or disables Arrow JSON strict mode.
    #[must_use]
    pub fn with_strict_mode(mut self, strict_mode: bool) -> Self {
        self.strict_mode = strict_mode;
        self
    }

    /// Decodes line-delimited JSON object payloads into one Arrow batch.
    pub fn decode_payload_slices<I, P>(&self, payloads: I) -> Result<RecordBatch>
    where
        I: IntoIterator<Item = P>,
        P: AsRef<[u8]>,
    {
        let payloads = payloads.into_iter().collect::<Vec<_>>();
        if payloads.is_empty() {
            return Ok(RecordBatch::new_empty(Arc::clone(&self.schema)));
        }

        let mut json = Vec::with_capacity(
            payloads
                .iter()
                .map(|payload| payload.as_ref().len() + 1)
                .sum(),
        );
        for payload in &payloads {
            json.extend_from_slice(payload.as_ref());
            json.push(b'\n');
        }

        let mut decoder = ReaderBuilder::new(Arc::clone(&self.schema))
            .with_batch_size(payloads.len())
            .with_strict_mode(self.strict_mode)
            .build_decoder()
            .map_err(DataFusionError::from)?;
        let decoded = decoder.decode(&json).map_err(DataFusionError::from)?;
        if decoded != json.len() {
            return Err(DataFusionError::Plan(format!(
                "JSON MQ decoder consumed {decoded} of {} bytes",
                json.len()
            )));
        }
        decoder
            .flush()
            .map_err(DataFusionError::from)?
            .ok_or_else(|| DataFusionError::Plan("JSON MQ decoder produced no rows".into()))
    }

    /// Encodes each row of an Arrow batch as one JSON object payload.
    pub fn encode_record_batch(&self, batch: &RecordBatch) -> Result<Vec<Vec<u8>>> {
        if batch.num_rows() == 0 {
            return Ok(Vec::new());
        }

        if batch.schema().fields() != self.schema.fields() {
            return Err(DataFusionError::Plan(format!(
                "JSON MQ encoder expected schema {:?}, found {:?}",
                self.schema,
                batch.schema()
            )));
        }

        let mut json = Vec::new();
        let mut writer = LineDelimitedWriter::new(&mut json);
        writer
            .write_batches(&[batch])
            .map_err(DataFusionError::from)?;
        writer.finish().map_err(DataFusionError::from)?;
        drop(writer);

        Ok(json
            .split(|byte| *byte == b'\n')
            .filter(|line| !line.is_empty())
            .map(<[u8]>::to_vec)
            .collect())
    }
}

impl MqPayloadFormat for JsonRowFormat {
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    fn schema_revision(&self) -> u64 {
        self.schema_revision
    }

    fn decode_payloads(&self, payloads: &[&[u8]]) -> Result<RecordBatch> {
        self.decode_payload_slices(payloads.iter().copied())
    }
}

/// Kafka-backed Arrow batch with source position metadata.
pub type KafkaRecordBatch = EnvelopedRecordBatch<KafkaSourcePosition>;
/// Committable Kafka-backed SQL batch.
pub type KafkaCommittableRecordBatch = CommittableRecordBatch;

impl crate::SqlEventPayload for KafkaRecordBatch {
    fn event_time_batch(&self) -> &RecordBatch {
        self.batch()
    }

    fn event_time_partition(&self, row: usize) -> Option<i64> {
        self.envelope()
            .source_position()
            .partition_for_row(row)
            .map(i64::from)
    }

    fn event_time_active_partitions(&self) -> Option<Vec<i64>> {
        if !self
            .envelope()
            .source_position()
            .has_event_time_partition_mapping()
        {
            return None;
        }
        let partitions = self
            .envelope()
            .source_position()
            .active_partitions()
            .iter()
            .map(|partition| i64::from(partition.partition))
            .collect::<Vec<_>>();
        (!partitions.is_empty()).then_some(partitions)
    }
}

#[must_use]
/// Creates a Kafka JSON source that preserves source position metadata.
pub fn kafka_json_source(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: JsonRowFormat,
) -> Source<KafkaRecordBatch, KafkaControl> {
    kafka_source(settings, topic, format)
}

#[must_use]
/// Creates a Kafka JSON source that emits plain Arrow batches.
pub fn kafka_json_record_batch_source(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: JsonRowFormat,
) -> Source<RecordBatch, KafkaControl> {
    kafka_json_source(settings, topic, format).map(|batch| batch.into_batch())
}

#[must_use]
/// Creates a Kafka JSON source with the materialized Kafka control discarded.
pub fn kafka_json_record_batch_source_uncontrolled(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: JsonRowFormat,
) -> Source<RecordBatch, NotUsed> {
    kafka_json_record_batch_source(settings, topic, format).map_materialized_value(|_| NotUsed)
}

#[must_use]
/// Creates a committable Kafka JSON source for SQL sink pipelines.
pub fn kafka_json_committable_source(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: JsonRowFormat,
) -> Source<KafkaCommittableRecordBatch, KafkaControl> {
    let topic = topic.into();
    KafkaSource::committable_payload_batches(settings, Subscription::topics([topic.clone()]))
        .try_map(move |batch| {
            decode_kafka_committable_batch(&topic, &format, &batch).map_err(crate::stream_error)
        })
}

#[must_use]
/// Creates a committable Kafka JSON source and discards the materialized control.
pub fn kafka_json_committable_source_uncontrolled(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: JsonRowFormat,
) -> Source<KafkaCommittableRecordBatch, NotUsed> {
    kafka_json_committable_source(settings, topic, format).map_materialized_value(|_| NotUsed)
}

#[must_use]
/// Creates a Kafka source with a custom payload format.
pub fn kafka_source<F>(
    settings: KafkaConsumerSettings,
    topic: impl Into<String>,
    format: F,
) -> Source<EnvelopedRecordBatch<KafkaSourcePosition>, KafkaControl>
where
    F: MqPayloadFormat,
{
    let topic = topic.into();
    KafkaSource::committable_payload_batches(settings, Subscription::topics([topic.clone()]))
        .try_map(move |batch| {
            decode_kafka_batch(&topic, &format, &batch).map_err(crate::stream_error)
        })
}

/// Decodes one `datum-mq` payload batch into an enveloped Arrow batch.
pub fn decode_kafka_batch<F>(
    topic: &str,
    format: &F,
    batch: &KafkaPayloadBatch,
) -> Result<EnvelopedRecordBatch<KafkaSourcePosition>>
where
    F: MqPayloadFormat,
{
    let payloads = batch
        .records()
        .iter()
        .map(|record| batch.payload(record))
        .collect::<Vec<_>>();
    let record_batch = format.decode_payloads(&payloads)?;
    let position = KafkaSourcePosition::from_payload_batch(topic, batch);
    Ok(EnvelopedRecordBatch::new(
        record_batch,
        BatchEnvelope::new(position, format.schema_revision()),
    ))
}

/// Decodes one `datum-mq` payload batch into a committable SQL batch.
pub fn decode_kafka_committable_batch<F>(
    topic: &str,
    format: &F,
    batch: &KafkaPayloadBatch,
) -> Result<KafkaCommittableRecordBatch>
where
    F: MqPayloadFormat,
{
    let decoded = decode_kafka_batch(topic, format, batch)?;
    let (record_batch, envelope) = decoded.into_parts();
    let schema_revision = envelope.schema_revision();
    let position = envelope.into_source_position();
    let commit_batch = batch.clone();
    let commit = SourceCommit::from_fn("Kafka payload batch offset commit", move || {
        commit_batch.commit().map_err(crate::stream_error)
    });
    Ok(CommittableRecordBatch::new(
        record_batch,
        Some(SqlSourcePosition::Kafka(position)),
        schema_revision,
        commit,
    ))
}