//! Read-write lock mechanism implementation
use core::{
    cell::UnsafeCell,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicU32, Ordering},
};

use abi::process::MutexOperation;

/// Defines a read-write lock primitive.
pub struct RwLock<T> {
    value: UnsafeCell<T>,
    state: AtomicU32,
}

/// Represents a read-only access mechanism to a currently held [RwLock].
pub struct RwLockReadGuard<'a, T> {
    lock: &'a RwLock<T>,
    _pd: PhantomData<&'a T>,
}

/// Represents a read-write access mechanism to a currently held [RwLock].
pub struct RwLockWriteGuard<'a, T> {
    lock: &'a RwLock<T>,
    _pd: PhantomData<&'a mut T>,
}

impl<T> RwLock<T> {
    const LOCKED_WRITE: u32 = 1 << 0;
    const LOCKED_READ: u32 = 1 << 2;

    /// Constructs a new [RwLock], wrapping the `value`.
    pub const fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            state: AtomicU32::new(0),
        }
    }

    /// Tries to acquire a read-only lock on the [RwLock] instantly.
    ///
    /// # Panics
    ///
    /// Will panic if there're more than `1 << 30` readers present.
    pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
        let value = self.state.fetch_add(Self::LOCKED_READ, Ordering::Acquire);
        // Prevent overflows
        if value >= (u32::MAX - Self::LOCKED_READ * 2) {
            panic!("Too many readers holding a RwLock!");
        }
        if value & Self::LOCKED_WRITE == 0 {
            // Can't acquire
            Some(RwLockReadGuard {
                lock: self,
                _pd: PhantomData,
            })
        } else {
            self.state.fetch_sub(Self::LOCKED_READ, Ordering::Release);
            None
        }
    }

    /// Tries to acquire a read-write lock on the [RwLock] immediately.
    pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
        if self
            .state
            .compare_exchange(0, Self::LOCKED_WRITE, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            Some(RwLockWriteGuard {
                lock: self,
                _pd: PhantomData,
            })
        } else {
            None
        }
    }

    /// Acquires a read-only lock on the [RwLock], blocking for the lock to become read-available
    /// if necessary.
    ///
    /// # Panics
    ///
    /// Will panic if there're more than `1 << 30` readers present.
    pub fn read(&self) -> RwLockReadGuard<'_, T> {
        loop {
            if let Some(guard) = self.try_read() {
                return guard;
            }

            self.wait_readable();
        }
    }

    /// Acquires a read-write lock on the [RwLock], blocking for the lock to become write-available
    /// if necessary.
    pub fn write(&self) -> RwLockWriteGuard<'_, T> {
        loop {
            if let Some(guard) = self.try_write() {
                return guard;
            }

            self.wait_writeable();
        }
    }

    /// Releases a read-access lock to this [RwLock].
    ///
    /// # Safety
    ///
    /// Intended to be called from [RwLockReadGuard]'s [Drop] handler.
    pub unsafe fn release_read(&self) {
        if self.state.fetch_sub(Self::LOCKED_READ, Ordering::Release) == Self::LOCKED_READ {
            // Only one lock and only for read access was held, which meant the RwLock just became
            // both readable and writeable, notify a writer about that
            self.notify_writeable();
        }
    }

    /// Releases a read-write-access lock to this [RwLock].
    ///
    /// # Safety
    ///
    /// Intended to be called from [RwLockWriteGuard]'s [Drop] handler.
    pub unsafe fn release_write(&self) {
        self.state.fetch_and(!Self::LOCKED_WRITE, Ordering::Release);
        self.notify_readable();
    }

    fn wait_readable(&self) {
        // TODO bitwise wait operations in mutex syscall
        core::hint::spin_loop();
    }

    fn wait_writeable(&self) {
        unsafe { crate::sys::mutex(&self.state, &MutexOperation::WaitUntilEqual(0, None)) }
            .expect("wait_writeable() failed");
    }

    // nop
    fn notify_readable(&self) {}

    fn notify_writeable(&self) {
        unsafe { crate::sys::mutex(&self.state, &MutexOperation::Wake(1)) }.ok();
    }
}

impl<T> Deref for RwLockReadGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.lock.value.get() }
    }
}

impl<T> Drop for RwLockReadGuard<'_, T> {
    fn drop(&mut self) {
        unsafe { self.lock.release_read() };
    }
}

impl<T> Deref for RwLockWriteGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.lock.value.get() }
    }
}

impl<T> DerefMut for RwLockWriteGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.lock.value.get() }
    }
}

impl<T> Drop for RwLockWriteGuard<'_, T> {
    fn drop(&mut self) {
        unsafe { self.lock.release_write() };
    }
}

unsafe impl<T: Sync> Sync for RwLock<T> {}
