use std::sync::atomic::{AtomicBool, Ordering};
use async_event::Event;


/// An inhouse lock implementation
/// The purpose is to have a lightweight async lock that doesn't busywait under contention,
/// instead the task is suspended until lock is freed.
pub struct Lock {
    state: AtomicBool,
    listener: Event
}

pub struct LockGuard<'a> {
    inner: &'a Lock
}

impl Lock {
    pub fn new() -> Self {
        Self {
            state: AtomicBool::new(false),
            listener: Event::new()
        }
    }

    pub async fn lock(&'_ self) -> LockGuard<'_> {
        // We first start by checking if lock is acquired
        if self.state.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
            // Doing a lockspin in case we acquire lock before fallbacking to more expensive call
            for _ in 0..10 {
                if !self.state.load(Ordering::Relaxed) {
                    if self.state.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed).is_ok() {
                        return LockGuard { inner: self }
                    }
                    // If we couldn't change state anyways, we will fallback to waiting
                    break;
                }
                std::hint::spin_loop();
            }

            // We have contention, we suspend until event that lock is freed
            self.listener.wait_until(|| {
                if self.state.swap(true, Ordering::Acquire) {
                    None
                } else {
                    Some(())
                }
            }).await;
        }

        LockGuard { inner: self }
    }
}

impl<'a> Drop for LockGuard<'a> {
    fn drop(&mut self) {
        self.inner.state.store(false, Ordering::Release);
        self.inner.listener.notify_one();
    }
}

#[cfg(test)]
mod tests {
    use futures_lite::{future::block_on, FutureExt};
    use super::*;
    use std::{thread, cell::UnsafeCell, sync::Arc};

    struct Container {
        lock: Lock,
        inc: UnsafeCell<u64>
    }
    unsafe impl Sync for Container {}

    #[test]
    fn lock_contented() {
        let obj = Arc::new(Container {
            lock: Lock::new(),
            inc: UnsafeCell::new(0u64)
        });
        
        let mut threads = Vec::new();

        // Checking if there is no race condition or deadlocks between concurrent access to locks
        for _ in 0..9 {
            let objc = obj.clone();
            threads.push(thread::spawn(|| {
                block_on(async move {
                    for _ in 0..1000000 {
                        let guard = objc.lock.lock().await;
                        unsafe { *objc.inc.get() += 1 };
                        drop(guard);
                    }
                });
            }));
        }

        // We also need to check if notify_one is properly awaking a non cancelled future
        let objc = obj.clone();
        threads.push(thread::spawn(|| {
            block_on(async move {
                let lock = Lock::new();
                for _ in 0..1000000 {
                    // We randomly poll one of these two futures to check if we are properly
                    // notifying other other non cancelled futures
                    let fut = objc.lock.lock();
                    fut.race(lock.lock()).await;
                }
            })
        }));


        for thread in threads {
            _ = thread.join();
        }

        unsafe {
            assert_eq!(*obj.inc.get(), 9000000);
        }
    }
}
