use crate::transport::InputCommandReceiver;
use crate::transport::kafka::MemoryUseReporter;
use crate::transport::kafka::ft::count_partitions_in_topic;
use crate::transport::kafka::{generate_oauthbearer_token, validate_aws_msk_region};
use crate::{InputBuffer, Parser};
use crate::{
    InputConsumer, TransportInputEndpoint,
    transport::{
        InputReader,
        kafka::{DeferredLogging, rdkafka_loglevel_from, refine_kafka_error},
    },
};
use anyhow::{Error as AnyError, Result as AnyResult, anyhow, bail};
use chrono::{DateTime, Utc};
use crossbeam::queue::ArrayQueue;
use crossbeam::sync::{Parker, Unparker};
use dbsp::operator::StagedBuffers;
use feldera_adapterlib::ConnectorMetadata;
use feldera_adapterlib::format::BufferSize;
use feldera_adapterlib::transport::{
    InputEndpoint, InputReaderCommand, Resume, Watermark, parse_resume_info,
};
use feldera_sqllib::{ByteArray, SqlString, Timestamp, Variant};
use feldera_types::config::FtModel;
use feldera_types::program_schema::Relation;
use feldera_types::transport::kafka::{KafkaInputConfig, KafkaStartFromConfig};
use itertools::Itertools;
use rdkafka::client::OAuthToken;
use rdkafka::config::RDKafkaLogLevel;
use rdkafka::consumer::base_consumer::PartitionQueue;
use rdkafka::message::{BorrowedMessage, Headers};
use rdkafka::{
    ClientConfig, ClientContext, Message,
    config::FromClientConfigAndContext,
    consumer::{BaseConsumer, Consumer, ConsumerContext},
    error::{KafkaError, KafkaResult},
};
use rdkafka::{Offset, TopicPartitionList};
use serde::{Deserialize, Serialize};
use smallvec::{SmallVec, smallvec};
#[cfg(test)]
use std::collections::BTreeSet;
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::error::Error;
use std::hash::Hasher;
use std::ops::Range;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
use std::thread::{self, JoinHandle, Thread};
use std::{
    iter,
    sync::{Arc, Mutex, Weak},
    thread::spawn,
    time::Duration,
};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
use tracing::span::EnteredSpan;
use tracing::{debug, info_span, warn};
use xxhash_rust::xxh3::Xxh3Default;

/// Poll timeout must be low, as it bounds the amount of time it takes to resume the connector.
const POLL_TIMEOUT: Duration = Duration::from_millis(5);

const METADATA_TIMEOUT: Duration = Duration::from_secs(10);

// Size of the circular buffer used to pass errors from ClientContext
// to the worker thread.
const ERROR_BUFFER_SIZE: usize = 1000;

pub struct KafkaFtInputEndpoint {
    config: Arc<KafkaInputConfig>,
}

impl KafkaFtInputEndpoint {
    pub fn new(mut config: KafkaInputConfig) -> AnyResult<KafkaFtInputEndpoint> {
        config.validate()?;
        Ok(KafkaFtInputEndpoint {
            config: Arc::new(config),
        })
    }
}

struct KafkaFtInputReader {
    inner: Arc<KafkaFtInputReaderInner>,
    command_sender: UnboundedSender<InputReaderCommand>,
    poller_thread: Thread,
}

struct KafkaFtInputContext {
    // We keep a weak reference to the endpoint to avoid a reference cycle:
    // endpoint->BaseConsumer->context->endpoint.
    endpoint: Mutex<Weak<KafkaFtInputReaderInner>>,

    deferred_logging: DeferredLogging,

    oauthbearer_config: HashMap<String, String>,

    memory_use_reporter: Mutex<MemoryUseReporter>,

    topic: String,
}

impl KafkaFtInputContext {
    fn new(kafka_config: &KafkaInputConfig) -> AnyResult<Self> {
        let mut oauthbearer_config = HashMap::new();
        if let Some(region) =
            validate_aws_msk_region(&kafka_config.kafka_options, kafka_config.region.clone())?
        {
            oauthbearer_config.insert("region".to_owned(), region);
        };

        Ok(Self {
            endpoint: Mutex::new(Weak::new()),
            deferred_logging: DeferredLogging::new(),
            oauthbearer_config,
            topic: kafka_config.topic.clone(),
            memory_use_reporter: Mutex::new(MemoryUseReporter::new()),
        })
    }
}

impl ClientContext for KafkaFtInputContext {
    const ENABLE_REFRESH_OAUTH_TOKEN: bool = true;

    fn log(&self, level: RDKafkaLogLevel, fac: &str, log_message: &str) {
        self.deferred_logging.log(level, fac, log_message);
    }

    fn error(&self, error: KafkaError, reason: &str) {
        // eprintln!("Kafka error: {error}");
        if let Some(endpoint) = self.endpoint.lock().unwrap().upgrade() {
            endpoint.push_error(error, reason);
        }
    }

    fn generate_oauth_token(&self, _: Option<&str>) -> Result<OAuthToken, Box<dyn Error>> {
        generate_oauthbearer_token(&self.oauthbearer_config)
    }

    fn stats(&self, statistics: rdkafka::Statistics) {
        let _guard = span(&self.topic);
        self.memory_use_reporter.lock().unwrap().update(&statistics);
    }
}

impl ConsumerContext for KafkaFtInputContext {}

struct KafkaFtHasher(BTreeMap<i32, Xxh3Default>);

impl KafkaFtHasher {
    fn new(partitions: &[i32]) -> Self {
        Self(
            partitions
                .iter()
                .cloned()
                .zip(iter::repeat_with(Xxh3Default::new))
                .collect(),
        )
    }

    fn add<B>(&mut self, partition: &i32, buffer: &B)
    where
        B: InputBuffer,
    {
        buffer.hash(&mut self.0.get_mut(partition).unwrap_or_else(|| {
            panic!("failed to get the partition '{partition}' in KafkaFtHasher")
        }));
    }

    fn reset(&mut self) {
        for hasher in self.0.values_mut() {
            hasher.reset();
        }
    }

    fn finish(&self) -> u64 {
        let mut h = Xxh3Default::new();
        for (partition, hasher) in self.0.iter() {
            h.write_i32(*partition);
            h.write_u64(hasher.finish());
        }
        h.finish()
    }

    fn take(&mut self) -> u64 {
        let hash = self.finish();
        self.reset();
        hash
    }
}

struct KafkaFtInputReaderInner {
    kafka_consumer: Arc<BaseConsumer<KafkaFtInputContext>>,
    errors: ArrayQueue<(KafkaError, String)>,
}

/// Error returned by [KafkaFtInputReaderInner::check_offsets].
enum OffsetError {
    /// Error fetching watermarks for Kafka partitions.
    FetchError(AnyError),

    /// At least one provided partition offset does not exist in the partition
    /// (probably it was expired).
    MissingOffset(AnyError),
}

impl OffsetError {
    /// Returns the inner [AnyError].
    fn into_any_error(self) -> AnyError {
        match self {
            OffsetError::FetchError(error) => error,
            OffsetError::MissingOffset(error) => error,
        }
    }
}

impl KafkaFtInputReaderInner {
    fn check_offsets(
        &self,
        config: &KafkaInputConfig,
        offsets: &[i64],
    ) -> Result<Vec<Offset>, OffsetError> {
        let partitions = config
            .partitions
            .clone()
            .unwrap_or((0..offsets.len() as i32).collect());

        for (offset, partition) in offsets.iter().copied().zip(partitions) {
            let (low, high) = self
                .kafka_consumer
                .fetch_watermarks(&config.topic, partition, METADATA_TIMEOUT)
                .map_err(|e| {
                    OffsetError::FetchError(anyhow!(
                        "error fetching metadata for partition '{partition}': {e}"
                    ))
                })?;

            // Return an error if the specified offset doesn't exist.
            if !(low..=high).contains(&offset) {
                return Err(OffsetError::MissingOffset(anyhow!(
                    "configuration error: provided offset '{offset}' not currently in partition '{partition}'"
                )));
            }
        }
        Ok(offsets.iter().copied().map(Offset::Offset).collect())
    }

    #[allow(clippy::borrowed_box)]
    fn poller_thread(
        &self,
        config: Arc<KafkaInputConfig>,
        consumer: &Box<dyn InputConsumer>,
        mut parser: Box<dyn Parser>,
        n_partitions: usize,
        command_receiver: UnboundedReceiver<InputReaderCommand>,
        resume_info: Option<Metadata>,
    ) -> AnyResult<()> {
        let topic = &config.topic;

        // Start reading the partitions either at the resume point or at the
        // beginning.
        let mut command_receiver = InputCommandReceiver::<Metadata, ()>::new(command_receiver);
        let mut persisted_partitions = None;
        let initial_offsets = match resume_info {
            Some(metadata) => {
                let offsets = metadata
                    .parse(n_partitions, &mut persisted_partitions)?
                    .into_iter()
                    .map(|range| range.end)
                    .collect::<Vec<_>>();
                if config.resume_earliest_if_data_expires {
                    match self.check_offsets(&config, &offsets) {
                        Ok(offsets) => offsets,
                        Err(OffsetError::FetchError(error)) => return Err(error),
                        Err(OffsetError::MissingOffset(_error)) => {
                            warn!(
                                "Checkpoint for topic {topic} refers to offsets that no longer exist in the topic; as configured, the connector will resume from the earliest offsets now in the topic (checkpointed offsets were {offsets:?})"
                            );
                            iter::repeat_n(Offset::Beginning, n_partitions).collect()
                        }
                    }
                } else {
                    offsets.into_iter().map(Offset::Offset).collect()
                }
            }
            None => match &config.start_from {
                KafkaStartFromConfig::Earliest => {
                    iter::repeat_n(Offset::Beginning, n_partitions).collect()
                }
                KafkaStartFromConfig::Latest => iter::repeat_n(Offset::End, n_partitions).collect(),
                KafkaStartFromConfig::Offsets(offsets) => {
                    if n_partitions != offsets.len() {
                        bail!(
                            "Topic {topic} has {n_partitions} partitions but configuration specifies {} offsets.",
                            offsets.len()
                        );
                    }

                    self.check_offsets(&config, offsets)
                        .map_err(|e| e.into_any_error())?
                }
                KafkaStartFromConfig::Timestamp(timestamp) => {
                    // The partitions that we're reading.
                    let partitions = config
                        .partitions
                        .clone()
                        .unwrap_or((0..n_partitions as i32).collect());

                    // Assemble a TopicPartitionList with the desired timestamp
                    // as the "offset" for each partition.
                    let mut timestamps = TopicPartitionList::new();
                    for partition in partitions.iter().copied() {
                        timestamps
                            .add_partition_offset(topic, partition, Offset::Offset(*timestamp))
                            .map_err(|error| self.refine_error(error).1)?;
                    }

                    // Translate the timestamps to offsets.
                    let offsets = self
                        .kafka_consumer
                        .offsets_for_times(timestamps, METADATA_TIMEOUT)
                        .map_err(|error| self.refine_error(error).1)?;

                    // Extract the offsets from the returned value.
                    offsets
                        .elements()
                        .into_iter()
                        .zip_eq(partitions)
                        .map(|(elem, partition)| {
                            if elem.partition() == partition {
                                Ok(elem.offset())
                            } else {
                                Err(anyhow!(
                                    "rdkafka returned partition {} instead of expected {partition}",
                                    elem.partition()
                                ))
                            }
                        })
                        .collect::<Result<_, AnyError>>()?
                }
            },
        };
        let next_offsets = initial_offsets
            .iter()
            .map(|o| match o {
                Offset::Offset(offset) => *offset,
                _ => 0,
            })
            .collect::<Vec<_>>();

        let mut assignment = TopicPartitionList::new();

        let partitions = config
            .partitions
            .clone()
            .unwrap_or((0..initial_offsets.len() as i32).collect());

        for (partition, offset) in partitions.iter().zip(initial_offsets) {
            assignment
                .add_partition_offset(topic, *partition, offset)
                .map_err(|error| self.refine_error(error).1)?;
        }

        self.kafka_consumer
            .assign(&assignment)
            .map_err(|error| self.refine_error(error).1)?;

        // Prepare to launch threads.
        let exit = Arc::new(AtomicBool::new(false));
        let n_threads = config.poller_threads().min(n_partitions);
        let mut threads = (0..n_threads)
            .map(|_| RecvThread {
                topic: topic.clone(),
                exit: exit.clone(),
                main_thread: thread::current(),
                parker: Parker::new(),
                base_consumer: self.kafka_consumer.clone(),
                consumer: consumer.clone(),
                parser: parser.fork(),
                receivers: Vec::with_capacity(n_partitions.div_ceil(n_threads)),
            })
            .collect::<Vec<_>>();

        // Split every partition away as its own separate queue.
        let mut receivers = BTreeMap::new();

        for ((partition, thread), next_offset) in partitions
            .iter()
            .zip((0..n_threads).cycle())
            .zip(next_offsets)
        {
            let thread = &mut threads[thread];
            let mut queue = self
                .kafka_consumer
                .split_partition_queue(topic, *partition)
                .ok_or_else(|| anyhow!("could not split queue for partition {partition}"))?;

            let unparker = thread.parker.unparker().clone();
            queue.set_nonempty_callback({
                let unparker = unparker.clone();
                move || unparker.unpark()
            });

            let receiver = Arc::new(PartitionReceiver::new(
                *partition,
                queue,
                next_offset,
                &config,
                unparker,
            ));
            receivers.insert(partition, receiver.clone());
            thread.receivers.push(receiver);
        }

        // Poll the main consumer queue until it is empty.
        //
        // We have to do this because there's a race between assigning topics
        // and splitting the partition queues.  But (presumably), after we
        // process initial messages that might go to the main queue, the rest
        // should go to the split queues.
        while let Some(message) = self.kafka_consumer.poll(Duration::ZERO) {
            match message {
                Err(KafkaError::PartitionEOF(p)) if (0..n_partitions as i32).contains(&p) => {
                    receivers[&p].eof.store(true, Ordering::Relaxed);
                }
                Err(e) => {
                    let (fatal, e) = self.refine_error(e);
                    consumer.error(fatal, e, Some("kafka-poll"));
                    if fatal {
                        return Ok(());
                    }
                }
                Ok(message) => {
                    let partition = message.partition();
                    let receiver = receivers.get_mut(&partition).ok_or_else(|| {
                        anyhow!("received message for nonexistent partition {partition}")
                    })?;
                    receiver.handle_kafka_message(
                        &self.kafka_consumer,
                        &**consumer,
                        &mut *parser,
                        Ok(message),
                    );
                }
            }
        }

        // Launch threads.
        let threads = threads
            .into_iter()
            .map(RecvThreadHandle::new)
            .collect::<Vec<_>>();

        // Then replay as many steps as requested.
        while let Some((metadata, ())) = command_receiver.blocking_recv_replay()? {
            let n_offsets = metadata.offsets.len();
            let metadata = metadata.parse(n_partitions, &mut persisted_partitions)?;
            let mut incomplete_partitions = HashSet::new();
            for (offsets, (partition, receiver)) in metadata.iter().zip(receivers.iter()) {
                if !offsets.is_empty() {
                    receiver.set_max_offset(offsets.end - 1);
                    incomplete_partitions.insert(partition);
                }
            }
            for thread in &threads {
                thread.unparker.unpark();
            }
            let mut total = BufferSize::empty();
            let mut hasher = KafkaFtHasher::new(&partitions[..n_offsets]);
            loop {
                // Process messages for all partitions.
                for (partition, receiver) in receivers.iter() {
                    let max = receiver.max_offset();
                    while let Some((offset, (mut msg, _timestamp))) = receiver.read(max) {
                        let amt = msg.len();
                        total += amt;
                        hasher.add(partition, &msg);
                        consumer.buffered(amt);
                        msg.flush();
                        if offset == max {
                            incomplete_partitions.remove(partition);
                        }
                    }
                }
                if incomplete_partitions.is_empty() {
                    break;
                }

                // Poll the Kafka consumer. We don't want messages from it but it's obligatory.
                match self.kafka_consumer.poll(Duration::ZERO) {
                    Some(Err(e)) => {
                        let (fatal, e) = self.refine_error(e);
                        consumer.error(fatal, e, Some("kafka-poll"));
                        if fatal {
                            return Ok(());
                        }
                    }
                    Some(Ok(message)) => {
                        bail!(
                            "All partitions are split but the main consumer received a message anyway: {message:?}"
                        );
                    }
                    None => (),
                }

                if receivers.values().any(|receiver| receiver.fatal_error()) {
                    return Ok(());
                }

                thread::park_timeout(POLL_TIMEOUT);
            }
            consumer.replayed(total, hasher.finish());
        }

        // We're done replaying.

        let mut running = false;
        let mut kafka_paused = false;
        struct StagedBuffer {
            buffers: Box<dyn StagedBuffers>,
            timestamp: DateTime<Utc>,
            amt: BufferSize,
            hash: u64,
            offsets: Vec<Range<i64>>,
        }
        struct Stager {
            buffers: VecDeque<StagedBuffer>,
            offsets: Vec<Range<i64>>,
            hasher: KafkaFtHasher,
            inputs: Vec<Box<dyn InputBuffer>>,
            amt: BufferSize,
            timestamp: Option<DateTime<Utc>>,
        }
        impl Stager {
            fn new(receivers: &BTreeMap<&i32, Arc<PartitionReceiver>>, partitions: &[i32]) -> Self {
                Self {
                    buffers: VecDeque::new(),
                    offsets: receivers
                        .values()
                        .map(|r| {
                            let next_offset = r.next_offset();
                            next_offset..next_offset
                        })
                        .collect(),
                    hasher: KafkaFtHasher::new(partitions),
                    inputs: Vec::new(),
                    amt: BufferSize::default(),
                    timestamp: None,
                }
            }
            fn take_offsets(&mut self) -> Vec<Range<i64>> {
                let offsets = self.offsets.clone();
                for partition_offsets in &mut self.offsets {
                    partition_offsets.start = partition_offsets.end;
                }
                offsets
            }
            fn take_buffer(&mut self, parser: &dyn Parser) -> StagedBuffer {
                StagedBuffer {
                    buffers: parser.stage(std::mem::take(&mut self.inputs)),
                    timestamp: self.timestamp.take().unwrap_or_else(Utc::now),
                    amt: std::mem::take(&mut self.amt),
                    hash: self.hasher.take(),
                    offsets: self.take_offsets(),
                }
            }
            fn pop_staged_buffer(&mut self, parser: &dyn Parser) -> StagedBuffer {
                self.buffers
                    .pop_front()
                    .unwrap_or_else(|| self.take_buffer(parser))
            }
            fn push_record(
                &mut self,
                consumer: &dyn InputConsumer,
                partition: &i32,
                offset: i64,
                msg: Option<Box<dyn InputBuffer>>,
                index: usize,
            ) {
                // Set the time when we start getting new data to be staged as ingestion timestamp.
                if self.timestamp.is_none() {
                    self.timestamp = Some(Utc::now());
                }
                let amt = msg.len();
                consumer.buffered(amt);
                self.amt += amt;
                self.hasher.add(partition, &msg);
                if let Some(buffer) = msg {
                    self.inputs.push(buffer);
                }

                let range = &mut self.offsets[index];
                if range.is_empty() {
                    *range = offset..offset + 1;
                } else {
                    range.end = offset + 1;
                }
            }

            fn maybe_stage(&mut self, consumer: &dyn InputConsumer, parser: &dyn Parser) {
                if self.amt.records >= consumer.max_batch_size() {
                    let buffer = self.take_buffer(parser);
                    self.buffers.push_back(buffer);
                }
            }
        }
        let mut stager = Stager::new(&receivers, &partitions);

        loop {
            let was_running = running;
            while let Some(command) = command_receiver.try_recv()? {
                match command {
                    command @ InputReaderCommand::Replay { .. } => {
                        unreachable!("{command:?} must be at the beginning of the command stream")
                    }
                    InputReaderCommand::Extend => running = true,
                    InputReaderCommand::Pause => running = false,
                    InputReaderCommand::Queue { .. } => {
                        let mut buffer = stager.pop_staged_buffer(&*parser);
                        buffer.buffers.flush();
                        let metadata = serde_json::to_value(&Metadata {
                            offsets: buffer.offsets,
                        })
                        .unwrap();
                        consumer.extended(
                            buffer.amt,
                            Some(Resume::Replay {
                                hash: buffer.hash,
                                seek: metadata.clone(),
                                replay: rmpv::Value::Nil,
                            }),
                            vec![Watermark::new(buffer.timestamp, Some(metadata))],
                        );
                    }
                    InputReaderCommand::Disconnect => return Ok(()),
                }
            }

            if !running {
                if !kafka_paused {
                    self.pause_partitions()
                        .map_err(|error| self.refine_error(error).1)?;
                    for receiver in receivers.values() {
                        receiver.set_max_offset(i64::MIN);
                    }
                    kafka_paused = true;
                }
            } else {
                if kafka_paused {
                    self.resume_partitions()
                        .map_err(|error| self.refine_error(error).1)?;
                    kafka_paused = false;
                }
                if !was_running {
                    for receiver in receivers.values() {
                        receiver.set_max_offset(i64::MAX);
                    }
                    for thread in &threads {
                        thread.unparker.unpark();
                    }
                }
            }

            let mut n_read = 0;
            if running {
                if config.synchronize_partitions {
                    /// Returns the index(es) of the partition(s) in `receivers`
                    /// with the earliest timestamp(s), or `None` if any of the
                    /// receivers are out of received items.
                    fn find_earliest(
                        receivers: &BTreeMap<&i32, Arc<PartitionReceiver>>,
                    ) -> Option<SmallVec<[usize; 1]>> {
                        let mut iter = receivers.values().enumerate().map(|(index, receiver)| {
                            receiver.peek().map(|timestamp| (index, timestamp))
                        });
                        let (index, mut best_timestamp) = iter.next()??;
                        let mut best_indexes = smallvec![index];
                        for item in iter {
                            let (index, timestamp) = item?;
                            if timestamp < best_timestamp {
                                best_indexes.clear();
                                best_indexes.push(index);
                                best_timestamp = timestamp;
                            } else if timestamp == best_timestamp {
                                best_indexes.push(index);
                            }
                        }
                        Some(best_indexes)
                    }

                    while n_read < partitions.len()
                        && let Some(indexes) = find_earliest(&receivers)
                    {
                        n_read += indexes.len();
                        for index in indexes {
                            let partition = partitions[index];
                            let (offset, (msg, _timestamp)) =
                                receivers[&partition].read(i64::MAX).unwrap();
                            stager.push_record(&**consumer, &partition, offset, msg, index);
                        }
                    }
                } else {
                    for (index, (partition, receiver)) in receivers.iter().enumerate() {
                        if let Some((offset, (msg, _timestamp))) = receiver.read(i64::MAX) {
                            stager.push_record(&**consumer, partition, offset, msg, index);
                            n_read += 1;
                        }
                    }
                };

                stager.maybe_stage(&**consumer, &*parser);
            }

            // Keep polling even while the consumer is paused as `BaseConsumer`
            // processes control messages (including rebalancing and errors)
            // within the polling thread.
            match self.kafka_consumer.poll(Duration::ZERO) {
                None => (),
                Some(Err(e)) => {
                    let (fatal, e) = self.refine_error(e);
                    consumer.error(fatal, e, Some("kafka-poll"));
                    if fatal {
                        return Ok(());
                    }
                }
                Some(Ok(message)) => {
                    bail!(
                        "All partitions are split but the main consumer received a message anyway: {message:?}"
                    );
                }
            }

            while let Some((error, reason)) = self.pop_error() {
                let (fatal, _e) = self.refine_error(error);
                // `reason` contains a human-readable description of the
                // error.
                consumer.error(fatal, anyhow!(reason), None);
                if fatal {
                    return Ok(());
                }
            }

            if receivers.values().any(|r| r.fatal_error()) {
                return Ok(());
            }
            if receivers.values().all(|r| r.eof()) {
                tracing::info!("reached end of all partitions (`enable.partition.eof` configured)");
                consumer.eoi();
                return Ok(());
            }

            if n_read == 0 {
                thread::park_timeout(Duration::from_secs(1));
            }
        }
    }

    fn push_error(&self, error: KafkaError, reason: &str) {
        // `force_push` makes the queue operate as a circular buffer.
        self.errors.force_push((error, reason.to_string()));
    }

    fn pop_error(&self) -> Option<(KafkaError, String)> {
        self.errors.pop()
    }

    /// Pause all partitions assigned to the consumer.
    fn pause_partitions(&self) -> KafkaResult<()> {
        self.kafka_consumer
            .pause(&self.kafka_consumer.assignment()?)
    }

    /// Resume all partitions assigned to the consumer.
    fn resume_partitions(&self) -> KafkaResult<()> {
        self.kafka_consumer
            .resume(&self.kafka_consumer.assignment()?)
    }

    fn refine_error(&self, e: KafkaError) -> (bool, AnyError) {
        refine_kafka_error(self.kafka_consumer.client(), e)
    }
}

fn span(topic: &str) -> EnteredSpan {
    info_span!("kafka_input", topic = topic).entered()
}

impl KafkaFtInputReader {
    fn new(
        config: &Arc<KafkaInputConfig>,
        consumer: Box<dyn InputConsumer>,
        parser: Box<dyn Parser>,
        resume_info: Option<serde_json::Value>,
    ) -> AnyResult<Self> {
        let resume_info = if let Some(resume_info) = resume_info {
            Some(parse_resume_info::<Metadata>(&resume_info)?)
        } else {
            None
        };
        let _guard = span(&config.topic);

        // Create Kafka consumer configuration.
        debug!("Starting Kafka input endpoint: {:?}", config);

        let mut client_config = ClientConfig::new();

        for (key, value) in config.kafka_options.iter() {
            client_config.set(key, value);
        }

        if let Some(log_level) = config.log_level {
            client_config.set_log_level(rdkafka_loglevel_from(log_level));
        }

        // Create Kafka consumer and count the number of partitions in the topic.
        //
        // This has the desirable side effect of ensuring that we can reach the
        // broker and failing with an error if we cannot.
        let context = KafkaFtInputContext::new(config)?;
        let kafka_consumer = BaseConsumer::from_config_and_context(&client_config, context)?;

        // IMPORTANT: Poll before trying to fetch metadata. Necessary so that OAUTHBREAKER token
        // is set.
        kafka_consumer.poll(std::time::Duration::from_nanos(0));
        let partition_count = count_partitions_in_topic(&kafka_consumer, &config.topic)?;

        let inner = Arc::new(KafkaFtInputReaderInner {
            kafka_consumer: Arc::new(kafka_consumer),
            errors: ArrayQueue::new(ERROR_BUFFER_SIZE),
        });
        *inner.kafka_consumer.context().endpoint.lock().unwrap() = Arc::downgrade(&inner);

        let (command_sender, command_receiver) = unbounded_channel();
        let poller_handle = thread::Builder::new()
            .name("kafka-input-poller".to_string())
            .spawn({
                let endpoint = inner.clone();
                let config = config.clone();
                move || {
                    let _guard = span(&config.topic);
                    if let Err(e) = endpoint.poller_thread(
                        config.clone(),
                        &consumer,
                        parser,
                        config
                            .partitions
                            .as_deref()
                            .map(|p| p.len())
                            .unwrap_or(partition_count),
                        command_receiver,
                        resume_info,
                    ) {
                        consumer.error(true, e, None);
                    }
                    tracing::info!("kafka input complete")
                }
            })
            .expect("failed to spawn Kafka input poller thread");
        let poller_thread = poller_handle.thread().clone();
        Ok(KafkaFtInputReader {
            inner,
            command_sender,
            poller_thread,
        })
    }
}

impl InputEndpoint for KafkaFtInputEndpoint {
    fn fault_tolerance(&self) -> Option<FtModel> {
        Some(FtModel::ExactlyOnce)
    }
}

impl TransportInputEndpoint for KafkaFtInputEndpoint {
    fn open(
        &self,
        consumer: Box<dyn InputConsumer>,
        parser: Box<dyn Parser>,
        _schema: Relation,
        resume_info: Option<serde_json::Value>,
    ) -> AnyResult<Box<dyn InputReader>> {
        Ok(Box::new(KafkaFtInputReader::new(
            &self.config,
            consumer,
            parser,
            resume_info,
        )?))
    }
}

impl InputReader for KafkaFtInputReader {
    fn as_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> {
        self
    }

    fn request(&self, command: InputReaderCommand) {
        let _ = self.command_sender.send(command);
        self.poller_thread.unpark();
    }
    fn is_closed(&self) -> bool {
        self.command_sender.is_closed()
    }
    fn memory(&self) -> usize {
        self.inner
            .kafka_consumer
            .context()
            .memory_use_reporter
            .lock()
            .unwrap()
            .current()
    }
}

impl Drop for KafkaFtInputReader {
    fn drop(&mut self) {
        self.request(InputReaderCommand::Disconnect);
    }
}

#[derive(Serialize, Deserialize)]
pub(super) struct Metadata {
    /// Per-partition ranges of offsets.
    pub offsets: Vec<Range<i64>>,
}

impl Metadata {
    fn parse(
        self,
        n_partitions: usize,
        persisted_partitions: &mut Option<usize>,
    ) -> AnyResult<Vec<Range<i64>>> {
        // Kafka supports increasing the number of partitions in a topic, but
        // not decreasing it.  If we are restarted with more partitions than we
        // previously had, then we interpret checkpoints and journal entries as
        // having empty ranges for the new partitions.  We log the discrepancy
        // once at the beginning and once each time the number of journaled
        // partitions increases.
        //
        // We disallow and fail if the number of partitions decreases, because
        // Kafka doesn't support that either and we don't have a sensible way to
        // recover from it (probably the lost partitions contained data and we
        // can't recover it).
        //
        // We use `persisted_partitions` to ensure that the number of partitions
        // increases monotonically.
        let mut offsets = self.offsets;
        let n_offsets = offsets.len();
        if n_offsets > n_partitions {
            bail!(
                "metadata for replay has offsets for {n_offsets} partitions but topic has only {n_partitions} partitions; because Kafka does not allow partitions to be deleted, this implies that the topic has been deleted and recreated, which makes data loss likely"
            );
        }

        if let Some(persisted_partitions) = *persisted_partitions {
            if n_offsets > persisted_partitions {
                warn!("journaled partitions increased from {persisted_partitions} to {n_offsets}");
            } else if n_offsets < persisted_partitions {
                // Decreasing number of offsets indicates an inconsistency.
                bail!("journaled partitions decreased from {persisted_partitions} to {n_offsets}");
            }
        } else if n_offsets != n_partitions {
            warn!(
                "topic partitions increased since the checkpoint from {n_offsets} to {n_partitions}; reading additional partitions from the start",
            );
        }
        *persisted_partitions = Some(n_offsets);
        offsets.extend((n_offsets..n_partitions).map(|_| 0..0));
        Ok(offsets)
    }
}

/// For testing purposes, keeps track of partitions that we're not reading from
/// because there is already too much queued.
#[cfg(test)]
pub static BACKPRESSURE: Mutex<BTreeSet<(String, i32)>> = Mutex::new(BTreeSet::new());

#[cfg(test)]
fn update_backpressure(topic: &str, partition: i32, has_backpressure: bool) {
    let mut backpressure = BACKPRESSURE.lock().unwrap();
    let member = (topic.into(), partition);
    if has_backpressure {
        backpressure.insert(member);
    } else {
        backpressure.remove(&member);
    }
}

#[cfg(not(test))]
fn update_backpressure(_topic: &str, _partition: i32, _has_backpressure: bool) {}

/// Returns true if a partition queue that holds `n_bytes` should pause for
/// backpressure.
fn needs_backpressure(n_bytes: usize) -> bool {
    n_bytes >= 1_000_000
}

struct PartitionReceiver {
    partition: i32,
    queue: PartitionQueue<KafkaFtInputContext>,
    config: KafkaInputConfig,
    metadata_requested: bool,

    /// The maximum message offset that we want to receive, used as follows:
    ///
    /// - `i64::MIN`, the initial value, disables receiving messages entirely.
    ///
    /// - The offset of a specific message ensures that we will not read
    ///   messages beyond that one.  This is used during replay, set to the
    ///   offset of the final message to be replayed.
    ///
    /// - `i64::MAX`, used when replaying is complete, will read all messages.
    max_offset: AtomicI64,

    /// The minimum message offset that we could receive next in this partition
    /// and offset.
    ///
    /// We process messages in order, so when `messages` is non-empty, this is
    /// the largest offset (key) in `messages` plus 1.
    next_offset: AtomicI64,

    /// Initial value of `next_offset`.  Only used in log messages.
    initial_next_offset: i64,

    /// Parsed messages and errors, paired with timestamps.
    #[allow(clippy::type_complexity)]
    messages: Mutex<BTreeMap<i64, (Option<Box<dyn InputBuffer>>, i64)>>,

    /// Number of bytes of buffers in `messages`, available without taking the
    /// lock.
    n_bytes: AtomicUsize,

    /// Wakes up the [RecvThread] that receives into this partition.
    unparker: Unparker,

    eof: AtomicBool,
    fatal_error: AtomicBool,
}

impl PartitionReceiver {
    pub fn new(
        partition: i32,
        queue: PartitionQueue<KafkaFtInputContext>,
        next_offset: i64,
        config: &KafkaInputConfig,
        unparker: Unparker,
    ) -> Self {
        let metadata_requested = config.metadata_requested();

        Self {
            partition,
            queue,
            max_offset: AtomicI64::new(i64::MIN),
            next_offset: AtomicI64::new(next_offset),
            initial_next_offset: next_offset,
            messages: Mutex::new(BTreeMap::new()),
            n_bytes: AtomicUsize::new(0),
            eof: AtomicBool::new(false),
            fatal_error: AtomicBool::new(false),
            config: config.clone(),
            metadata_requested,
            unparker,
        }
    }

    #[allow(clippy::type_complexity)]
    pub fn read(&self, max: i64) -> Option<(i64, (Option<Box<dyn InputBuffer>>, i64))> {
        let mut messages = self.messages.lock().unwrap();
        match messages.first_key_value() {
            Some((offset, _)) if *offset <= max => {
                let (offset, (buffer, timestamp)) = messages.pop_first().unwrap();

                // Account the subtraction of the buffer from `self.n_bytes`.
                // If that releases backpressure, then wake up its receiver
                // thread.
                let buffer_len = buffer.len().bytes;
                let old_nbytes = self.n_bytes.fetch_sub(buffer_len, Ordering::Relaxed);
                let new_bytes = old_nbytes - buffer_len;
                if needs_backpressure(old_nbytes) && !needs_backpressure(new_bytes) {
                    self.unparker.unpark();
                }

                Some((offset, (buffer, timestamp)))
            }
            _ => None,
        }
    }

    pub fn peek(&self) -> Option<i64> {
        self.messages
            .lock()
            .unwrap()
            .first_entry()
            .map(|entry| entry.get().1)
    }

    pub fn fatal_error(&self) -> bool {
        self.fatal_error.load(Ordering::Relaxed)
    }

    pub fn eof(&self) -> bool {
        self.eof.load(Ordering::Relaxed)
    }

    fn set_max_offset(&self, max: i64) {
        self.max_offset.store(max, Ordering::Relaxed);
    }

    pub fn max_offset(&self) -> i64 {
        self.max_offset.load(Ordering::Relaxed)
    }

    fn next_offset(&self) -> i64 {
        self.next_offset.load(Ordering::Relaxed)
    }

    /// Create record metadata from Kafka message containing only properties specified in the connector config.
    fn create_metadata(&self, message: &BorrowedMessage<'_>) -> Option<ConnectorMetadata> {
        if !self.metadata_requested {
            return None;
        }

        let mut metadata = ConnectorMetadata::new();

        if self.config.include_topic == Some(true) {
            metadata.insert(
                "kafka_topic",
                Variant::String(SqlString::from(message.topic())),
            );
        };

        if self.config.include_timestamp == Some(true) {
            let timestamp = message.timestamp().to_millis();
            if let Some(timestamp) = timestamp {
                metadata.insert(
                    "kafka_timestamp",
                    Variant::Timestamp(Timestamp::from_milliseconds(timestamp)),
                );
            }
        }

        if self.config.include_partition == Some(true) {
            metadata.insert("kafka_partition", Variant::Int(message.partition()));
        }

        if self.config.include_offset == Some(true) {
            metadata.insert("kafka_offset", Variant::BigInt(message.offset()));
        }

        if self.config.include_headers == Some(true) {
            let mut kafka_headers = BTreeMap::new();

            if let Some(headers) = message.headers() {
                for i in 0..headers.count() {
                    let header = headers.get(i);
                    let value = if let Some(value) = header.value {
                        Variant::Binary(ByteArray::from(value))
                    } else {
                        Variant::SqlNull
                    };
                    kafka_headers.insert(Variant::String(SqlString::from(header.key)), value);
                }
            }

            metadata.insert("kafka_headers", Variant::Map(Arc::new(kafka_headers)));
        }

        Some(metadata)
    }

    fn handle_kafka_message(
        &self,
        base_consumer: &BaseConsumer<KafkaFtInputContext>,
        consumer: &dyn InputConsumer,
        parser: &mut dyn Parser,
        message: KafkaResult<BorrowedMessage<'_>>,
    ) {
        match message {
            Err(KafkaError::PartitionEOF(_)) => {
                self.eof.store(true, Ordering::Relaxed);
            }
            Err(e) => {
                let (fatal, e) = refine_kafka_error(base_consumer.client(), e);
                consumer.error(fatal, e, Some("kafka-msg"));
                if fatal {
                    self.fatal_error.store(true, Ordering::Relaxed);
                }
            }
            Ok(message) => {
                let offset = message.offset();
                let next_offset = self.next_offset();
                if offset >= next_offset {
                    self.next_offset.store(offset + 1, Ordering::Relaxed);
                    let timestamp = message.timestamp().to_millis().unwrap_or(i64::MIN);
                    let payload = message.payload().unwrap_or(&[]);
                    let metadata = self.create_metadata(&message);
                    let (buffer, errors) = parser.parse(payload, metadata);
                    self.n_bytes
                        .fetch_add(buffer.len().bytes, Ordering::Relaxed);
                    self.messages
                        .lock()
                        .unwrap()
                        .insert(offset, (buffer, timestamp));
                    consumer.parse_errors(errors);
                } else {
                    tracing::error!(
                        "Received message in partition {} at out-of-order offset {offset} (expected offset {next_offset} or greater; initial offset for this partition was {})",
                        self.partition,
                        self.initial_next_offset
                    );
                }
            }
        }
    }

    fn run(
        &self,
        base_consumer: &BaseConsumer<KafkaFtInputContext>,
        consumer: &dyn InputConsumer,
        parser: &mut dyn Parser,
    ) -> bool {
        // Limit per-partition queuing.  This should ultimately exert
        // backpressure on librdkafka reading from the partition (see [1]).
        //
        // The limit is pretty arbitrary.  It only needs to be high enough to
        // keep up with `KafkaFtInputReaderInner::poller_thread`, which serially
        // polls all of the partitions.
        //
        // For the unit tests only, we set an additional 1000-message limit
        // because that is easier to test precisely.  Generally speaking,
        // though, counting messages is undesirable because small messages will
        // be dequeued quicker than we can start queuing them again.
        //
        // [1]: https://github.com/confluentinc/librdkafka/wiki/FAQ#what-are-partition-queues-and-why-are-some-partitions-slower-than-others
        let backpressure = needs_backpressure(self.n_bytes.load(Ordering::Relaxed));
        #[cfg(test)]
        let backpressure = backpressure || self.messages.lock().unwrap().len() >= 1000;
        update_backpressure(&self.config.topic, self.partition, backpressure);
        if backpressure {
            return false;
        }

        let next_offset = self.next_offset();
        let max_offset = self.max_offset();
        if next_offset > max_offset {
            return false;
        }

        match self.queue.poll(Duration::ZERO) {
            Some(message) => {
                self.handle_kafka_message(base_consumer, consumer, parser, message);
                true
            }
            None => false,
        }
    }
}

struct RecvThreadHandle {
    exit: Arc<AtomicBool>,
    join_handle: Option<JoinHandle<()>>,
    unparker: Unparker,
}

impl RecvThreadHandle {
    fn new(mut thread: RecvThread) -> RecvThreadHandle {
        Self {
            exit: thread.exit.clone(),
            unparker: thread.parker.unparker().clone(),
            join_handle: {
                Some(spawn(move || {
                    //let _guard = span(&endpoint.config);
                    thread.run();
                }))
            },
        }
    }
}

impl Drop for RecvThreadHandle {
    /// When we're dropped, make the thread exit.
    ///
    /// *Careful*: `exit` is shared with all the helper threads, so if one gets
    /// dropped, the others will exit too. Currently this is OK because we
    /// always drop all of them together. It is in fact desirable because if
    /// they had separate `should_exit` flags then we'd have to block up to
    /// `POLL_TIMEOUT` per thread whereas since it is shared we will only block
    /// that long once.
    fn drop(&mut self) {
        self.exit.store(true, Ordering::Release);

        let _ = self.join_handle.take().map(|handle| {
            self.unparker.unpark();
            handle.join()
        });
    }
}

struct RecvThread {
    exit: Arc<AtomicBool>,
    topic: String,
    main_thread: Thread,
    parker: Parker,
    base_consumer: Arc<BaseConsumer<KafkaFtInputContext>>,
    consumer: Box<dyn InputConsumer>,
    parser: Box<dyn Parser>,
    receivers: Vec<Arc<PartitionReceiver>>,
}

impl RecvThread {
    fn run(&mut self) {
        let _guard = span(&self.topic);
        while !self.exit.load(Ordering::Relaxed) {
            let mut did_work = false;
            for receiver in &self.receivers {
                if receiver.run(&self.base_consumer, &*self.consumer, &mut *self.parser) {
                    did_work = true;
                }
            }
            if did_work {
                self.main_thread.unpark();
            } else {
                self.parker.park_timeout(Duration::from_secs(1));
            }
        }
    }
}
