//! Batching tasks.
//!
//! [`Batch`] represents a prepared batch of tasks.  It offers functionality for
//! preparing new batches and extracting tasks from them, but does not consider
//! the collection of tasks from external sources.

use alloc::vec::Vec;
use core::{num::NonZeroUsize, ops::Range};

use crate::{Queue, Task, TaskPriority, config::Config, pubqueue::PubQueue};

//----------- Batch ------------------------------------------------------------

/// A worker's batch of tasks.
pub struct Batch<T: Task> {
    /// The batch counter.
    ///
    /// This counts down the number of tasks in the batch and decides where the
    /// next task will be retrieved from (the local or the public queue).
    pub counter: BatchCounter,

    /// The local queue of tasks.
    pub local: Vec<T>,

    /// The public queue of tasks.
    pub pub_queue: PubQueue,

    /// The priority of the public queue.
    ///
    /// This identifies the priority of the tasks in the public queue, if it is
    /// non-empty.  Otherwise, it is [`None`].
    pub pq_priority: Option<T::Priority>,
}

impl<T: Task> Batch<T> {
    /// Construct a new [`Batch`].
    ///
    /// ## Safety
    ///
    /// - `id` is the worker's ID.
    /// - The worker's public queue has not been modified since initialization.
    pub unsafe fn new(id: usize, config: &Config) -> Self {
        debug_assert!(id < config.num_workers.get());

        Self {
            counter: BatchCounter::empty(),
            local: Vec::with_capacity(config.batch_size.get()),
            // SAFETY:
            // - 'id' is the worker's ID, and is thus a valid ID.
            // - 'config.batch_size <= config.batch_size'.
            pub_queue: unsafe { PubQueue::new(id, config.batch_size, config) },
            pq_priority: None,
        }
    }

    /// Fill the batch.
    ///
    /// ## Safety
    ///
    /// - `id` is the worker's ID.
    /// - `global` is the worker's global queue.
    #[inline]
    pub unsafe fn fill(
        &mut self,
        reserved: usize,
        tasks: &mut Vec<T>,
        id: usize,
        global: &Queue<T>,
    ) {
        // Initialize the batch.
        self.counter = BatchCounter::new(tasks.len(), reserved, &global.config);

        // Determine which tasks are part of the batch.
        let batch_offset = tasks.len() - self.counter.total();
        let priority = tasks.get(batch_offset + 1).map(|t| t.priority());
        let mut tasks = tasks.drain(batch_offset..);

        debug_assert!(self.pq_priority.is_none());
        let pq = self.pub_queue.id(&global.config);
        let pq = &global.pq_contents[pq];

        if <T::Priority as TaskPriority>::TRIVIAL {
            // Put half the tasks in the public queue.
            // SAFETY:
            // - 'tasks' has 'batch.total() <= pq_size * 2' elements, of which
            //   half are filtered out and sent to the local queue.  Thus, the
            //   iterator will produce at most 'pq_size' elements.
            // - 'pq' is owned by this worker, and it has not published a
            //   priority.
            unsafe { pq.fill(tasks.by_ref().take(self.counter.public_num())) };

            // Put half the tasks in the local queue.
            self.local.extend(tasks);
        } else {
            // Divide the tasks between the local and public queues.
            let mut local = false;
            let tasks = tasks.filter_map(|task| {
                local = !local;
                if local {
                    self.local.push(task);
                    None
                } else {
                    Some(task)
                }
            });

            // Fill the local and public queues.
            // SAFETY:
            // - 'tasks' has 'batch.total() <= pq_size * 2' elements, of which
            //   half are filtered out and sent to the local queue.  Thus, the
            //   iterator will produce at most 'pq_size' elements.
            // - 'pq' is owned by this worker, and it has not published a
            //   priority.
            unsafe { pq.fill(tasks) };

            debug_assert_eq!(self.local.len(), self.counter.local_num());
        }

        // If the public queue is non-empty, publish it.
        if let Some(priority) = priority {
            let pq_len = self.counter.public_num();
            // SAFETY: As per the batch, the public queue is non-empty.
            let pq_len = unsafe { NonZeroUsize::new_unchecked(pq_len) };

            let stealer = &global.stealer[id];
            // SAFETY:
            // - 'pq_len <= pq_size' as per 'Batch::public_num()'.
            let new_pubqueue =
                unsafe { self.pub_queue.with_len(pq_len, &global.config) };
            // SAFETY:
            // - This worker owns 'stealer'.
            // - The published priority is 'None'.
            // - 'last_pubqueue' is the last known value of 'stealer'.
            unsafe { stealer.set(self.pub_queue, new_pubqueue) };
            self.pub_queue = new_pubqueue;

            // SAFETY:
            // - This worker owns this priority.
            // - The corresponding stealer has been initialized.
            unsafe { global.priority[id].set(priority) };
            self.pq_priority = Some(priority);

            // Wake up a sleeping thread, if any.
            global.wake();
        }
    }

    /// Retrieve the next task in the batch.
    ///
    /// As per the batch counter, a task will be retrieved from the local or
    /// public queue.  If the batch is depleted, [`None`] is returned instead.
    ///
    /// ## Safety
    ///
    /// - `id` is the worker's ID.
    /// - `global` is the worker's global queue.
    #[inline]
    pub unsafe fn next(&mut self, id: usize, global: &Queue<T>) -> Option<T> {
        if <T::Priority as TaskPriority>::TRIVIAL {
            // NOTE: We don't keep the batch counter updated here.

            // All tasks have equal priority, so it doesn't matter whether we
            // read from the local queue or the public queue.  The local queue
            // is easier and more efficient to read from.
            return self.local.pop();
        }

        debug_assert_eq!(self.counter.local_num(), self.local.len());
        if self.counter.public_num() != 0 {
            debug_assert_eq!(
                self.counter.public_num(),
                self.pub_queue.len(&global.config).get()
            );
            debug_assert!(self.pq_priority.is_some());
        }

        // Tick off the counter and pick the queue to read from.
        if self.counter.take()? {
            // Read the task from the local queue.

            // SAFETY:
            // - Before 'take()', 'counter.local_num() == local.len()'.
            // - After 'take()', 'local_num()' decreased by 1.
            // - Thus 'local_num()' was non-zero.
            // - Thus 'local.len()' is still non-zero.
            // - Thus 'local' is non-empty.
            Some(unsafe { self.local.pop().unwrap_unchecked() })
        } else {
            // Read the task from the public queue.

            let index = self.counter.public_num();
            let pq = self.pub_queue.id(&global.config);
            let len = self.pub_queue.len(&global.config);
            let pq = &global.pq_contents[pq];

            // Pre-emptively read this task from the public queue.
            debug_assert_eq!(index + 1, len.get());
            // SAFETY: 'index < len <= pq_size'
            let task = unsafe { pq.read(index) };

            // Take ownership of this task.
            if index > 0 {
                // Try updating the public queue length.
                let stealer = &global.stealer[id];
                // SAFETY:
                // - 'stealer' belongs to this worker.
                // - 'last_pubqueue' is the last known stealer value.
                // - 'index = len - 1 > 0' so 'len > 0'.
                match unsafe { stealer.try_dec(self.pub_queue, &global.config) }
                {
                    Ok(pq) => {
                        // The public queue was updated successfully.
                        self.pub_queue = pq;
                    }

                    Err(pq) => {
                        // The public queue has been stolen.
                        self.pq_priority = None;
                        self.pub_queue = pq;
                        return None;
                    }
                }
            } else {
                // Taking this task would empty the public queue, which
                // means ownership of the whole public queue is needed.

                // SAFETY: As per the counter, the public queue contains tasks
                // and has not yet been stolen.
                let last_priority =
                    unsafe { self.pq_priority.take().unwrap_unchecked() };

                let priority = &global.priority[id];
                match unsafe { priority.steal_self(last_priority) } {
                    Ok(()) => {
                        // Ownership of the public queue has been taken.
                    }

                    Err(()) => {
                        // Another worker is in the middle of stealing from
                        // us.
                        let stealer = &global.stealer[id];

                        // Wait until the public queue is stolen.
                        // SAFETY:
                        // - As per 'Priority::steal_self()', the priority
                        //   of this worker was set to 'None' by another
                        //   worker.
                        // - 'last_pubqueue' is the last value for 'stealer'
                        //   seen by this worker.
                        self.pub_queue =
                            unsafe { stealer.wait_for_theft(self.pub_queue) };

                        return None;
                    }
                }
            }

            // SAFETY: Taking ownership of the task was successful.  This
            // means that nobody could have overwritten the task since we
            // had filled the public queue, so our read of the task was
            // coherent and read initialized data.
            Some(unsafe { task.assume_init() })
        }
    }

    /// Drain the batch.
    ///
    /// ## Safety
    ///
    /// - `id` is the worker's ID.
    /// - `global` is the worker's global queue.
    #[inline]
    pub unsafe fn drain(
        &mut self,
        tasks: &mut Vec<T>,
        id: usize,
        global: &Queue<T>,
    ) {
        // Move out of the local queue.
        tasks.append(&mut self.local);

        // Collect tasks from the public queue.
        if let Some(priority) = self.pq_priority.take() {
            // SAFETY:
            // - This 'Priority' belongs to the current worker.
            // - 'last_priority' was 'Some(priority)'.
            match unsafe { global.priority[id].steal_self(priority) } {
                Ok(()) => {
                    // We have taken ownership of the public queue.
                    let pq = self.pub_queue.id(&global.config);
                    let len = self.pub_queue.len(&global.config);
                    let pq = &global.pq_contents[pq];

                    // Move the tasks out of the public queue.
                    // SAFETY:
                    // - 'len <= pq_size' as per 'PubQueue::len()'.
                    // - We set the length of the public queue to 'len', and
                    //   nobody stole it, so the length is unchanged.
                    // - This worker owns 'pq' so nobody will write to it.
                    unsafe { pq.move_to(len.get(), tasks) };
                }

                Err(()) => {
                    // Another worker is in the middle of stealing from us.
                    let stealer = &global.stealer[id];

                    // Wait until the public queue is stolen.
                    // SAFETY:
                    // - As per 'Priority::steal_self()', the priority of this
                    //   worker was set to 'None' by another worker.
                    // - 'last_pubqueue' is the last value for 'stealer' seen by
                    //   this worker.
                    self.pub_queue =
                        unsafe { stealer.wait_for_theft(self.pub_queue) };
                }
            }
        }

        // Reset the batch counter.
        self.counter = BatchCounter::empty();
    }
}

//----------- BatchCounter -----------------------------------------------------

/// A counter for a [`Batch`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BatchCounter {
    /// The range of pending task numbers.
    ///
    /// `range.end - range.start` is the number of pending tasks.
    ///
    /// Invariants:
    /// - `range.start <= range.end`.
    /// - `range.end <= pq_size * 2`.
    range: Range<usize>,
}

impl BatchCounter {
    /// Construct an empty [`BatchCounter`].
    pub const fn empty() -> Self {
        Self { range: 0..0 }
    }

    /// Construct a new [`BatchCounter`].
    ///
    /// `remaining` is the number of tasks to be divided between the local and
    /// public queues.  If some tasks have already been removed (for special
    /// handling) but should be counted against the maximum batch size, their
    /// count should be passed in `reserved`.  The batch will contain at most
    /// `config.batch_size - reserved` tasks.
    pub const fn new(
        remaining: usize,
        reserved: usize,
        config: &Config,
    ) -> Self {
        let batch_size = config.batch_size.get();
        debug_assert!(reserved <= batch_size * 2);

        // Determine the total number of tasks for the batch.
        //let num = remaining.max(pq_size * 2 - reserved);
        let total = if remaining < batch_size * 2 - reserved {
            remaining
        } else {
            batch_size - reserved
        };

        // Separate the tasks into pending and not-pending sets.
        let not_pending = total / 2;

        Self {
            range: not_pending..total,
        }
    }

    /// Whether the batch is empty.
    pub const fn is_empty(&self) -> bool {
        self.range.start >= self.range.end
    }

    /// The total number of tasks in the batch.
    pub const fn total(&self) -> usize {
        self.range.end
    }

    /// The number of batch tasks in the local queue.
    ///
    /// There may be additional tasks in the local queue before the batch.
    pub const fn local_num(&self) -> usize {
        // 0, 1, 2, 3, 4, 5, 6, 7, ... -> 0, 1, 1, 2, 2, 3, 3, 4, ...
        self.range.end.div_ceil(2)
    }

    /// The number of batch tasks in the public queue.
    ///
    /// There are no other tasks in the public queue.
    pub const fn public_num(&self) -> usize {
        // 0, 1, 2, 3, 4, 5, 6, 7, ... -> 0, 0, 1, 1, 2, 2, 3, 3, ...
        self.range.end / 2
    }

    /// Extract a pending task.
    ///
    /// If the task is part of the local queue, `true` is returned.
    pub const fn take(&mut self) -> Option<bool> {
        if self.range.start >= self.range.end {
            return None;
        }

        // SAFETY: '0 <= range.start < range.end' so 'range.end > 0'.
        self.range.end -= 1;

        // 0, 1, 2, 3, 4, 5, 6, 7, ... -> L, P, L, P, L, P, L, P, ...
        Some(self.range.end.is_multiple_of(2))
    }
}
