//! Multi-producer, single-consumer channels.

use std::collections::VecDeque;
use std::future::poll_fn;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

use crate::op::completion::{CompletionFuture, CompletionHandle};
use crate::sys::current::channel::runtime_waiter;

/// Creates a bounded channel with room for at most `capacity` queued messages.
///
/// Bounded senders provide both [`Sender::try_send`] and async [`Sender::send`] backpressure.
///
/// # Panics
///
/// Panics if `capacity == 0`.
pub fn channel<T: Send + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
    assert!(capacity > 0, "bounded channels require capacity > 0");
    let shared = Arc::new(Mutex::new(State::new(Some(capacity))));
    (
        Sender {
            shared: Arc::clone(&shared),
        },
        Receiver { shared },
    )
}

/// Creates an unbounded channel.
///
/// Unbounded senders never wait for capacity, but the single receiver is still asynchronous.
pub fn unbounded_channel<T: Send + 'static>() -> (UnboundedSender<T>, Receiver<T>) {
    let shared = Arc::new(Mutex::new(State::new(None)));
    (
        UnboundedSender {
            shared: Arc::clone(&shared),
        },
        Receiver { shared },
    )
}

/// Bounded multi-producer sender.
pub struct Sender<T: Send + 'static> {
    shared: Arc<Mutex<State<T>>>,
}

/// Unbounded multi-producer sender.
pub struct UnboundedSender<T: Send + 'static> {
    shared: Arc<Mutex<State<T>>>,
}

/// Single consumer for both bounded and unbounded MPSC channels.
pub struct Receiver<T: Send + 'static> {
    shared: Arc<Mutex<State<T>>>,
}

struct State<T: Send + 'static> {
    queue: VecDeque<T>,
    capacity: Option<usize>,
    sender_count: usize,
    receiver_closed: bool,
    recv_waiter: Option<CompletionHandle<Option<T>>>,
    send_waiters: VecDeque<SendWaiter<T>>,
    next_waiter_id: usize,
}

struct SendWaiter<T: Send + 'static> {
    id: usize,
    value: T,
    handle: CompletionHandle<Result<(), SendError<T>>>,
}

#[derive(Debug, Eq, PartialEq)]
/// Error returned when sending fails because the receiver has been closed or dropped.
pub struct SendError<T>(pub T);

#[derive(Debug, Eq, PartialEq)]
/// Error returned by [`Sender::try_send`] when a message cannot be queued immediately.
pub enum TrySendError<T> {
    /// The bounded queue is currently full.
    Full(T),
    /// The receiver has been closed or dropped.
    Closed(T),
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
/// Error returned by [`Receiver::try_recv`] when no message is available immediately.
pub enum TryRecvError {
    /// The channel is still open, but currently empty.
    Empty,
    /// The channel is closed and no more messages can arrive.
    Disconnected,
}

/// A wakeup deferred until the channel mutex has been released.
///
/// Waking a waiter while holding the channel lock can be expensive (cross-thread
/// wakeups go through the io_uring ring notification path) and risks priority
/// inversion.  All `State` methods collect these instead of calling
/// `CompletionHandle::complete` directly; the caller fires them after dropping
/// the `MutexGuard`.
enum PendingCompletion<T: Send + 'static> {
    RecvSome(CompletionHandle<Option<T>>, T),
    RecvNone(CompletionHandle<Option<T>>),
    SendOk(CompletionHandle<Result<(), SendError<T>>>),
    SendErr(CompletionHandle<Result<(), SendError<T>>>, T),
}

fn fire_completions<T: Send + 'static>(completions: Vec<PendingCompletion<T>>) {
    for c in completions {
        match c {
            PendingCompletion::RecvSome(h, v) => h.complete(Some(v)),
            PendingCompletion::RecvNone(h) => h.complete(None),
            PendingCompletion::SendOk(h) => h.complete(Ok(())),
            PendingCompletion::SendErr(h, v) => h.complete(Err(SendError(v))),
        }
    }
}

impl<T: Send + 'static> State<T> {
    fn new(capacity: Option<usize>) -> Self {
        Self {
            queue: VecDeque::new(),
            capacity,
            sender_count: 1,
            receiver_closed: false,
            recv_waiter: None,
            send_waiters: VecDeque::new(),
            next_waiter_id: 1,
        }
    }

    fn try_send_now(
        &mut self,
        value: T,
        completions: &mut Vec<PendingCompletion<T>>,
    ) -> Result<(), TrySendError<T>> {
        if self.receiver_closed {
            return Err(TrySendError::Closed(value));
        }

        if let Some(waiter) = self.recv_waiter.take() {
            completions.push(PendingCompletion::RecvSome(waiter, value));
            return Ok(());
        }

        if self
            .capacity
            .is_some_and(|capacity| self.queue.len() >= capacity)
        {
            return Err(TrySendError::Full(value));
        }

        self.queue.push_back(value);
        Ok(())
    }

    fn enqueue_send_waiter(
        &mut self,
        value: T,
        handle: CompletionHandle<Result<(), SendError<T>>>,
    ) -> usize {
        let id = self.next_waiter_id;
        self.next_waiter_id = self.next_waiter_id.wrapping_add(1);
        self.send_waiters
            .push_back(SendWaiter { id, value, handle });
        id
    }

    fn remove_send_waiter(&mut self, waiter_id: usize) -> bool {
        let Some(index) = self
            .send_waiters
            .iter()
            .position(|waiter| waiter.id == waiter_id)
        else {
            return false;
        };
        self.send_waiters.remove(index);
        true
    }

    fn pump_senders(&mut self, completions: &mut Vec<PendingCompletion<T>>) {
        loop {
            if self.receiver_closed {
                self.fail_pending_senders(completions);
                break;
            }

            let has_capacity = self
                .capacity
                .is_none_or(|capacity| self.queue.len() < capacity);
            if !has_capacity {
                break;
            }

            let Some(waiter) = self.send_waiters.pop_front() else {
                break;
            };

            if let Some(receiver) = self.recv_waiter.take() {
                completions.push(PendingCompletion::RecvSome(receiver, waiter.value));
            } else {
                self.queue.push_back(waiter.value);
            }
            completions.push(PendingCompletion::SendOk(waiter.handle));
        }

        if self.queue.is_empty()
            && self.sender_count == 0
            && let Some(waiter) = self.recv_waiter.take()
        {
            completions.push(PendingCompletion::RecvNone(waiter));
        }
    }

    fn fail_pending_senders(&mut self, completions: &mut Vec<PendingCompletion<T>>) {
        while let Some(waiter) = self.send_waiters.pop_front() {
            completions.push(PendingCompletion::SendErr(waiter.handle, waiter.value));
        }
    }

    fn close_receiver(&mut self, completions: &mut Vec<PendingCompletion<T>>) {
        self.receiver_closed = true;
        self.fail_pending_senders(completions);
        if self.queue.is_empty()
            && let Some(waiter) = self.recv_waiter.take()
        {
            completions.push(PendingCompletion::RecvNone(waiter));
        }
    }

    fn drop_sender(&mut self, completions: &mut Vec<PendingCompletion<T>>) {
        self.sender_count = self
            .sender_count
            .checked_sub(1)
            .expect("sender count underflow: more drops than creates");
        if self.sender_count == 0
            && self.queue.is_empty()
            && let Some(waiter) = self.recv_waiter.take()
        {
            completions.push(PendingCompletion::RecvNone(waiter));
        }
    }
}

impl<T: Send + 'static> Clone for Sender<T> {
    fn clone(&self) -> Self {
        self.shared
            .lock()
            .expect("mpsc state should not be poisoned")
            .sender_count += 1;
        Self {
            shared: Arc::clone(&self.shared),
        }
    }
}

impl<T: Send + 'static> Clone for UnboundedSender<T> {
    fn clone(&self) -> Self {
        self.shared
            .lock()
            .expect("mpsc state should not be poisoned")
            .sender_count += 1;
        Self {
            shared: Arc::clone(&self.shared),
        }
    }
}

impl<T: Send + 'static> Sender<T> {
    /// Waits until the message can be queued.
    ///
    /// When the bounded channel is full, this future waits until the receiver frees capacity.
    ///
    /// # Panics
    ///
    /// Panics if this future is first polled outside a runtime-managed thread.
    pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
        let mut value = Some(value);
        let mut wait = None;
        poll_fn(|cx| self.poll_send(cx, &mut value, &mut wait)).await
    }

    /// Attempts to queue a message immediately.
    pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
        let mut completions = Vec::new();
        let result = {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.try_send_now(value, &mut completions)
        };
        fire_completions(completions);
        result
    }

    /// Returns `true` if the receiver has been closed or dropped.
    pub fn is_closed(&self) -> bool {
        self.shared
            .lock()
            .expect("mpsc state should not be poisoned")
            .receiver_closed
    }

    fn poll_send(
        &self,
        cx: &mut Context<'_>,
        value_slot: &mut Option<T>,
        wait: &mut Option<CompletionFuture<Result<(), SendError<T>>>>,
    ) -> Poll<Result<(), SendError<T>>> {
        if let Some(future) = wait.as_mut() {
            match Pin::new(future).poll(cx) {
                Poll::Ready(result) => {
                    wait.take();
                    Poll::Ready(result)
                }
                Poll::Pending => Poll::Pending,
            }
        } else {
            let mut completions = Vec::new();
            let first_result = {
                let mut state = self
                    .shared
                    .lock()
                    .expect("mpsc state should not be poisoned");
                state.try_send_now(
                    value_slot.take().expect("send value should be present"),
                    &mut completions,
                )
            };
            fire_completions(completions);
            match first_result {
                Ok(()) => Poll::Ready(Ok(())),
                Err(TrySendError::Closed(value)) => Poll::Ready(Err(SendError(value))),
                Err(TrySendError::Full(returned)) => {
                    let (future, handle) = runtime_waiter::<Result<(), SendError<T>>>();
                    let state_shared = Arc::clone(&self.shared);
                    let mut completions = Vec::new();
                    let registration = {
                        let mut state = state_shared
                            .lock()
                            .expect("mpsc state should not be poisoned");
                        match state.try_send_now(returned, &mut completions) {
                            Ok(()) => Ok(None),
                            Err(TrySendError::Closed(value)) => Err(SendError(value)),
                            Err(TrySendError::Full(value)) => {
                                Ok(Some(state.enqueue_send_waiter(value, handle.clone())))
                            }
                        }
                    };
                    fire_completions(completions);
                    match registration {
                        Ok(None) => {
                            handle.complete(Ok(()));
                            *wait = Some(future);
                            self.poll_send(cx, value_slot, wait)
                        }
                        Err(error) => {
                            handle.complete(Err(error));
                            *wait = Some(future);
                            self.poll_send(cx, value_slot, wait)
                        }
                        Ok(Some(waiter_id)) => {
                            let cancel_shared = Arc::clone(&self.shared);
                            let cancel_handle = handle.clone();
                            handle.set_cancel(move || {
                                let mut state = cancel_shared
                                    .lock()
                                    .expect("mpsc state should not be poisoned");
                                let _ = state.remove_send_waiter(waiter_id);
                                drop(state);
                                cancel_handle.finish(None);
                            });
                            *wait = Some(future);
                            self.poll_send(cx, value_slot, wait)
                        }
                    }
                }
            }
        }
    }
}

impl<T: Send + 'static> UnboundedSender<T> {
    /// Queues a message immediately.
    pub fn send(&self, value: T) -> Result<(), SendError<T>> {
        let mut completions = Vec::new();
        let result = {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.try_send_now(value, &mut completions)
        };
        fire_completions(completions);
        result.map_err(|error| match error {
            TrySendError::Full(value) | TrySendError::Closed(value) => SendError(value),
        })
    }

    /// Returns `true` if the receiver has been closed or dropped.
    pub fn is_closed(&self) -> bool {
        self.shared
            .lock()
            .expect("mpsc state should not be poisoned")
            .receiver_closed
    }
}

impl<T: Send + 'static> Receiver<T> {
    /// Waits for the next message.
    ///
    /// Returns `None` when the channel is closed and all buffered messages have been drained.
    ///
    /// # Panics
    ///
    /// Panics if this future is first polled outside a runtime-managed thread.
    pub async fn recv(&mut self) -> Option<T> {
        let mut wait = None;
        poll_fn(|cx| self.poll_recv(cx, &mut wait)).await
    }

    /// Attempts to receive a message immediately.
    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
        let mut completions = Vec::new();
        let result = {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            if let Some(value) = state.queue.pop_front() {
                state.pump_senders(&mut completions);
                Ok(value)
            } else if state.sender_count == 0 || state.receiver_closed {
                Err(TryRecvError::Disconnected)
            } else {
                Err(TryRecvError::Empty)
            }
        };
        fire_completions(completions);
        result
    }

    /// Closes the channel to future sends.
    ///
    /// Already-buffered messages remain available to [`recv`](Self::recv) and
    /// [`try_recv`](Self::try_recv).
    pub fn close(&mut self) {
        let mut completions = Vec::new();
        {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.close_receiver(&mut completions);
        }
        fire_completions(completions);
    }

    /// Returns `true` if the channel is closed or all senders have been dropped.
    pub fn is_closed(&self) -> bool {
        let state = self
            .shared
            .lock()
            .expect("mpsc state should not be poisoned");
        state.receiver_closed || state.sender_count == 0
    }

    fn poll_recv(
        &mut self,
        cx: &mut Context<'_>,
        wait: &mut Option<CompletionFuture<Option<T>>>,
    ) -> Poll<Option<T>> {
        if let Some(future) = wait.as_mut() {
            match Pin::new(future).poll(cx) {
                Poll::Ready(result) => {
                    wait.take();
                    Poll::Ready(result)
                }
                Poll::Pending => Poll::Pending,
            }
        } else {
            let (future, handle) = runtime_waiter::<Option<T>>();
            let cancel_shared = Arc::clone(&self.shared);
            let cancel_handle = handle.clone();
            handle.set_cancel(move || {
                let mut state = cancel_shared
                    .lock()
                    .expect("mpsc state should not be poisoned");
                let _ = state.recv_waiter.take();
                drop(state);
                cancel_handle.finish(None);
            });

            let mut completions = Vec::new();
            {
                let mut state = self
                    .shared
                    .lock()
                    .expect("mpsc state should not be poisoned");
                if let Some(value) = state.queue.pop_front() {
                    state.pump_senders(&mut completions);
                    completions.push(PendingCompletion::RecvSome(handle.clone(), value));
                } else if state.receiver_closed || state.sender_count == 0 {
                    completions.push(PendingCompletion::RecvNone(handle.clone()));
                } else {
                    assert!(
                        state.recv_waiter.is_none(),
                        "only one mpsc receive operation may wait at a time"
                    );
                    state.recv_waiter = Some(handle.clone());
                }
            }
            fire_completions(completions);

            *wait = Some(future);
            self.poll_recv(cx, wait)
        }
    }
}

impl<T: Send + 'static> Drop for Sender<T> {
    fn drop(&mut self) {
        let mut completions = Vec::new();
        {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.drop_sender(&mut completions);
        }
        fire_completions(completions);
    }
}

impl<T: Send + 'static> Drop for UnboundedSender<T> {
    fn drop(&mut self) {
        let mut completions = Vec::new();
        {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.drop_sender(&mut completions);
        }
        fire_completions(completions);
    }
}

impl<T: Send + 'static> Drop for Receiver<T> {
    fn drop(&mut self) {
        let mut completions = Vec::new();
        {
            let mut state = self
                .shared
                .lock()
                .expect("mpsc state should not be poisoned");
            state.close_receiver(&mut completions);
        }
        fire_completions(completions);
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};
    use std::time::Duration;

    use crate::time::sleep;
    use crate::{queue_future, queue_task, run, spawn_worker};

    use super::{TryRecvError, TrySendError, channel, unbounded_channel};

    #[test]
    fn bounded_channel_applies_backpressure() {
        let log = Arc::new(Mutex::new(Vec::<String>::new()));
        let log_for_task = Arc::clone(&log);

        queue_task(move || {
            let (sender, mut receiver) = channel(1);
            let log_for_sender = Arc::clone(&log_for_task);
            let log_for_receiver = Arc::clone(&log_for_task);

            queue_future(async move {
                sender
                    .send("first")
                    .await
                    .expect("first send should succeed");
                log_for_sender
                    .lock()
                    .unwrap()
                    .push("sent first".to_string());
                sender
                    .send("second")
                    .await
                    .expect("second send should succeed");
                log_for_sender
                    .lock()
                    .unwrap()
                    .push("sent second".to_string());
            });

            queue_future(async move {
                sleep(Duration::from_millis(5)).await;
                let first = receiver.recv().await.expect("first recv should succeed");
                log_for_receiver
                    .lock()
                    .unwrap()
                    .push(format!("received {first}"));
                let second = receiver.recv().await.expect("second recv should succeed");
                log_for_receiver
                    .lock()
                    .unwrap()
                    .push(format!("received {second}"));
            });
        });
        run();

        let log = log.lock().unwrap();
        let sent_first = log.iter().position(|entry| entry == "sent first").unwrap();
        let received_first = log
            .iter()
            .position(|entry| entry == "received first")
            .unwrap();
        let sent_second = log.iter().position(|entry| entry == "sent second").unwrap();
        let received_second = log
            .iter()
            .position(|entry| entry == "received second")
            .unwrap();

        assert!(
            sent_first < received_first,
            "first send should happen before first recv"
        );
        assert!(
            received_first < sent_second,
            "second send should not complete before capacity is freed"
        );
        assert!(
            received_first < received_second,
            "receiver should observe messages in FIFO order"
        );
    }

    #[test]
    fn unbounded_channel_moves_messages_across_worker_threads() {
        let log = Arc::new(Mutex::new(Vec::new()));
        let log_for_task = Arc::clone(&log);

        queue_task(move || {
            let (sender, mut receiver) = unbounded_channel::<String>();
            let worker_sender = sender.clone();
            let log_for_receiver = Arc::clone(&log_for_task);

            let _worker = spawn_worker(
                move || {
                    queue_task(move || {
                        worker_sender
                            .send("worker boot".into())
                            .expect("worker boot send should succeed");
                        worker_sender
                            .send("worker done".into())
                            .expect("worker done send should succeed");
                    });
                },
                || {},
            );
            drop(sender);

            queue_future(async move {
                while let Some(message) = receiver.recv().await {
                    log_for_receiver.lock().unwrap().push(message);
                }
            });
        });
        run();

        assert_eq!(
            log.lock().unwrap().as_slice(),
            ["worker boot", "worker done"]
        );
    }

    #[test]
    fn try_send_try_recv_and_close_semantics_work() {
        let (sender, mut receiver) = channel(1);
        sender
            .try_send(1usize)
            .expect("initial send should succeed");
        assert_eq!(sender.try_send(2usize), Err(TrySendError::Full(2)));
        assert_eq!(receiver.try_recv(), Ok(1));
        assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
        receiver.close();
        assert!(sender.is_closed(), "sender should observe closed receiver");
        assert_eq!(sender.try_send(3usize), Err(TrySendError::Closed(3)));
        assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected));
    }
}