use core::{cell::UnsafeCell, hint, ops::{Deref, DerefMut}, sync::atomic::{AtomicBool, Ordering}};

pub struct Spinlock<T> {
  data: UnsafeCell<T>,
  locked: AtomicBool,
}

impl<T> Spinlock<T> {
  pub const fn new(data: T) -> Self {
    Self {
      data: UnsafeCell::new(data),
      locked: AtomicBool::new(false),
    }
  }

  pub fn lock(&self) -> SpinlockGuard<T> {
    // If we observe `locked` was `false`, then:
    // - Use acquire ordering, so nothing inside the critical section gets reordered before we
    //   observed the `false` value.
    // - Store `true`, so nothing else can enter the critical section until we exit it. 
    // Otherwise, spin until we observe a `false` value.
    while self.locked
      .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
      .is_err()
    {
      hint::spin_loop();
    }

    SpinlockGuard { lock: self }
  }

  /// # Safety
  /// There must be no "active" `SpinlockGuards` for this lock, i.e. a `SpinlockGuard` which can be
  /// used to obtain a reference to the spinlock-protected data.
  unsafe fn unlock(&self) {
    // Unset `locked` with release ordering so that nothing inside the critical section gets
    // reordered to after we stored `false`.
    self.locked.store(false, Ordering::Release);
  }

  /// # Safety
  /// No mutable references to the spinlock-protected data may exist.
  unsafe fn get<'s, 'a>(&'s self) -> &'a T
  where
    's: 'a,
  {
    unsafe { &*self.data.get() }
  }

  /// # Safety
  /// No references to the spinlock-protected data may exist.
  unsafe fn get_mut<'s, 'a>(&'s self) -> &'a mut T
  where
    's: 'a,
  {
    unsafe { &mut *self.data.get() }
  }
}

unsafe impl<T> Sync for Spinlock<T>
where
  T: Send,
{}

pub struct SpinlockGuard<'a, T> {
  lock: &'a Spinlock<T>,
}

impl<'a, T> Deref for SpinlockGuard<'a, T> {
  type Target = T;

  #[inline]
  fn deref(&self) -> &Self::Target {
    // SAFETY:
    // For the entire lifetime of the `SpinlockGuard`, `locked` remains `true`, so we have
    // exclusive access to `data`, so no mutable references to `data` can exist.
    unsafe { self.lock.get() }
  }
}

impl<'a, T> DerefMut for SpinlockGuard<'a, T> {
  #[inline]
  fn deref_mut(&mut self) -> &mut Self::Target {
    // SAFETY:
    // For the entire lifetime of the `SpinlockGuard`, `locked` remains `true`, so we have
    // exclusive access to `data`, so no other references to `data` can exist.
    unsafe { self.lock.get_mut() }
  }
}

impl<'a, T> Drop for SpinlockGuard<'a, T> {
  fn drop(&mut self) {
    // SAFETY:
    // Only one `SpinlockGuard` can exist at a time for a particular lock, since we set `locked`
    // to true before creating a guard and refuse to create any new ones until it is `false` again.
    // Therefore, we are the only `SpinlockGuard` for the lock. Since this is the destructor, and
    // we don't access the spinlock-protected data here, there are therefore no "active"
    // `SpinlockGuard`s remaining for the lock.
    unsafe { self.lock.unlock() }
  }
}