use core::sync::atomic::Ordering;

use cfg_if::cfg_if;

cfg_if! {
	if #[cfg(any(
		target_os = "linux",
		target_os = "android",
		//all(target_os = "emscripten", target_feature = "atomics"),
		//target_os = "freebsd",
		//target_os = "openbsd",
		//target_os = "dragonfly",
		//target_os = "fuchsia",
	))] {
		mod linux;
		use linux::{Futex, Primitive, SmallFutex, SmallPrimitive};
	} else if #[cfg(all(target_os = "windows", not(target_vendor = "win7")))] {
		mod windows;
		use windows::Futex;
	}
}

const SPIN_COUNT: u32 = 100;

const UNLOCKED: SmallPrimitive = 0;
const LOCKED: SmallPrimitive = 1; // locked, no other threads waiting
const CONTENDED: SmallPrimitive = 2; // locked, and other threads waiting (contended)

pub struct Mutex(SmallFutex);

impl Mutex {
	#[inline]
	pub const fn new() -> Self {
		Self(Futex::new(UNLOCKED))
	}

	/// Locks the mutex
	///
	/// # Safety
	///
	/// UB occurs if the mutex is already locked by the current thread and the
	/// `unsafe_lock` feature is enabled.
	#[inline]
	pub unsafe fn lock(&self) {
		if self
			.0
			.compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
			.is_err()
		{
			self.lock_contended();
		}
	}

	/// If the mutex is unlocked, it is locked, and this function returns
	/// `true'. Otherwise, `false` is returned.
	#[inline]
	pub unsafe fn try_lock(&self) -> bool {
		self.0
			.compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
			.is_ok()
	}

	/// Unlocks the mutex
	///
	/// # Safety
	///
	/// UB occurs if the mutex is already unlocked or if it has been locked on
	/// a different thread.
	#[inline]
	pub unsafe fn unlock(&self) {
		if self.0.swap(UNLOCKED, Ordering::Release) == CONTENDED {
			// We only wake up one thread. When that thread locks the mutex, it
			// will mark the mutex as CONTENDED (see lock_contended above),
			// which makes sure that any other waiting threads will also be
			// woken up eventually.
			self.0.wake();
		}
	}

	#[inline]
	pub unsafe fn is_locked(&self) -> bool {
		self.0.load(Ordering::Relaxed) == UNLOCKED
	}

	#[cold]
	fn lock_contended(&self) {
		// Spin first to speed things up if the lock is released quickly.
		let mut state = self.spin();

		// If it's unlocked now, attempt to take the lock
		// without marking it as contended.
		if state == UNLOCKED {
			match self
				.0
				.compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
			{
				Ok(_) => return, // Locked!
				Err(s) => state = s,
			}
		}

		loop {
			// Put the lock in contended state.
			// We avoid an unnecessary write if it as already set to CONTENDED,
			// to be friendlier for the caches.
			if state != CONTENDED && self.0.swap(CONTENDED, Ordering::Acquire) == UNLOCKED {
				// We changed it from UNLOCKED to CONTENDED, so we just successfully locked it.
				return;
			}

			// Wait for the futex to change state, assuming it is still CONTENDED.
			self.0.wait(CONTENDED);

			// Spin again after waking up.
			state = self.spin();
		}
	}

	fn spin(&self) -> SmallPrimitive {
		let mut spin = SPIN_COUNT;
		loop {
			// We only use `load` (and not `swap` or `compare_exchange`)
			// while spinning, to be easier on the caches.
			let state = self.0.load(Ordering::Acquire);

			// We stop spinning when the mutex is UNLOCKED,
			// but also when it's CONTENDED.
			if state != LOCKED || spin == 0 {
				return state;
			}

			core::hint::spin_loop();
			spin -= 1;
		}
	}
}

pub struct RwLock {
	// The state consists of a 30-bit reader counter, a 'readers waiting' flag, and a 'writers waiting' flag.
	// Bits 0.30:
	//   0: Unlocked
	//   1.=0x3FFF_FFFE: Locked by N readers
	//   0x3FFF_FFFF: Write locked
	// Bit 30: Readers are waiting on this futex.
	// Bit 31: Writers are waiting on the writer_notify futex.
	state: Futex,
	// The 'condition variable' to notify writers through.
	// Incremented on every signal.
	writer_notify: Futex,
}

const READ_LOCKED: Primitive = 1;
const MASK: Primitive = (1 << 30) - 1;
const WRITE_LOCKED: Primitive = MASK;
const DOWNGRADE: Primitive = READ_LOCKED.wrapping_sub(WRITE_LOCKED); // READ_LOCKED - WRITE_LOCKED
const MAX_READERS: Primitive = MASK - 1;
const READERS_WAITING: Primitive = 1 << 30;
const WRITERS_WAITING: Primitive = 1 << 31;

#[inline]
fn is_unlocked(state: Primitive) -> bool {
	state & MASK == 0
}

#[inline]
fn is_write_locked(state: Primitive) -> bool {
	state & MASK == WRITE_LOCKED
}

#[inline]
fn has_readers_waiting(state: Primitive) -> bool {
	state & READERS_WAITING != 0
}

#[inline]
fn has_writers_waiting(state: Primitive) -> bool {
	state & WRITERS_WAITING != 0
}

#[inline]
fn is_read_lockable(state: Primitive) -> bool {
	// This also returns false if the counter could overflow if we tried to read lock it.
	//
	// We don't allow read-locking if there's readers waiting, even if the lock is unlocked
	// and there's no writers waiting. The only situation when this happens is after unlocking,
	// at which point the unlocking thread might be waking up writers, which have priority over readers.
	// The unlocking thread will clear the readers waiting bit and wake up readers, if necessary.
	state & MASK < MAX_READERS && !has_readers_waiting(state) && !has_writers_waiting(state)
}

#[inline]
fn is_read_lockable_after_wakeup(state: Primitive) -> bool {
	// We make a special case for checking if we can read-lock _after_ a reader thread that went to
	// sleep has been woken up by a call to `downgrade`.
	//
	// `downgrade` will wake up all readers and place the lock in read mode. Thus, there should be
	// no readers waiting and the lock should be read-locked (not write-locked or unlocked).
	//
	// Note that we do not check if any writers are waiting. This is because a call to `downgrade`
	// implies that the caller wants other readers to read the value protected by the lock. If we
	// did not allow readers to acquire the lock before writers after a `downgrade`, then only the
	// original writer would be able to read the value, thus defeating the purpose of `downgrade`.
	state & MASK < MAX_READERS
		&& !has_readers_waiting(state)
		&& !is_write_locked(state)
		&& !is_unlocked(state)
}

#[inline]
fn has_reached_max_readers(state: Primitive) -> bool {
	state & MASK == MAX_READERS
}

impl RwLock {
	#[inline]
	pub const fn new() -> Self {
		Self {
			state: Futex::new(0),
			writer_notify: Futex::new(0),
		}
	}

	#[inline]
	pub fn try_read(&self) -> bool {
		self.state
			.fetch_update(Ordering::Acquire, Ordering::Relaxed, |s| {
				is_read_lockable(s).then(|| s + READ_LOCKED)
			})
			.is_ok()
	}

	#[inline]
	pub fn read(&self) {
		let state = self.state.load(Ordering::Relaxed);
		if !is_read_lockable(state)
			|| self
				.state
				.compare_exchange_weak(
					state,
					state + READ_LOCKED,
					Ordering::Acquire,
					Ordering::Relaxed,
				)
				.is_err()
		{
			self.read_contended();
		}
	}

	/// # Safety
	///
	/// The `RwLock` must be read-locked (N readers) in order to call this.
	#[inline]
	pub unsafe fn read_unlock(&self) {
		let state = self.state.fetch_sub(READ_LOCKED, Ordering::Release) - READ_LOCKED;

		// It's impossible for a reader to be waiting on a read-locked RwLock,
		// except if there is also a writer waiting.
		debug_assert!(!has_readers_waiting(state) || has_writers_waiting(state));

		// Wake up a writer if we were the last reader and there's a writer waiting.
		if is_unlocked(state) && has_writers_waiting(state) {
			self.wake_writer_or_readers(state);
		}
	}

	#[cold]
	fn read_contended(&self) {
		let mut has_slept = false;
		let mut state = self.spin_read();

		loop {
			// If we have just been woken up, first check for a `downgrade` call.
			// Otherwise, if we can read-lock it, lock it.
			if (has_slept && is_read_lockable_after_wakeup(state)) || is_read_lockable(state) {
				match self.state.compare_exchange_weak(
					state,
					state + READ_LOCKED,
					Ordering::Acquire,
					Ordering::Relaxed,
				) {
					Ok(_) => return, // Locked!
					Err(s) => {
						state = s;
						continue;
					}
				}
			}

			// Check for overflow.
			assert!(
				!has_reached_max_readers(state),
				"too many active read locks on RwLock"
			);

			// Make sure the readers waiting bit is set before we go to sleep.
			if !has_readers_waiting(state) {
				if let Err(s) = self.state.compare_exchange(
					state,
					state | READERS_WAITING,
					Ordering::Relaxed,
					Ordering::Relaxed,
				) {
					state = s;
					continue;
				}
			}

			// Wait for the state to change.
			self.state.wait(state | READERS_WAITING);
			has_slept = true;

			// Spin again after waking up.
			state = self.spin_read();
		}
	}

	#[inline]
	pub fn try_write(&self) -> bool {
		self.state
			.fetch_update(Ordering::Acquire, Ordering::Relaxed, |s| {
				is_unlocked(s).then(|| s + WRITE_LOCKED)
			})
			.is_ok()
	}

	#[inline]
	pub fn write(&self) {
		if self
			.state
			.compare_exchange_weak(0, WRITE_LOCKED, Ordering::Acquire, Ordering::Relaxed)
			.is_err()
		{
			self.write_contended();
		}
	}

	/// # Safety
	///
	/// The `RwLock` must be write-locked (single writer) in order to call this.
	#[inline]
	pub unsafe fn write_unlock(&self) {
		let state = self.state.fetch_sub(WRITE_LOCKED, Ordering::Release) - WRITE_LOCKED;

		debug_assert!(is_unlocked(state));

		if has_writers_waiting(state) || has_readers_waiting(state) {
			self.wake_writer_or_readers(state);
		}
	}

	/// # Safety
	///
	/// The `RwLock` must be write-locked (single writer) in order to call this.
	#[inline]
	pub unsafe fn downgrade(&self) {
		// Removes all write bits and adds a single read bit.
		let state = self.state.fetch_add(DOWNGRADE, Ordering::Release);
		debug_assert!(
			is_write_locked(state),
			"RwLock must be write locked to call `downgrade`"
		);

		if has_readers_waiting(state) {
			// Since we had the exclusive lock, nobody else can unset this bit.
			self.state.fetch_sub(READERS_WAITING, Ordering::Relaxed);
			self.state.wake_all();
		}
	}

	#[cold]
	fn write_contended(&self) {
		let mut state = self.spin_write();

		let mut other_writers_waiting = 0;

		loop {
			// If it's unlocked, we try to lock it.
			if is_unlocked(state) {
				match self.state.compare_exchange_weak(
					state,
					state | WRITE_LOCKED | other_writers_waiting,
					Ordering::Acquire,
					Ordering::Relaxed,
				) {
					Ok(_) => return, // Locked!
					Err(s) => {
						state = s;
						continue;
					}
				}
			}

			// Set the waiting bit indicating that we're waiting on it.
			if !has_writers_waiting(state) {
				if let Err(s) = self.state.compare_exchange(
					state,
					state | WRITERS_WAITING,
					Ordering::Relaxed,
					Ordering::Relaxed,
				) {
					state = s;
					continue;
				}
			}

			// Other writers might be waiting now too, so we should make sure
			// we keep that bit on once we manage lock it.
			other_writers_waiting = WRITERS_WAITING;

			// Examine the notification counter before we check if `state` has changed,
			// to make sure we don't miss any notifications.
			let seq = self.writer_notify.load(Ordering::Acquire);

			// Don't go to sleep if the lock has become available,
			// or if the writers waiting bit is no longer set.
			state = self.state.load(Ordering::Relaxed);
			if is_unlocked(state) || !has_writers_waiting(state) {
				continue;
			}

			// Wait for the state to change.
			self.writer_notify.wait(seq);

			// Spin again after waking up.
			state = self.spin_write();
		}
	}

	/// Wakes up waiting threads after unlocking.
	///
	/// If both are waiting, this will wake up only one writer, but will fall
	/// back to waking up readers if there was no writer to wake up.
	#[cold]
	fn wake_writer_or_readers(&self, mut state: Primitive) {
		assert!(is_unlocked(state));

		// The readers waiting bit might be turned on at any point now,
		// since readers will block when there's anything waiting.
		// Writers will just lock the lock though, regardless of the waiting bits,
		// so we don't have to worry about the writer waiting bit.
		//
		// If the lock gets locked in the meantime, we don't have to do
		// anything, because then the thread that locked the lock will take
		// care of waking up waiters when it unlocks.

		// If only writers are waiting, wake one of them up.
		if state == WRITERS_WAITING {
			match self
				.state
				.compare_exchange(state, 0, Ordering::Relaxed, Ordering::Relaxed)
			{
				Ok(_) => {
					self.wake_writer();
					return;
				}
				Err(s) => {
					// Maybe some readers are now waiting too. So, continue to the next `if`.
					state = s;
				}
			}
		}

		// If both writers and readers are waiting, leave the readers waiting
		// and only wake up one writer.
		if state == READERS_WAITING + WRITERS_WAITING {
			if self
				.state
				.compare_exchange(state, READERS_WAITING, Ordering::Relaxed, Ordering::Relaxed)
				.is_err()
			{
				// The lock got locked. Not our problem anymore.
				return;
			}
			if self.wake_writer() {
				return;
			}
			// No writers were actually blocked on futex_wait, so we continue
			// to wake up readers instead, since we can't be sure if we notified a writer.
			state = READERS_WAITING;
		}

		// If readers are waiting, wake them all up.
		if state == READERS_WAITING
			&& self
				.state
				.compare_exchange(state, 0, Ordering::Relaxed, Ordering::Relaxed)
				.is_ok()
		{
			self.state.wake_all();
		}
	}

	/// This wakes one writer and returns true if we woke up a writer that was
	/// blocked on futex_wait.
	///
	/// If this returns false, it might still be the case that we notified a
	/// writer that was about to go to sleep.
	fn wake_writer(&self) -> bool {
		self.writer_notify.fetch_add(1, Ordering::Release);
		self.writer_notify.wake()
		// Note that FreeBSD and DragonFlyBSD don't tell us whether they woke
		// up any threads or not, and always return `false` here. That still
		// results in correct behavior: it just means readers get woken up as
		// well in case both readers and writers were waiting.
	}

	/// Spin for a while, but stop directly at the given condition.
	#[inline]
	fn spin_until(&self, f: impl Fn(Primitive) -> bool) -> Primitive {
		let mut spin = SPIN_COUNT; // Chosen by fair dice roll.
		loop {
			let state = self.state.load(Ordering::Relaxed);
			if f(state) || spin == 0 {
				return state;
			}
			core::hint::spin_loop();
			spin -= 1;
		}
	}

	#[inline]
	fn spin_write(&self) -> Primitive {
		// Stop spinning when it's unlocked or when there's waiting writers, to keep things somewhat fair.
		self.spin_until(|state| is_unlocked(state) || has_writers_waiting(state))
	}

	#[inline]
	fn spin_read(&self) -> Primitive {
		// Stop spinning when it's unlocked or read locked, or when there's waiting threads.
		self.spin_until(|state| {
			!is_write_locked(state) || has_readers_waiting(state) || has_writers_waiting(state)
		})
	}
}