#![forbid(unsafe_code)]

use core::future::Future;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::task::{Context, Poll};
use std::collections::VecDeque;
use std::sync::TryLockError;
use std::task::Waker;

/// An [RAII](https://doc.rust-lang.org/rust-by-example/scope/raii.html)
/// scoped lock of a [`Mutex`](struct.Mutex.html).
/// It automatically unlocks the mutex when dropped (falls out of scope).
///
/// You can access the data in the mutex through this guard's
/// [`Deref`](https://doc.rust-lang.org/stable/std/ops/trait.Deref.html)
/// and
/// [`DerefMut`](https://doc.rust-lang.org/stable/std/ops/trait.DerefMut.html)
/// implementations.
///
/// The struct is not
/// [`Send`](https://doc.rust-lang.org/stable/std/marker/trait.Send.html)
/// so you cannot await while holding it.
///
/// If a task panics while holding the struct, the underlying mutex becomes
/// ["poisoned"](https://doc.rust-lang.org/stable/std/sync/struct.Mutex.html#poisoning)
/// and subsequent calls to `lock` will panic.
#[allow(clippy::module_name_repetitions)]
pub struct MutexGuard<'a, T> {
    mutex: &'a Mutex<T>,
    value_guard: Option<std::sync::MutexGuard<'a, T>>,
}
impl<'a, T> MutexGuard<'a, T> {
    pub fn new(
        mutex: &'a Mutex<T>,
        value_guard: std::sync::MutexGuard<'a, T>,
    ) -> MutexGuard<'a, T> {
        let mut inner_guard = mutex.inner.lock().unwrap();
        assert!(!inner_guard.locked);
        inner_guard.locked = true;
        MutexGuard {
            mutex,
            value_guard: Some(value_guard),
        }
    }
}
impl<'a, T> Drop for MutexGuard<'a, T> {
    fn drop(&mut self) {
        let mut wakers = VecDeque::new();
        {
            let mut inner_guard = self.mutex.inner.lock().unwrap();
            assert!(inner_guard.locked);
            inner_guard.locked = false;
            std::mem::swap(&mut inner_guard.wakers, &mut wakers);
        }
        self.value_guard.take();
        for waker in wakers {
            waker.wake();
        }
    }
}
impl<'a, T> Deref for MutexGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &*self.value_guard.as_ref().unwrap()
    }
}
impl<'a, T> DerefMut for MutexGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut *self.value_guard.as_mut().unwrap()
    }
}

#[doc(hidden)]
pub struct LockFuture<'a, T> {
    mutex: &'a Mutex<T>,
}
impl<'a, T> Future for LockFuture<'a, T> {
    type Output = MutexGuard<'a, T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match self.mutex.value.try_lock() {
                Ok(guard) => return Poll::Ready(MutexGuard::new(&self.mutex, guard)),
                Err(TryLockError::Poisoned(e)) => panic!("{}", e),
                Err(TryLockError::WouldBlock) => {}
            }
            let mut guard = self.mutex.inner.lock().unwrap();
            if guard.locked {
                // Mutex is locked.  Add our waker and sleep.
                guard.wakers.push_back(cx.waker().clone());
                return Poll::Pending;
            }
            // Mutex is now unlocked.  Try to acquire it again.
        }
    }
}

struct Inner {
    wakers: VecDeque<Waker>,
    locked: bool,
}

/// A wrapper around
/// [`std::sync::Mutex`](https://doc.rust-lang.org/stable/std/sync/struct.Mutex.html)
/// with an async [`lock`](#method.lock) method.
pub struct Mutex<T> {
    inner: std::sync::Mutex<Inner>,
    value: std::sync::Mutex<T>,
}

impl<T> Mutex<T> {
    pub fn new(value: T) -> Mutex<T> {
        Self {
            inner: std::sync::Mutex::new(Inner {
                wakers: VecDeque::new(),
                locked: false,
            }),
            value: std::sync::Mutex::new(value),
        }
    }

    /// Acquires the mutex, sleeping the current task until it is able to do so.
    ///
    /// After this function returns, is the only task holding the lock.
    ///
    /// The returned guard is an
    /// [RAII](https://doc.rust-lang.org/rust-by-example/scope/raii.html)
    /// scoped lock.
    /// It automatically unlocks the mutex when dropped (falls out of scope).
    ///
    /// The returned guard is not
    /// [`Send`](https://doc.rust-lang.org/stable/std/marker/trait.Send.html)
    /// so you cannot `await` while holding it.
    ///
    /// If a task panics while holding the guard, the underlying mutex becomes
    /// ["poisoned"](https://doc.rust-lang.org/stable/std/sync/struct.Mutex.html#poisoning)
    /// and subsequent calls to `lock` will panic.
    pub async fn lock(&self) -> MutexGuard<'_, T> {
        LockFuture { mutex: self }.await
    }
}

#[cfg(test)]
mod tests {
    use super::super::tests::expect_elapsed;
    use super::*;
    use core::time::Duration;
    use safina_async_test::async_test;
    use safina_timer::sleep_for;
    use std::sync::Arc;
    use std::time::Instant;

    #[async_test]
    async fn should_unlock_immediately() {
        let before = Instant::now();
        let counter = Mutex::new(0_u8);
        {
            let mut guard = counter.lock().await;
            *guard += 1;
        }
        let guard = counter.lock().await;
        assert_eq!(1_u8, *guard);
        expect_elapsed(before, 0..1);
    }

    #[async_test]
    async fn unlocking_wakes_waiters() {
        let executor = safina_executor::Executor::new(1);
        safina_timer::start_timer_thread();
        let before = Instant::now();
        let counter = Arc::new(Mutex::new(0_u8));
        let counter_clone = counter.clone();
        executor.spawn(async move {
            let mut guard = counter_clone.lock().await;
            *guard += 1;
        });
        // Delay so other task can get the lock
        sleep_for(Duration::from_millis(100)).await;
        let guard = counter.lock().await;
        assert_eq!(1_u8, *guard);
        expect_elapsed(before, 100..200);
    }
    #[async_test]
    async fn should_handle_many_waiters() {
        const NUM_WAITERS: u32 = 100_000;
        let executor = safina_executor::Executor::new(10);
        safina_timer::start_timer_thread();
        let before = Instant::now();
        let counter = Arc::new(Mutex::new(0_u32));
        let mut receivers = Vec::new();
        for _ in 0..NUM_WAITERS {
            let counter_clone = counter.clone();
            let (sender, receiver) = std::sync::mpsc::channel::<()>();
            receivers.push(receiver);
            executor.spawn(async move {
                {
                    let mut guard = counter_clone.lock().await;
                    *guard += 1;
                }
                sender.send(()).unwrap();
            });
        }
        for receiver in receivers {
            receiver.recv().unwrap();
        }
        let guard = counter.lock().await;
        assert_eq!(NUM_WAITERS, *guard);
        expect_elapsed(before, 0..5000);
    }
}
