use std::marker::PhantomData;

use crate::{
	context::{LockContext, LockingIterator, LockingTuple},
	lockable::{Lockable, OwnedLockable},
	ThreadKey,
};

impl<'l, L> LockContext<'l, L> {
	/// Safety: Don't lock the locks in a different order than what the
	///         `LockContext` uses
	pub(crate) const unsafe fn new(lockable: &'l L) -> Self
	where
		L: OwnedLockable,
	{
		Self {
			key: None,
			lockable,
		}
	}

	/// Unlocks all locks in the collection, returning the [`ThreadKey`].
	///
	/// This requires a mutable reference to the context, so it cannot be called
	/// without first dropping any [`ContextGuard`]s that reference this context.
	/// This method will also return `None` if the context has not been locked
	/// with a `ThreadKey`.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = (Mutex::new(42), Mutex::new(true));
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let tuple = ctx.tuple(key);
	///
	/// let (use_other, tuple) = tuple.lock_1();
	/// if **use_other {
	///     drop(use_other);
	///     drop(tuple);
	///     let key = ctx.unlock().unwrap();
	///     let tuple = ctx.tuple(key);
	///     let (mut item, _) = tuple.lock_0();
	///     **item = 67;
	/// } else {
	///     drop(use_other);
	///     drop(tuple);
	/// };
	///
	/// let key = ctx.unlock().unwrap();
	/// let tuple = ctx.tuple(key);
	/// let (number, _) = tuple.lock_0();
	/// assert_eq!(**number, 67);
	/// ```
	///
	/// [`ContextGuard`]: `crate::context::ContextGuard`
	// The mutable reference ensures that all of the guards, which have a shared
	// borrow to the context, must be dropped first
	pub fn unlock(&mut self) -> Option<ThreadKey> {
		self.key.take()
	}
}

impl<L: Lockable> LockContext<'_, L> {
	/// Creates a [`LockingTuple`], which can lock a subset of a tuple of locks,
	/// in a specific order.
	///
	/// Sometimes, partial allocation of locks is useful. For example, you may want
	/// to acquire a lock on one item before deciding if the second item should be
	/// locked. If the locks can be organized into a tuple, [`LockingTuple`] is
	/// capable of doing exactly that.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = (Mutex::new(true), Mutex::new(42), Mutex::new(67));
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let tuple = ctx.tuple(key);
	///
	/// let (use_other, tuple) = tuple.lock_0();
	/// let number = if **use_other {
	///     tuple.lock_2().0
	/// } else {
	///     tuple.lock_1().0
	/// };
	/// assert_eq!(**number, 67);
	/// ```
	pub fn tuple(&mut self, key: ThreadKey) -> LockingTuple<'_, L, L> {
		unsafe {
			self.key = Some(key);

			LockingTuple {
				_lockable: PhantomData,
				// safety: we just inserted a key
				key: self.key.as_ref().unwrap_unchecked(),
				tuple: self.lockable,
				outer: (),
			}
		}
	}
}

impl<'l, L> LockContext<'l, L>
where
	&'l L: IntoIterator,
{
	/// Creates a [`LockingIterator`] to iterate through a collection of locks
	/// without locking everything at once.
	///
	/// Sometimes, partial allocation of locks is useful. For example, you may
	/// want to acquire a lock on the first element of a list before deciding if
	/// the second element should be locked. If the list is iterable, then a
	/// [`LockingIterator`] is capable of doing exactly that.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [Mutex::new(1), Mutex::new(3), Mutex::new(8)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key);
	///
	/// let mut sum = 0;
	/// while let Some(item) = iter.lock_next() {
	///     sum += **item;
	/// }
	///
	/// assert_eq!(sum, 12);
	/// ```
	// TODO: support scoped locks
	// TODO: implement get_disjoint
	// TODO: support some sort of index tower thing
	#[expect(clippy::iter_not_returning_iterator)]
	pub fn iter(
		&mut self,
		key: ThreadKey,
	) -> LockingIterator<'_, <&'l L as IntoIterator>::IntoIter> {
		unsafe {
			self.key = Some(key);

			LockingIterator {
				// safety: we just inserted a key
				key: self.key.as_ref().unwrap_unchecked(),
				iterator: self.lockable.into_iter(),
				outer: (),
			}
		}
	}
}