batch-aint-one 0.15.1

use std::{collections::VecDeque, fmt::Debug, time::Duration};

use tokio::sync::mpsc;

use crate::{
    BatchError, Limits,
    batch::{Batch, BatchItem},
    batch_inner::Generation,
    processor::Processor,
    worker::Message,
};

/// A double-ended queue for queueing up multiple batches for later processing.
pub(crate) struct BatchQueue<P: Processor> {
    batcher_name: String,

    queue: VecDeque<Batch<P>>,

    limits: Limits,

    /// The number of batches with this key that are currently pre-acquiring resources.
    pre_acquiring: usize,

    /// The number of batches with this key that are currently processing.
    processing: usize,
}

impl<P: Processor> BatchQueue<P> {
    pub(crate) fn new(batcher_name: String, key: P::Key, limits: Limits) -> Self {
        let mut queue = VecDeque::with_capacity(limits.max_batch_queue_size);

        let processing = 0;
        let pre_acquiring = 0;
        queue.push_back(Batch::new(batcher_name.clone(), key));

        Self {
            batcher_name,
            queue,
            limits,
            pre_acquiring,
            processing,
        }
    }

    /// Is the next batch – the one at the front of the queue – full?
    pub(crate) fn is_next_batch_full(&self) -> bool {
        let next = self.queue.front().expect("Should always be non-empty");
        next.is_full(self.limits.max_batch_size)
    }

    pub(crate) fn has_last_batch_reached_size(&self, size: usize) -> bool {
        let last = self.queue.back().expect("Should always be non-empty");
        last.len() >= size
    }

    pub(crate) fn is_last_batch_acquiring_resources(&self) -> bool {
        let last = self.queue.back().expect("Should always be non-empty");
        last.has_started_acquiring()
    }

    pub(crate) fn has_next_batch_timeout_expired(&self) -> bool {
        let next = self.queue.front().expect("Should always be non-empty");
        next.has_timeout_expired()
    }

    /// Is this batch queue full?
    pub(crate) fn is_full(&self) -> bool {
        let back = self.queue.back().expect("Should always be non-empty");
        self.queue.len() >= self.limits.max_batch_queue_size
            && back.len() >= self.limits.max_batch_size
    }

    pub(crate) fn is_empty(&self) -> bool {
        // We always have at least one (possibly empty) batch in the queue.
        self.queue.len() == 1
            && self
                .queue
                .front()
                .expect("Should always be non-empty")
                .is_empty()
    }

    pub(crate) fn last_space_in_batch(&self) -> bool {
        let back = self.queue.back().expect("Should always be non-empty");
        back.has_single_space(self.limits.max_batch_size)
    }

    pub(crate) fn adding_to_new_batch(&self) -> bool {
        let back = self.queue.back().expect("Should always be non-empty");
        back.is_new_batch() || back.is_full(self.limits.max_batch_size)
    }

    /// Are we currently processing any batches for this key?
    pub(crate) fn is_processing(&self) -> bool {
        self.processing > 0
    }

    /// Is this key idle, i.e. no queued items and no batches in flight?
    pub(crate) fn is_idle(&self) -> bool {
        self.is_empty() && self.processing == 0 && self.pre_acquiring == 0
    }

    pub(crate) fn processing(&self) -> usize {
        self.processing
    }

    pub(crate) fn queued(&self) -> usize {
        self.queue.len()
    }

    pub(crate) fn queued_items(&self) -> usize {
        self.queue.iter().map(Batch::len).sum()
    }

    pub(crate) fn mark_processed(&mut self) {
        soft_assert!(
            self.processing > 0,
            "processing count should never go below zero"
        );
        self.processing = self.processing.saturating_sub(1);
    }

    fn increment_processing_count(&mut self) {
        self.processing += 1;

        soft_assert!(
            self.processing <= self.limits.max_key_concurrency,
            "Processing count should not exceed max key concurrency"
        );
    }

    /// Store pre-acquired resources on the batch for the given generation, making it ready for
    /// processing.
    pub(crate) fn resources_acquired(
        &mut self,
        generation: Generation,
        resources: P::Resources,
        span: tracing::Span,
    ) {
        self.mark_resource_acquisition_finished();

        let Some(batch) = self
            .queue
            .iter_mut()
            .find(|batch| batch.is_generation(generation))
        else {
            soft_assert!(
                false,
                "No batch found for generation {:?} in batch queue '{}'",
                generation,
                self.batcher_name
            );
            return;
        };

        batch.resources_acquired(resources, span);
    }

    pub(crate) fn mark_resource_acquisition_finished(&mut self) {
        soft_assert!(
            self.pre_acquiring > 0,
            "pre-acquiring count should never go below zero"
        );
        self.pre_acquiring = self.pre_acquiring.saturating_sub(1);
    }

    fn increment_resource_acquisition_count(&mut self) {
        self.pre_acquiring += 1;

        soft_assert!(
            self.pre_acquiring <= self.limits.max_key_concurrency,
            "pre-acquiring count should not exceed max key concurrency"
        );
    }

    /// Are we currently at maximum total capacity for this key?
    ///
    /// Includes both processing and pre-acquiring batches.
    pub(crate) fn at_max_total_processing_capacity(&self) -> bool {
        self.pre_acquiring + self.processing >= self.limits.max_key_concurrency
    }

    pub(crate) fn push(&mut self, item: BatchItem<P>) {
        let back = self.queue.back_mut().expect("Should always be non-empty");

        if back.is_full(self.limits.max_batch_size) {
            let mut new_back = back.new_generation();
            new_back.push(item);
            self.queue.push_back(new_back);
        } else {
            back.push(item);
        }
    }

    pub(crate) fn has_batch_ready(&self) -> bool {
        for batch in &self.queue {
            if batch.is_ready() {
                return true;
            }
        }

        false
    }

    pub(crate) fn is_generation_ready(&self, generation: Generation) -> bool {
        for batch in &self.queue {
            if batch.is_generation(generation) {
                return batch.is_ready();
            }
        }

        false
    }

    fn take_next_batch(&mut self) -> Option<Batch<P>> {
        let batch = self.queue.pop_front().expect("Should always be non-empty");

        if self.queue.is_empty() {
            self.queue.push_back(batch.new_generation())
        }

        Some(batch)
    }

    pub(crate) fn take_next_ready_batch(&mut self) -> Option<Batch<P>> {
        self.take_first(|batch| batch.is_ready())
    }

    /// Take the first batch matching the predicate, ensuring the queue remains non-empty.
    fn take_first(&mut self, predicate: impl Fn(&Batch<P>) -> bool) -> Option<Batch<P>> {
        for (index, batch) in self.queue.iter().enumerate() {
            if predicate(batch) {
                let batch = self
                    .queue
                    .remove(index)
                    .expect("Should exist, we just found it");

                if self.queue.is_empty() {
                    self.queue.push_back(batch.new_generation())
                }

                return Some(batch);
            }
        }

        None
    }

    pub(crate) fn process_next_ready_batch(
        &mut self,
        processor: P,
        on_finished: mpsc::Sender<Message<P>>,
    ) {
        let Some(batch) = self.take_next_ready_batch() else {
            soft_assert!(
                false,
                "No ready batch found in batch queue '{}'",
                self.batcher_name
            );
            return;
        };

        self.increment_processing_count();

        batch.process(processor, on_finished);
    }

    pub(crate) fn process_next_batch(
        &mut self,
        processor: P,
        on_finished: mpsc::Sender<Message<P>>,
    ) {
        let Some(batch) = self.take_next_batch() else {
            soft_assert!(
                false,
                "No next batch found in batch queue '{}'",
                self.batcher_name
            );
            return;
        };

        self.increment_processing_count();

        batch.process(processor, on_finished);
    }

    fn take_generation(&mut self, generation: Generation) -> Option<Batch<P>> {
        self.take_first(|batch| batch.is_generation(generation))
    }

    pub(crate) fn process_generation(
        &mut self,
        generation: Generation,
        processor: P,
        tx: mpsc::Sender<Message<P>>,
    ) {
        let Some(batch) = self.take_generation(generation) else {
            soft_assert!(
                false,
                "No batch found for generation {:?} in batch queue '{}'",
                generation,
                self.batcher_name
            );
            return;
        };

        self.increment_processing_count();

        batch.process(processor, tx);
    }

    pub(crate) fn fail_generation(&mut self, generation: Generation, error: BatchError<P::Error>) {
        self.mark_resource_acquisition_finished();

        let Some(batch) = self.take_generation(generation) else {
            soft_assert!(
                false,
                "No batch found for generation {:?} in batch queue '{}'",
                generation,
                self.batcher_name
            );
            return;
        };

        batch.fail(error);
    }

    /// Acquire resources for the first batch that hasn't yet acquired resources.
    pub(crate) fn pre_acquire_resources(&mut self, processor: P, tx: mpsc::Sender<Message<P>>) {
        let Some(batch) = self
            .queue
            .iter_mut()
            .find(|batch| !batch.has_started_acquiring())
        else {
            soft_assert!(
                false,
                "No batch found needing resource acquisition in batch queue '{}'",
                self.batcher_name
            );
            return;
        };

        batch.pre_acquire_resources(processor, tx);

        self.increment_resource_acquisition_count();
    }

    /// Process the last batch after a delay.
    pub(crate) fn process_after(&mut self, duration: Duration, tx: mpsc::Sender<Message<P>>) {
        let back = self.queue.back_mut().expect("Should always be non-empty");
        back.process_after(duration, tx);
    }
}

impl<P: Processor> Debug for BatchQueue<P> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self {
            batcher_name,
            queue,
            limits,
            processing,
            pre_acquiring,
        } = self;
        f.debug_struct("BatchQueue")
            .field("batcher_name", &batcher_name)
            .field("queue", &queue)
            .field("processing", &processing)
            .field("pre_acquiring", &pre_acquiring)
            .field("limits", limits)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use tokio::sync::oneshot;
    use tracing::Span;

    use super::*;

    #[derive(Clone)]
    struct DummyProcessor;
    impl Processor for DummyProcessor {
        type Key = String;
        type Input = String;
        type Output = String;
        type Error = String;
        type Resources = ();
        async fn acquire_resources(&self, _key: String) -> Result<(), String> {
            Ok(())
        }
        async fn process(
            &self,
            _key: String,
            _inputs: impl Iterator<Item = String> + Send,
            _resources: (),
        ) -> Result<Vec<String>, String> {
            Ok(vec![])
        }
    }

    fn item(key: &str) -> BatchItem<DummyProcessor> {
        let (tx, _rx) = oneshot::channel();
        BatchItem {
            key: key.to_string(),
            input: "item".to_string(),
            submitted_at: tokio::time::Instant::now(),
            received_at: None,
            tx,
            requesting_span: Span::none(),
        }
    }

    #[tokio::test]
    async fn queued_items_sums_across_batches() {
        let limits = Limits::builder().max_batch_size(2).build();
        let mut queue: BatchQueue<DummyProcessor> =
            BatchQueue::new("test".to_string(), "key".to_string(), limits);

        for _ in 0..3 {
            queue.push(item("key"));
        }

        // The 3rd item can't fit in the first (now full) batch, so it starts a new one.
        assert_eq!(queue.queued(), 2, "should have split into 2 batches");
        assert_eq!(
            queue.queued_items(),
            3,
            "should count every item, not just batches"
        );
    }
}