//! Instrumented wrappers for `tokio::sync::mpsc` channels.
//!
//! [`Sender`] and [`Receiver`] wrap the tokio bounded mpsc channel
//! and report to a [`ResourceObserver`] before/after every `send()` and `recv()`.
//!
//! For lockdep purposes:
//!
//! - **`send()`** on a bounded channel can block (back-pressure).  The useful
//!   ordering edge is `held_lock → channel`, recorded at wait time — before send
//!   completes.  After the value is enqueued, the sender isn't holding
//!   anything, so the channel resource is immediately released from the
//!   held-lock stack.
//!
//! - **`recv()`** blocks until a message arrives.  Once it completes, the
//!   channel resource
//!   is pushed onto the held-lock stack **and never removed**.  Any resource
//!   acquired later in the same task automatically records
//!   `channel → resource`.
//!   This is correct because the ordering relationship "this task receives
//!   from the channel and also acquires `resource`" is a permanent fact about the
//!   task's behaviour — it doesn't expire when "processing" ends.
use crate::observer::{DefaultObserver, Id, Resource, ResourceObserver, observer};
use std::future::Future;
use std::panic::Location;

/// Create a bounded mpsc channel using the crate-wide default observer.
#[track_caller]
pub fn channel<T>(buffer: usize) -> (Sender<T, DefaultObserver>, Receiver<T, DefaultObserver>) {
    let caller = Location::caller();
    let observer = observer();
    let (tx, rx) = tokio::sync::mpsc::channel(buffer);
    let id = Id::unnamed(caller);
    (
        Sender {
            id,
            inner: tx,
            observer: observer.clone(),
        },
        Receiver {
            id,
            inner: rx,
            observer: observer.clone(),
        },
    )
}

/// Create a named bounded mpsc channel using the crate-wide default observer.
#[track_caller]
pub fn named_channel<T>(
    name: &'static str,
    buffer: usize,
) -> (Sender<T, DefaultObserver>, Receiver<T, DefaultObserver>) {
    let observer = observer();
    let (tx, rx) = tokio::sync::mpsc::channel(buffer);
    let id = Id::named(name);
    (
        Sender {
            id,
            inner: tx,
            observer: observer.clone(),
        },
        Receiver {
            id,
            inner: rx,
            observer,
        },
    )
}

/// Instrumented wrapper around `tokio::sync::mpsc::Sender`.
///
/// On a bounded channel, `send()` may block waiting for capacity — this is
/// reported as a wait-for / acquired on the channel resource.
#[derive(Clone, Debug)]
pub struct Sender<T, O: ResourceObserver = DefaultObserver> {
    id: Id,
    inner: tokio::sync::mpsc::Sender<T>,
    observer: O,
}

impl<T, O: ResourceObserver> Sender<T, O> {
    /// Send a value, blocking if the channel is full.
    ///
    /// For lockdep, the ordering edge `held_lock → channel` is recorded at wait
    /// time (before the send completes).  Once the value is enqueued, nothing
    /// is "held" — the channel resource is immediately released.
    #[track_caller]
    pub fn send(
        &self,
        value: T,
    ) -> impl Future<Output = Result<(), tokio::sync::mpsc::error::SendError<T>>> + '_ {
        let caller = Location::caller();
        async move {
            let id = Resource::Channel(self.id);
            self.observer.on_waiting(&id, caller);
            let result = self.inner.send(value).await;
            match &result {
                Ok(()) => {
                    self.observer.on_acquired(&id, caller);
                    self.observer.on_released(&id, caller);
                }
                Err(_) => {}
            }
            result
        }
    }
    /// Try to send without blocking (not instrumented for lockdep since it never waits).
    pub fn try_send(&self, value: T) -> Result<(), tokio::sync::mpsc::error::TrySendError<T>> {
        self.inner.try_send(value)
    }

    /// Get the resource id
    pub fn id(&self) -> Id {
        self.id
    }
}

/// Instrumented wrapper around `tokio::sync::mpsc::Receiver`.
///
/// `recv()` blocks until a message is available — this is reported as a
/// wait-for / acquired on the channel resource.
///
/// Once `recv()` completes, the channel resource stays on the held-lock stack
/// permanently. Any resource acquired later in the same task records
/// `channel → resource`.
/// This correctly models the ordering constraint: if this task receives
/// from the channel and then acquires a lock, that ordering is a permanent fact
/// about the task's behaviour, not something that expires.
pub struct Receiver<T, O: ResourceObserver = DefaultObserver> {
    id: Id,
    inner: tokio::sync::mpsc::Receiver<T>,
    observer: O,
}
impl<T, O: ResourceObserver> Receiver<T, O> {
    /// Receive the next value, blocking until one is available.
    ///
    /// After a successful receive, the channel resource is permanently marked
    /// as held for this task.  Any subsequent resource acquisition (lock, send, recv)
    /// records the ordering edge `channel → resource`.
    ///
    /// Returns `None` if the channel is closed.
    #[track_caller]
    pub fn recv(&mut self) -> impl Future<Output = Option<T>> + '_ {
        let caller = Location::caller();
        async move {
            let id = Resource::Channel(self.id);
            self.observer.on_waiting(&id, caller);
            let result = self.inner.recv().await;
            if result.is_some() {
                self.observer.on_acquired(&id, caller);
                // No on_released: channel stays on the held-lock stack so that
                // any later resource acquisition records channel -> resource.
            }
            result
        }
    }

    /// Get the resource id
    pub fn id(&self) -> Id {
        self.id
    }
}
