//! Mutual exclusion primitive implementation

use core::{
    cell::UnsafeCell,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicU32, Ordering},
    time::Duration,
};

use abi::{error::Error, process::MutexOperation};

/// Raw implementation of the mutual exclusion lock.
pub struct RawMutex {
    state: AtomicU32,
}

/// Guards a value of type `T` by a mutual exclusion lock.
pub struct Mutex<T> {
    value: UnsafeCell<T>,
    lock: RawMutex,
}

/// Represents a method to access a [Mutex], currently locked by this structure's owner.
pub struct MutexGuard<'a, T> {
    mutex: &'a Mutex<T>,
    _pd: PhantomData<&'a mut T>,
}

impl RawMutex {
    const UNLOCKED: u32 = 0;
    const LOCKED: u32 = 1;

    /// Constructs a new unlocked mutex.
    pub const fn new() -> Self {
        Self {
            state: AtomicU32::new(Self::UNLOCKED),
        }
    }

    /// Attempts to gain a lock on a mutex immediately.
    ///
    /// Returns `true` if mutex can be locked instantly.
    pub fn try_lock(&self) -> bool {
        self.state
            .compare_exchange(
                Self::UNLOCKED,
                Self::LOCKED,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_ok()
    }

    /// Gains a lock on a mutex. Blocks if mutex is currently locked by someone else.
    pub fn lock(&self) {
        loop {
            if self.try_lock() {
                return;
            }

            self.wait(None);
        }
    }

    /// Gains a lock on a mutex. Block if mutex is currently locked by someone else, but waits
    /// for the mutex to become available no more than `timeout`.
    ///
    /// Returns [Error::TimedOut] if the function could not gain a lock within the requested time.
    pub fn lock_timeout(&self, timeout: Duration) -> Result<(), Error> {
        if self.try_lock() {
            return Ok(());
        }

        self.wait(Some(timeout));

        if self.try_lock() {
            Ok(())
        } else {
            Err(Error::TimedOut)
        }
    }

    /// Releases the mutex lock.
    ///
    /// # Safety
    ///
    /// The caller must ensure the lock is owned/held by themselves.
    pub unsafe fn release(&self) {
        if self.state.swap(Self::UNLOCKED, Ordering::Release) == Self::LOCKED {
            self.notify_one();
        }
    }

    fn wait(&self, duration: Option<Duration>) {
        unsafe {
            crate::sys::mutex(
                &self.state,
                &MutexOperation::WaitWhileEqual(Self::LOCKED, duration),
            )
        }
        .expect("call to mutex wait failed");
    }

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

impl<T> Mutex<T> {
    /// Constructs a new [Mutex], wrapping the given `value`.
    pub const fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            lock: RawMutex::new(),
        }
    }

    /// Attempts to immediately gain a lock on a mutex.
    ///
    /// Returns [Some] if lock can be acquired instantly, [None] if someone else is already holding
    /// the lock at the moment.
    pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
        if self.lock.try_lock() {
            Some(MutexGuard {
                mutex: self,
                _pd: PhantomData,
            })
        } else {
            None
        }
    }

    /// Gains a lock on a mutex, returning a [MutexGuard].
    pub fn lock(&self) -> MutexGuard<'_, T> {
        self.lock.lock();
        MutexGuard {
            mutex: self,
            _pd: PhantomData,
        }
    }

    /// Same at [Mutex::lock], except blocking for no more than `timeout` specified.
    pub fn lock_timeout(&self, timeout: Duration) -> Result<MutexGuard<'_, T>, Error> {
        self.lock.lock_timeout(timeout)?;
        Ok(MutexGuard {
            mutex: self,
            _pd: PhantomData,
        })
    }
}

unsafe impl<T: Send> Send for Mutex<T> {}
unsafe impl<T: Sync> Sync for Mutex<T> {}

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

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

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

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

unsafe impl<T: Sync> Sync for MutexGuard<'_, T> {}
