use std::collections::{HashMap, HashSet, VecDeque};
use std::time::{Duration, Instant};

use anyhow::{Context, Result as AnyResult, anyhow, bail};
use kafka_protocol::error::{ParseResponseErrorCode, ResponseError};
use kafka_protocol::messages::find_coordinator_response::FindCoordinatorResponse;
use kafka_protocol::messages::init_producer_id_response::InitProducerIdResponse;
use kafka_protocol::messages::produce_response::ProduceResponse;
use kafka_protocol::messages::txn_offset_commit_response::TxnOffsetCommitResponse;
use kafka_protocol::messages::{
    EndTxnRequest, FindCoordinatorRequest, InitProducerIdRequest, ProduceRequest, ProducerId,
    TransactionalId, TxnOffsetCommitRequest,
};
use kafka_protocol::protocol::StrBytes;
use kafka_protocol::records::{NO_PRODUCER_EPOCH, NO_PRODUCER_ID};
use tokio::sync::{mpsc, oneshot};
use tokio::time::sleep;
use tracing::{Instrument, debug, trace, warn};

mod event;
mod metadata;

use super::ProducerRuntimeEvent;
use super::accumulator::{BatchProducerState, ProducerBatch, RecordAccumulator, fail_batch};
use super::partitioner::DefaultPartitionerState;
use super::request::{acks_for_request, build_produce_request};
use super::transaction::{
    ProducerIdentity, TransactionCoordinator, TransactionCoordinatorConnection,
    TransactionFailureDisposition, TransactionManager, TransactionResult,
    build_find_coordinator_request, build_group_find_coordinator_request,
    build_txn_offset_commit_request, classify_transactional_error, ensure_transaction_v2_feature,
    find_coordinator_error, parse_find_coordinator_response, validate_group_metadata,
};
use crate::config::ProducerConfig;
use crate::constants::{
    END_TXN_TRANSACTION_V2_MIN_VERSION, END_TXN_VERSION_CAP, FIND_COORDINATOR_VERSION_CAP,
    INIT_PRODUCER_ID_MIN_VERSION, INIT_PRODUCER_ID_VERSION_CAP, PRODUCE_TRANSACTION_V2_MIN_VERSION,
    PRODUCE_VERSION_CAP, TXN_OFFSET_COMMIT_TRANSACTION_V2_MIN_VERSION,
    TXN_OFFSET_COMMIT_VERSION_CAP,
};
use crate::network::{BrokerConnection, connect_to_any_bootstrap, duration_to_i32_ms};
use crate::types::{CommitOffset, ConsumerGroupMetadata, ProduceAck, TopicPartitionKey};
use crate::{Error, ProducerError, Result as ClientResult};
use event::producer_event_name;
use metadata::ProducerMetadata;

pub struct Sender {
    config: ProducerConfig,
    metadata: ProducerMetadata,
    connections: HashMap<i32, BrokerConnection>,
    accumulator: RecordAccumulator,
    application_events: VecDeque<ProducerRuntimeEvent>,
    pending_flushes: Vec<oneshot::Sender<ClientResult<()>>>,
    shutting_down: bool,
    close_reply: Option<oneshot::Sender<ClientResult<()>>>,
    transaction: Option<TransactionManager>,
    idempotent_producer: Option<ProducerIdentity>,
    idempotent_sequences: HashMap<TopicPartitionKey, i32>,
    transaction_coordinator_connection: Option<TransactionCoordinatorConnection>,
    partitioner: DefaultPartitionerState,
}

impl Sender {
    pub fn new(config: ProducerConfig) -> Self {
        let transaction = config.transactional_id.clone().map(|transactional_id| {
            TransactionManager::new(transactional_id, config.transaction_timeout)
        });

        Self {
            config,
            metadata: ProducerMetadata::default(),
            connections: HashMap::new(),
            accumulator: RecordAccumulator::default(),
            application_events: VecDeque::new(),
            pending_flushes: Vec::new(),
            shutting_down: false,
            close_reply: None,
            transaction,
            idempotent_producer: None,
            idempotent_sequences: HashMap::new(),
            transaction_coordinator_connection: None,
            partitioner: DefaultPartitionerState::default(),
        }
    }

    pub async fn run(mut self, mut rx: mpsc::Receiver<ProducerRuntimeEvent>) {
        let client_id = self.config.client_id.clone();
        let transactional = self.transaction.is_some();
        async move {
            let mut running = true;

            while running {
                let wake = self.next_wakeup();
                tokio::select! {
                    biased;
                    maybe_event = rx.recv() => {
                        match maybe_event {
                            Some(event) => self.application_events.push_back(event),
                            None => {
                                debug!("producer application event channel closed");
                                self.shutting_down = true;
                            }
                        }
                    }
                    _ = sleep(wake) => {}
                }

                running = self.run_once(&mut rx).await;
            }

            self.accumulator.fail_all("producer sender thread stopped");
            if let Some(transaction) = self.transaction.as_mut() {
                transaction.fail_pending("producer sender thread stopped");
            }
            self.fail_pending_flushes("producer sender thread stopped");
            if let Some(reply) = self.close_reply.take() {
                let _ = reply.send(Err(Error::Internal(anyhow!(
                    "producer sender thread stopped"
                ))));
            }
        }
        .instrument(tracing::debug_span!(
            "producer_sender_thread",
            %client_id,
            transactional
        ))
        .await;
    }

    async fn run_once(&mut self, rx: &mut mpsc::Receiver<ProducerRuntimeEvent>) -> bool {
        self.drain_application_events(rx);
        if self.pending_flushes.is_empty() {
            self.process_application_events().await;
        }
        self.fail_expired_batches();
        self.maybe_refresh_tracked_topics().await;
        if let Err(error) = self.ensure_idempotent_producer_initialized().await {
            let message = format!("{error:#}");
            self.accumulator.fail_all(&message);
        }
        self.send_ready_batches().await;
        self.maybe_complete_flushes();
        self.maybe_complete_transaction().await;
        self.finish_shutdown_if_ready()
    }

    fn drain_application_events(&mut self, rx: &mut mpsc::Receiver<ProducerRuntimeEvent>) {
        while let Ok(event) = rx.try_recv() {
            self.application_events.push_back(event);
        }
    }

    async fn process_application_events(&mut self) {
        while let Some(event) = self.application_events.pop_front() {
            trace!(
                event = producer_event_name(&event),
                pending_events = self.application_events.len(),
                "processing producer application event"
            );
            match event {
                ProducerRuntimeEvent::WarmUp { reply } => {
                    let result = connect_to_any_bootstrap(
                        &self.config.bootstrap_servers,
                        &self.config.client_id,
                        self.config.request_timeout,
                        self.config.security_protocol,
                        &self.config.tls,
                        &self.config.sasl,
                    )
                    .await
                    .map(|_| ())
                    .map_err(Error::from);
                    let _ = reply.send(result);
                }
                ProducerRuntimeEvent::BeginTransaction { reply } => {
                    let _ = reply.send(self.begin_transaction().await);
                }
                ProducerRuntimeEvent::InitTransactions { reply } => {
                    let _ = reply.send(self.ensure_transaction_initialized().await);
                }
                ProducerRuntimeEvent::AppendRecord { record, reply } => {
                    if let Some(transaction) = self.transaction.as_ref()
                        && let Err(error) = transaction.ensure_append_allowed()
                    {
                        let _ = reply.send(Err(error.into()));
                        continue;
                    }

                    if let Err(error) = self.append_record(record, reply).await {
                        warn!(error = %error, "failed to append produce record");
                    }
                }
                ProducerRuntimeEvent::Flush { reply } => {
                    if self.accumulator.has_undrained() {
                        self.pending_flushes.push(reply);
                        break;
                    }
                    let _ = reply.send(Ok(()));
                }
                ProducerRuntimeEvent::CommitTransaction { reply } => {
                    let mut reply = Some(reply);
                    if let Err(error) =
                        self.request_transaction_completion(TransactionResult::Commit, reply.take())
                        && let Some(reply) = reply
                    {
                        let _ = reply.send(Err(error));
                    }
                }
                ProducerRuntimeEvent::AbortTransaction { reply } => {
                    let mut reply = Some(reply);
                    if let Err(error) =
                        self.request_transaction_completion(TransactionResult::Abort, reply.take())
                        && let Some(reply) = reply
                    {
                        let _ = reply.send(Err(error));
                    }
                }
                ProducerRuntimeEvent::SendOffsetsToTransaction {
                    offsets,
                    group_metadata,
                    reply,
                } => {
                    if self.accumulator.has_undrained() {
                        self.application_events.push_front(
                            ProducerRuntimeEvent::SendOffsetsToTransaction {
                                offsets,
                                group_metadata,
                                reply,
                            },
                        );
                        break;
                    }
                    let _ = reply.send(
                        self.send_offsets_to_transaction(offsets, group_metadata)
                            .await,
                    );
                }
                ProducerRuntimeEvent::Shutdown { reply } => {
                    self.shutting_down = true;
                    self.close_reply = Some(reply);
                    if let Some(transaction) = self.transaction.as_mut() {
                        transaction.ensure_shutdown_abort();
                    }
                }
            }
        }
    }

    async fn begin_transaction(&mut self) -> ClientResult<()> {
        if self.transaction.is_none() {
            return Err(ProducerError::NotTransactional.into());
        }

        self.ensure_transaction_initialized().await?;
        debug!("beginning producer transaction");
        self.transaction
            .as_mut()
            .context("transaction manager disappeared")?
            .begin()?;
        Ok(())
    }

    fn request_transaction_completion(
        &mut self,
        result: TransactionResult,
        reply: Option<oneshot::Sender<ClientResult<()>>>,
    ) -> ClientResult<()> {
        let transaction = self
            .transaction
            .as_mut()
            .ok_or(ProducerError::NotTransactional)?;
        transaction.request_completion(result, reply)?;
        Ok(())
    }

    async fn append_record(
        &mut self,
        mut record: crate::types::ProduceRecord,
        reply: oneshot::Sender<ClientResult<ProduceAck>>,
    ) -> ClientResult<()> {
        self.metadata.track(record.topic.clone());
        let partition = self.resolve_partition_for_record(&record).await?;
        record.partition = Some(partition);
        self.accumulator
            .append(record, reply, self.config.batch_size);
        Ok(())
    }

    async fn resolve_partition_for_record(
        &mut self,
        record: &crate::types::ProduceRecord,
    ) -> ClientResult<i32> {
        if let Some(partition) = record.partition {
            return Ok(partition);
        }

        let partitions = self.partitions_for_topic(&record.topic).await?;
        if partitions.is_empty() {
            return Err(Error::Internal(anyhow!(
                "topic '{}' has no partitions in metadata",
                record.topic
            )));
        }

        Ok(self.partitioner.partition(
            self.config.partitioner,
            record,
            &partitions,
            &self.accumulator,
            self.config.batch_size,
        ))
    }

    async fn partitions_for_topic(
        &mut self,
        topic: &str,
    ) -> ClientResult<Vec<(i32, crate::metadata::PartitionMetadata)>> {
        if self.metadata.cache().partitions_for(topic).is_none()
            || self
                .metadata
                .cache()
                .needs_refresh(topic, self.config.metadata_max_age)
        {
            self.metadata
                .refresh_topics(&self.config, vec![topic.to_owned()])
                .await?;
            self.retain_connections_for_metadata();
        }

        self.metadata
            .cache()
            .partitions_for(topic)
            .with_context(|| format!("missing partition metadata for topic '{topic}'"))
            .map_err(Error::from)
    }

    async fn ensure_transaction_initialized(&mut self) -> ClientResult<()> {
        if self
            .transaction
            .as_ref()
            .and_then(|transaction| transaction.producer)
            .is_some()
        {
            return Ok(());
        }

        let transactional_id = self.transactional_id()?.to_owned();
        let transaction_timeout = self
            .transaction
            .as_ref()
            .map(|transaction| transaction.transaction_timeout)
            .context("transaction manager disappeared")?;

        let coordinator = self.find_transaction_coordinator(&transactional_id).await?;
        let producer = self
            .init_transactional_producer(&transactional_id, transaction_timeout)
            .await?;

        let transaction = self
            .transaction
            .as_mut()
            .context("transaction manager disappeared")?;
        transaction.coordinator = Some(coordinator);
        transaction.reset_with_new_producer(producer);
        debug!(%transactional_id, producer_id = producer.id, producer_epoch = producer.epoch, "initialized transactional producer");
        Ok(())
    }

    async fn ensure_idempotent_producer_initialized(&mut self) -> ClientResult<()> {
        if self.transaction.is_some() || !self.config.is_idempotent() {
            return Ok(());
        }

        if self.idempotent_producer.is_some() {
            return Ok(());
        }

        let producer = self.init_idempotent_producer().await?;
        self.idempotent_producer = Some(producer);
        self.idempotent_sequences.clear();
        debug!(
            producer_id = producer.id,
            producer_epoch = producer.epoch,
            "initialized idempotent producer"
        );
        Ok(())
    }

    async fn send_offsets_to_transaction(
        &mut self,
        offsets: Vec<CommitOffset>,
        group_metadata: ConsumerGroupMetadata,
    ) -> ClientResult<()> {
        validate_group_metadata(&group_metadata)?;
        let transaction = self
            .transaction
            .as_ref()
            .ok_or(ProducerError::NotTransactional)?;
        transaction.ensure_send_offsets_allowed()?;
        self.ensure_transaction_initialized().await?;

        let producer = self
            .transaction
            .as_ref()
            .and_then(|transaction| transaction.producer)
            .context("transactional producer was not initialized")?;

        self.commit_transactional_offsets(&offsets, &group_metadata, producer)
            .await?;
        Ok(())
    }

    async fn commit_transactional_offsets(
        &mut self,
        offsets: &[CommitOffset],
        group_metadata: &ConsumerGroupMetadata,
        producer: ProducerIdentity,
    ) -> ClientResult<()> {
        let transactional_id = self.transactional_id()?.to_owned();
        let attempts = self.config.max_retries.max(10);

        for attempt in 0..=attempts {
            let coordinator = self
                .find_group_coordinator(&group_metadata.group_id)
                .await?;
            let mut connection = BrokerConnection::connect_with_transport(
                &coordinator.address,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;
            ensure_transaction_v2_feature(&connection)?;
            let version = connection
                .version_with_cap::<TxnOffsetCommitRequest>(TXN_OFFSET_COMMIT_VERSION_CAP)?;
            if version < TXN_OFFSET_COMMIT_TRANSACTION_V2_MIN_VERSION {
                return Err(ProducerError::UnsupportedApiVersion {
                    api: "TxnOffsetCommit",
                    min_version: TXN_OFFSET_COMMIT_TRANSACTION_V2_MIN_VERSION,
                    broker_version: version,
                }
                .into());
            }

            let request = build_txn_offset_commit_request(
                &transactional_id,
                producer,
                offsets,
                group_metadata,
            );
            let response: TxnOffsetCommitResponse = connection
                .send_request::<TxnOffsetCommitRequest>(&self.config.client_id, version, &request)
                .await?;

            let response_error = response.topics.iter().find_map(|topic| {
                topic
                    .partitions
                    .iter()
                    .find_map(|partition| partition.error_code.err())
            });

            let Some(error) = response_error else {
                return Ok(());
            };

            if attempt < attempts
                && matches!(
                    error,
                    ResponseError::CoordinatorLoadInProgress
                        | ResponseError::CoordinatorNotAvailable
                        | ResponseError::NotCoordinator
                )
            {
                sleep(self.config.retry_backoff).await;
                continue;
            }

            let message = format!("TxnOffsetCommit failed: {error}");
            match classify_transactional_error(error) {
                TransactionFailureDisposition::AbortOnly => {
                    self.mark_transaction_abort_only(message.clone());
                }
                TransactionFailureDisposition::Fatal => {
                    self.mark_transaction_fatal(message.clone())
                }
            }
            return Err(Error::Internal(anyhow!(message)));
        }

        Err(Error::Internal(anyhow!(
            "failed to commit transactional offsets for group '{}' after {} attempts",
            group_metadata.group_id,
            attempts + 1
        )))
    }

    async fn maybe_refresh_tracked_topics(&mut self) {
        if !self
            .metadata
            .needs_periodic_refresh(self.config.metadata_max_age)
        {
            return;
        }

        if let Err(error) = self.metadata.refresh_tracked(&self.config).await {
            warn!(error = %error, "producer metadata refresh failed");
            return;
        }
        self.retain_connections_for_metadata();
    }

    async fn send_ready_batches(&mut self) {
        loop {
            let ready =
                self.accumulator
                    .ready(self.metadata.cache(), &self.config, self.shutting_down);

            if !ready.unknown_leaders.is_empty() {
                self.resolve_unknown_leaders(ready.unknown_leaders).await;
                continue;
            }

            let drained = self.accumulator.drain_ready(ready.ready_by_leader);
            if drained.is_empty() {
                break;
            }

            debug!(
                leader_count = drained.len(),
                batch_count = drained.values().map(Vec::len).sum::<usize>(),
                "sending ready producer batches"
            );

            for (leader_id, mut batches) in drained {
                if let Err(error) = self.prepare_batches_for_send(&mut batches) {
                    let message = format!("{error:#}");
                    self.fail_batches(batches, &message);
                    self.mark_transaction_abort_only(message);
                    continue;
                }

                self.send_batches_to_leader(leader_id, batches).await;
            }
        }
    }

    fn prepare_batches_for_send(&mut self, batches: &mut [ProducerBatch]) -> AnyResult<()> {
        for batch in batches {
            if batch.producer_state.is_some() {
                continue;
            }

            batch.producer_state =
                self.producer_state_for_batch(&batch.key, batch.records.len())?;
        }

        Ok(())
    }

    fn producer_state_for_batch(
        &mut self,
        key: &TopicPartitionKey,
        record_count: usize,
    ) -> AnyResult<Option<BatchProducerState>> {
        if let Some(transaction) = self.transaction.as_mut() {
            let producer = transaction
                .producer
                .context("transactional producer was not initialized")?;
            let next_sequence = transaction.sequences.entry(key.clone()).or_insert(0);
            let batch_state = BatchProducerState {
                producer_id: producer.id,
                producer_epoch: producer.epoch,
                base_sequence: *next_sequence,
                transactional: true,
            };
            *next_sequence = next_sequence.saturating_add(i32::try_from(record_count).unwrap_or(0));
            return Ok(Some(batch_state));
        }

        if let Some(producer) = self.idempotent_producer {
            let next_sequence = self.idempotent_sequences.entry(key.clone()).or_insert(0);
            let batch_state = BatchProducerState {
                producer_id: producer.id,
                producer_epoch: producer.epoch,
                base_sequence: *next_sequence,
                transactional: false,
            };
            *next_sequence = next_sequence.saturating_add(i32::try_from(record_count).unwrap_or(0));
            return Ok(Some(batch_state));
        }

        Ok(None)
    }

    async fn resolve_unknown_leaders(&mut self, partitions: Vec<TopicPartitionKey>) {
        let topics: HashSet<String> = partitions.iter().map(|key| key.topic.clone()).collect();
        debug!(
            topic_count = topics.len(),
            partition_count = partitions.len(),
            "refreshing metadata for producer partitions without leaders"
        );
        if let Err(error) = self
            .metadata
            .refresh_topics(&self.config, topics.iter().cloned().collect())
            .await
        {
            warn!(error = %error, "producer metadata refresh for unknown leaders failed");
        } else {
            self.retain_connections_for_metadata();
        }

        let unresolved = partitions
            .into_iter()
            .filter(|key| {
                self.metadata
                    .leader_for(&key.topic, key.partition)
                    .is_none()
            })
            .collect();

        for batch in self.accumulator.drain_unknown_leader(unresolved) {
            let topic = batch.key.topic.clone();
            let partition = batch.key.partition;
            self.metadata.invalidate_topic(&batch.key.topic);
            self.retry_or_fail_batch(
                batch,
                format!(
                    "no leader metadata for topic '{}' partition {}",
                    topic, partition
                ),
            );
        }
    }

    async fn send_batches_to_leader(&mut self, leader_id: i32, batches: Vec<ProducerBatch>) {
        debug!(
            leader_id,
            batch_count = batches.len(),
            record_count = batches
                .iter()
                .map(|batch| batch.records.len())
                .sum::<usize>(),
            "dispatching produce request to leader"
        );
        let transactional_id = self.transactional_id_if_active().map(str::to_owned);
        let request = match build_produce_request(
            &batches,
            acks_for_request(self.config.acks, transactional_id.is_some()),
            self.config.request_timeout,
            self.config.compression,
            transactional_id.as_deref(),
        ) {
            Ok(request) => request,
            Err(error) => {
                self.fail_batches(batches, &format!("{error:#}"));
                self.mark_transaction_abort_only(format!("{error:#}"));
                return;
            }
        };

        let client_id = self.config.client_id.clone();
        let response = match self
            .send_produce_request(leader_id, &client_id, request, transactional_id.is_some())
            .await
        {
            Ok(response) => response,
            Err(error) => {
                self.handle_transport_failure(leader_id, batches, error);
                return;
            }
        };

        self.handle_produce_response(leader_id, batches, response);
    }

    async fn send_produce_request(
        &mut self,
        leader_id: i32,
        client_id: &str,
        request: ProduceRequest,
        transactional: bool,
    ) -> AnyResult<ProduceResponse> {
        let connection = self.connection_for(leader_id).await?;
        let version = connection.version_with_cap::<ProduceRequest>(PRODUCE_VERSION_CAP)?;
        if transactional && version < PRODUCE_TRANSACTION_V2_MIN_VERSION {
            bail!(
                "broker only supports Produce v{}, but transaction v2 requires Produce v{}+",
                version,
                PRODUCE_TRANSACTION_V2_MIN_VERSION
            );
        }

        connection
            .send_request::<ProduceRequest>(client_id, version, &request)
            .await
    }

    fn handle_produce_response(
        &mut self,
        leader_id: i32,
        batches: Vec<ProducerBatch>,
        response: ProduceResponse,
    ) {
        let mut in_flight = batches
            .into_iter()
            .map(|batch| (batch.key.clone(), batch))
            .collect::<HashMap<_, _>>();

        for topic_response in response.responses {
            let topic = topic_response.name.0.to_string();
            for partition_response in topic_response.partition_responses {
                let key = TopicPartitionKey::new(topic.clone(), partition_response.index);
                let Some(batch) = in_flight.remove(&key) else {
                    continue;
                };

                if let Some(error) = partition_response.error_code.err() {
                    self.handle_kafka_failure(leader_id, batch, error);
                    continue;
                }

                self.complete_batch_success(
                    batch,
                    partition_response.base_offset,
                    partition_response.log_append_time_ms,
                );
            }
        }

        for batch in in_flight.into_values() {
            let topic = batch.key.topic.clone();
            let partition = batch.key.partition;
            self.handle_transport_failure(
                leader_id,
                vec![batch],
                anyhow!("missing produce response for {}:{}", topic, partition),
            );
        }
    }

    fn complete_batch_success(
        &self,
        mut batch: ProducerBatch,
        base_offset: i64,
        log_append_time_ms: i64,
    ) {
        for (index, reply) in batch.replies.drain(..).enumerate() {
            let _ = reply.send(Ok(ProduceAck {
                topic: batch.key.topic.clone(),
                partition: batch.key.partition,
                base_offset,
                offset: base_offset + i64::try_from(index).unwrap_or(0),
                log_append_time_ms,
            }));
        }
    }

    fn handle_kafka_failure(&mut self, leader_id: i32, batch: ProducerBatch, error: ResponseError) {
        warn!(
            leader_id,
            topic = %batch.key.topic,
            partition = batch.key.partition,
            retriable = error.is_retriable(),
            error = %error,
            "producer batch failed with broker error"
        );
        if error.is_retriable() {
            self.connections.remove(&leader_id);
            self.metadata.invalidate_topic(&batch.key.topic);
            self.retry_or_fail_batch(batch, format!("produce rejected by broker: {error}"));
            return;
        }

        let message = format!("produce rejected by broker: {error}");
        if self.transaction.is_some() && batch.producer_state.is_some() {
            fail_batch(batch, &message);
            self.accumulator.fail_all(&message);
            match classify_transactional_error(error) {
                TransactionFailureDisposition::AbortOnly => {
                    self.mark_transaction_abort_only(message)
                }
                TransactionFailureDisposition::Fatal => self.mark_transaction_fatal(message),
            }
            return;
        }

        fail_batch(batch, &message);
    }

    fn handle_transport_failure(
        &mut self,
        leader_id: i32,
        batches: Vec<ProducerBatch>,
        error: anyhow::Error,
    ) {
        warn!(
            leader_id,
            batch_count = batches.len(),
            error = %error,
            "producer transport failure while sending batches"
        );
        self.connections.remove(&leader_id);
        let message = format!("{error:#}");
        for batch in batches {
            self.metadata.invalidate_topic(&batch.key.topic);
            self.retry_or_fail_batch(batch, message.clone());
        }
    }

    fn retry_or_fail_batch(&mut self, mut batch: ProducerBatch, message: String) {
        if batch.created_at.elapsed() >= self.config.delivery_timeout {
            let timeout_message = format!(
                "delivery timed out after {:?}: {}",
                self.config.delivery_timeout, message
            );
            warn!(
                topic = %batch.key.topic,
                partition = batch.key.partition,
                attempts = batch.attempts,
                "producer batch exceeded delivery timeout"
            );
            self.fail_delivery_timeout(batch, timeout_message);
            return;
        }

        if batch.attempts >= self.config.max_retries {
            warn!(
                topic = %batch.key.topic,
                partition = batch.key.partition,
                attempts = batch.attempts,
                max_retries = self.config.max_retries,
                "producer batch exhausted retries"
            );
            if self.transaction.is_some() && batch.producer_state.is_some() {
                fail_batch(batch, &message);
                self.accumulator.fail_all(&message);
                self.mark_transaction_abort_only(message);
            } else {
                fail_batch(batch, &message);
            }
            return;
        }

        batch.attempts += 1;
        batch.retry_at = Some(Instant::now() + self.config.retry_backoff);
        debug!(
            topic = %batch.key.topic,
            partition = batch.key.partition,
            attempts = batch.attempts,
            retry_backoff_ms = self.config.retry_backoff.as_millis(),
            "reenqueuing producer batch for retry"
        );
        self.accumulator.reenqueue_front(batch);
    }

    fn fail_expired_batches(&mut self) {
        let expired = self
            .accumulator
            .expire_batches(self.config.delivery_timeout);
        for batch in expired {
            let message = format!(
                "delivery timed out after {:?} for {}:{}",
                self.config.delivery_timeout, batch.key.topic, batch.key.partition
            );
            self.fail_delivery_timeout(batch, message);
        }
    }

    fn fail_delivery_timeout(&mut self, batch: ProducerBatch, message: String) {
        if self.transaction.is_some() && batch.producer_state.is_some() {
            fail_batch(batch, &message);
            self.accumulator.fail_all(&message);
            self.mark_transaction_abort_only(message);
            return;
        }

        fail_batch(batch, &message);
    }

    fn fail_batches(&self, batches: Vec<ProducerBatch>, message: &str) {
        for batch in batches {
            fail_batch(batch, message);
        }
    }

    fn maybe_complete_flushes(&mut self) {
        if self.pending_flushes.is_empty() || self.accumulator.has_undrained() {
            return;
        }

        let completed = self.pending_flushes.len();
        for reply in std::mem::take(&mut self.pending_flushes) {
            let _ = reply.send(Ok(()));
        }
        debug!(completed, "completed producer flush waiters");
    }

    fn fail_pending_flushes(&mut self, message: &str) {
        for reply in std::mem::take(&mut self.pending_flushes) {
            let _ = reply.send(Err(Error::Internal(anyhow!(message.to_owned()))));
        }
    }

    async fn connection_for(&mut self, broker_id: i32) -> AnyResult<&mut BrokerConnection> {
        if !self.connections.contains_key(&broker_id) {
            let broker = self
                .metadata
                .broker(broker_id)
                .cloned()
                .with_context(|| format!("broker {broker_id} is missing from metadata"))?;
            let address = broker.address();
            let connection = BrokerConnection::connect_with_transport(
                &address,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;
            self.connections.insert(broker_id, connection);
            trace!(broker_id, %address, "opened producer broker connection");
        }

        self.connections
            .get_mut(&broker_id)
            .with_context(|| format!("broker connection {broker_id} was not created"))
    }

    fn retain_connections_for_metadata(&mut self) {
        self.connections
            .retain(|broker_id, _| self.metadata.contains_broker(*broker_id));
    }

    async fn maybe_complete_transaction(&mut self) {
        let Some(result) = self.transaction.as_ref().and_then(|transaction| {
            transaction
                .pending_completion
                .as_ref()
                .map(|pending| pending.result)
        }) else {
            return;
        };

        if self.accumulator.has_undrained() {
            return;
        }

        debug!(
            result = if matches!(result, TransactionResult::Commit) {
                "commit"
            } else {
                "abort"
            },
            "completing producer transaction"
        );

        let reply = self
            .transaction
            .as_mut()
            .and_then(|transaction| transaction.pending_completion.take())
            .and_then(|pending| pending.reply);

        let outcome = self.end_transaction(result).await;
        match outcome {
            Ok(()) => {
                debug!("producer transaction completed");
                if let Some(reply) = reply {
                    let _ = reply.send(Ok(()));
                }
            }
            Err(error) => {
                let message = format!("{error:#}");
                self.mark_transaction_fatal(message.clone());
                warn!(error = %message, "producer transaction completion failed");
                if let Some(reply) = reply {
                    let _ = reply.send(Err(Error::Internal(anyhow!(message))));
                }
            }
        }
    }

    async fn end_transaction(&mut self, result: TransactionResult) -> AnyResult<()> {
        let transactional_id = self.transactional_id()?.to_owned();
        let producer = self
            .transaction
            .as_ref()
            .and_then(|transaction| transaction.producer)
            .context("transactional producer was not initialized")?;

        let version = {
            let connection = self.transaction_coordinator_connection().await?;
            let version = connection.version_with_cap::<EndTxnRequest>(END_TXN_VERSION_CAP)?;
            if version < END_TXN_TRANSACTION_V2_MIN_VERSION {
                bail!(
                    "transaction v2 requires EndTxn v{}+, broker only supports EndTxn v{}",
                    END_TXN_TRANSACTION_V2_MIN_VERSION,
                    version
                );
            }
            version
        };

        let request = EndTxnRequest::default()
            .with_transactional_id(TransactionalId(StrBytes::from_string(transactional_id)))
            .with_producer_id(ProducerId(producer.id))
            .with_producer_epoch(producer.epoch)
            .with_committed(matches!(result, TransactionResult::Commit));

        let response = {
            let client_id = self.config.client_id.clone();
            let connection = self.transaction_coordinator_connection().await?;
            connection
                .send_request::<EndTxnRequest>(&client_id, version, &request)
                .await?
        };

        if let Some(error) = response.error_code.err() {
            if matches!(
                error,
                ResponseError::CoordinatorNotAvailable | ResponseError::NotCoordinator
            ) {
                self.invalidate_transaction_coordinator();
            }
            bail!("EndTxn failed: {error}");
        }

        let next_producer = ProducerIdentity {
            id: *response.producer_id,
            epoch: response.producer_epoch,
        };
        let transaction = self
            .transaction
            .as_mut()
            .context("transaction manager disappeared")?;
        transaction.finish_success(next_producer);
        Ok(())
    }

    async fn find_transaction_coordinator(
        &mut self,
        transactional_id: &str,
    ) -> AnyResult<TransactionCoordinator> {
        let attempts = self.config.max_retries.max(10);
        for attempt in 0..=attempts {
            let mut bootstrap = connect_to_any_bootstrap(
                &self.config.bootstrap_servers,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;

            let version = bootstrap
                .version_with_cap::<FindCoordinatorRequest>(FIND_COORDINATOR_VERSION_CAP)?;
            let request = build_find_coordinator_request(transactional_id, version);
            let response: FindCoordinatorResponse = bootstrap
                .send_request::<FindCoordinatorRequest>(&self.config.client_id, version, &request)
                .await?;

            if let Some(error) = find_coordinator_error(&response, version)
                && attempt < attempts
                && matches!(
                    error,
                    ResponseError::CoordinatorNotAvailable
                        | ResponseError::CoordinatorLoadInProgress
                        | ResponseError::NotCoordinator
                )
            {
                sleep(self.config.retry_backoff).await;
                continue;
            }

            let coordinator = parse_find_coordinator_response(response, version)?;
            self.transaction_coordinator_connection = None;
            if let Some(transaction) = self.transaction.as_mut() {
                transaction.coordinator = Some(coordinator.clone());
            }
            debug!(
                %transactional_id,
                coordinator = %coordinator.address,
                "resolved transaction coordinator"
            );
            return Ok(coordinator);
        }

        bail!(
            "failed to resolve transaction coordinator for '{}' after {} attempts",
            transactional_id,
            attempts + 1
        )
    }

    async fn find_group_coordinator(
        &mut self,
        group_id: &str,
    ) -> AnyResult<TransactionCoordinator> {
        let attempts = self.config.max_retries.max(10);
        for attempt in 0..=attempts {
            let mut bootstrap = connect_to_any_bootstrap(
                &self.config.bootstrap_servers,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;

            let version = bootstrap
                .version_with_cap::<FindCoordinatorRequest>(FIND_COORDINATOR_VERSION_CAP)?;
            let request = build_group_find_coordinator_request(group_id, version);
            let response: FindCoordinatorResponse = bootstrap
                .send_request::<FindCoordinatorRequest>(&self.config.client_id, version, &request)
                .await?;

            if let Some(error) = find_coordinator_error(&response, version)
                && attempt < attempts
                && matches!(
                    error,
                    ResponseError::CoordinatorNotAvailable
                        | ResponseError::CoordinatorLoadInProgress
                        | ResponseError::NotCoordinator
                )
            {
                sleep(self.config.retry_backoff).await;
                continue;
            }

            let coordinator = parse_find_coordinator_response(response, version)?;
            debug!(%group_id, coordinator = %coordinator.address, "resolved group coordinator");
            return Ok(coordinator);
        }

        bail!(
            "failed to resolve group coordinator for '{}' after {} attempts",
            group_id,
            attempts + 1
        )
    }

    async fn init_transactional_producer(
        &mut self,
        transactional_id: &str,
        transaction_timeout: Duration,
    ) -> ClientResult<ProducerIdentity> {
        let attempts = self.config.max_retries.max(10);
        for attempt in 0..=attempts {
            let version = {
                let connection = self.transaction_coordinator_connection().await?;
                ensure_transaction_v2_feature(connection)?;
                let version = connection
                    .version_with_cap::<InitProducerIdRequest>(INIT_PRODUCER_ID_VERSION_CAP)?;
                if version < INIT_PRODUCER_ID_MIN_VERSION {
                    return Err(ProducerError::UnsupportedApiVersion {
                        api: "InitProducerId",
                        min_version: INIT_PRODUCER_ID_MIN_VERSION,
                        broker_version: version,
                    }
                    .into());
                }
                version
            };

            let request = InitProducerIdRequest::default()
                .with_transactional_id(Some(TransactionalId(StrBytes::from_string(
                    transactional_id.to_owned(),
                ))))
                .with_transaction_timeout_ms(duration_to_i32_ms(transaction_timeout)?)
                .with_producer_id(ProducerId(NO_PRODUCER_ID))
                .with_producer_epoch(NO_PRODUCER_EPOCH);

            let response: InitProducerIdResponse = {
                let client_id = self.config.client_id.clone();
                let connection = self.transaction_coordinator_connection().await?;
                connection
                    .send_request::<InitProducerIdRequest>(&client_id, version, &request)
                    .await?
            };

            if let Some(error) = response.error_code.err() {
                if attempt < attempts
                    && matches!(
                        error,
                        ResponseError::NotCoordinator
                            | ResponseError::CoordinatorNotAvailable
                            | ResponseError::CoordinatorLoadInProgress
                    )
                {
                    self.invalidate_transaction_coordinator();
                    sleep(self.config.retry_backoff).await;
                    continue;
                }

                return Err(Error::Internal(anyhow!("InitProducerId failed: {error}")));
            }

            return Ok(ProducerIdentity {
                id: *response.producer_id,
                epoch: response.producer_epoch,
            });
        }

        Err(Error::Internal(anyhow!(
            "failed to initialize producer id for transactional.id '{}' after {} attempts",
            transactional_id,
            attempts + 1
        )))
    }

    async fn init_idempotent_producer(&mut self) -> ClientResult<ProducerIdentity> {
        let attempts = self.config.max_retries.max(10);
        for attempt in 0..=attempts {
            let mut bootstrap = connect_to_any_bootstrap(
                &self.config.bootstrap_servers,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;
            let version = bootstrap
                .version_with_cap::<InitProducerIdRequest>(INIT_PRODUCER_ID_VERSION_CAP)?;
            if version < INIT_PRODUCER_ID_MIN_VERSION {
                return Err(ProducerError::UnsupportedApiVersion {
                    api: "InitProducerId",
                    min_version: INIT_PRODUCER_ID_MIN_VERSION,
                    broker_version: version,
                }
                .into());
            }

            let request = InitProducerIdRequest::default()
                .with_transactional_id(None)
                .with_transaction_timeout_ms(duration_to_i32_ms(self.config.transaction_timeout)?)
                .with_producer_id(ProducerId(NO_PRODUCER_ID))
                .with_producer_epoch(NO_PRODUCER_EPOCH);

            let response: InitProducerIdResponse = bootstrap
                .send_request::<InitProducerIdRequest>(&self.config.client_id, version, &request)
                .await?;

            if let Some(error) = response.error_code.err() {
                if attempt < attempts && error.is_retriable() {
                    sleep(self.config.retry_backoff).await;
                    continue;
                }

                return Err(Error::Internal(anyhow!("InitProducerId failed: {error}")));
            }

            return Ok(ProducerIdentity {
                id: *response.producer_id,
                epoch: response.producer_epoch,
            });
        }

        Err(Error::Internal(anyhow!(
            "failed to initialize producer id after {} attempts",
            attempts + 1
        )))
    }

    async fn transaction_coordinator_connection(&mut self) -> AnyResult<&mut BrokerConnection> {
        if self.transaction.is_none() {
            bail!("producer is not configured with a transactional_id");
        }

        let transactional_id = self.transactional_id()?.to_owned();
        let coordinator = match self
            .transaction
            .as_ref()
            .and_then(|transaction| transaction.coordinator.clone())
        {
            Some(coordinator) => coordinator,
            None => self.find_transaction_coordinator(&transactional_id).await?,
        };

        let reconnect = !matches!(
            self.transaction_coordinator_connection.as_ref(),
            Some(existing)
                if existing.broker_id == coordinator.broker_id && existing.address == coordinator.address
        );

        if reconnect {
            let connection = BrokerConnection::connect_with_transport(
                &coordinator.address,
                &self.config.client_id,
                self.config.request_timeout,
                self.config.security_protocol,
                &self.config.tls,
                &self.config.sasl,
            )
            .await?;
            self.transaction_coordinator_connection = Some(TransactionCoordinatorConnection {
                broker_id: coordinator.broker_id,
                address: coordinator.address.clone(),
                connection,
            });
            trace!(
                broker_id = coordinator.broker_id,
                address = %coordinator.address,
                "opened transaction coordinator connection"
            );
        }

        self.transaction_coordinator_connection
            .as_mut()
            .map(|coordinator| &mut coordinator.connection)
            .context("transaction coordinator connection was not created")
    }

    fn invalidate_transaction_coordinator(&mut self) {
        self.transaction_coordinator_connection = None;
        if let Some(transaction) = self.transaction.as_mut() {
            transaction.coordinator = None;
        }
        debug!("invalidated transaction coordinator connection");
    }

    fn mark_transaction_abort_only(&mut self, message: String) {
        if let Some(transaction) = self.transaction.as_mut() {
            transaction.mark_abort_only(message);
        }
        warn!("producer transaction moved to abort-only state");
    }

    fn mark_transaction_fatal(&mut self, message: String) {
        if let Some(transaction) = self.transaction.as_mut() {
            transaction.mark_fatal(message);
        }
        warn!("producer transaction moved to fatal state");
    }

    fn transactional_id_if_active(&self) -> Option<&str> {
        self.transaction
            .as_ref()
            .and_then(TransactionManager::transactional_id_if_active)
    }

    fn transactional_id(&self) -> AnyResult<&str> {
        self.transaction
            .as_ref()
            .map(|transaction| transaction.transactional_id.as_str())
            .context("producer is not configured with a transactional_id")
    }

    fn finish_shutdown_if_ready(&mut self) -> bool {
        let has_pending_txn_completion = self
            .transaction
            .as_ref()
            .and_then(|transaction| transaction.pending_completion.as_ref())
            .is_some();

        if !self.shutting_down || self.accumulator.has_undrained() || has_pending_txn_completion {
            return true;
        }

        if let Some(reply) = self.close_reply.take() {
            let result = self
                .transaction
                .as_ref()
                .map(|transaction| transaction.shutdown_result().map_err(Error::from))
                .unwrap_or(Ok(()));
            let _ = reply.send(result);
        }
        false
    }

    fn next_wakeup(&self) -> Duration {
        if !self.application_events.is_empty() {
            return Duration::ZERO;
        }

        let mut wake = Duration::from_millis(250);
        wake = wake.min(self.metadata.next_refresh_in(self.config.metadata_max_age));
        if let Some(expiration_wake) = self
            .accumulator
            .next_expiration_in(self.config.delivery_timeout)
        {
            wake = wake.min(expiration_wake);
        }

        let ready = self
            .accumulator
            .ready(self.metadata.cache(), &self.config, self.shutting_down);
        if !ready.ready_by_leader.is_empty() || !ready.unknown_leaders.is_empty() {
            return Duration::ZERO;
        }

        wake.min(ready.next_ready_check_delay)
    }
}
