use std::{
	iter::{Fuse, Peekable, Skip, Take},
	marker::PhantomData,
};

use super::{ContextGuard, LockingIterator};

use crate::{
	context::LockingTuple,
	lockable::{Lockable, RawLock, Sharable},
	ThreadKey,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TryLockNextError {
	FinishedIteration,
	WouldBlock,
}

impl<'l, I, O> LockingIterator<'l, I, O> {
	fn with_iterator<M>(self, f: impl FnOnce(I) -> M) -> LockingIterator<'l, M, O> {
		LockingIterator {
			key: self.key,
			iterator: f(self.iterator),
			outer: self.outer,
		}
	}

	/// Exit out of the current scope of the locking iterator into the parent.
	///
	/// After using one the recurse methods, it is possible to regain access to
	/// the parent by exiting out of the scope of the child. Doing this will make
	/// it impossible to re-enter this scope again.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = ([Mutex::new(1), Mutex::new(2), Mutex::new(3)], Mutex::new(true));
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let tuple = ctx.tuple(key);
	/// let mut iter = tuple.recurse_0_iter();
	///
	/// let mut sum = 0;
	/// while let Some(item) = iter.lock_next() {
	///     sum += **item;
	/// }
	///
	/// let tuple = iter.exit();
	/// let (should_assert, _) = tuple.lock_1();
	/// if **should_assert {
	///     assert_eq!(sum, 6);
	/// }
	/// ```
	pub fn exit(self) -> O {
		self.outer
	}
}

impl<'c, L: Iterator<Item = I>, I: IntoIterator, O> LockingIterator<'c, L, O> {
	/// Create a new `LockingIterator` based on the next element in the iterator.
	///
	/// If a list contains a list of locks, then this method can be used to
	/// recurse into the next element of the list. To go back to the parent scope,
	/// use [`LockingIterator::exit`] on the new list.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [
	///     [Mutex::new(1), Mutex::new(2), Mutex::new(3)],
	///     [Mutex::new(4), Mutex::new(5), Mutex::new(6)],
	/// ];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key);
	///
	/// let mut sums = Vec::new();
	/// while let Some(mut list) = iter.recurse_next() {
	///     let mut sum = 0;
	///     while let Some(item) = list.lock_next() {
	///         sum += **item;
	///     }
	///     sums.push(sum);
	///     iter = list.exit();
	/// }
	///
	/// assert_eq!(sums, vec![6, 15]);
	/// ```
	pub fn recurse_next(
		mut self,
	) -> Option<LockingIterator<'c, <I as IntoIterator>::IntoIter, Self>> {
		if let Some(iterator) = self.iterator.next() {
			Some(LockingIterator {
				key: self.key,
				iterator: iterator.into_iter(),
				outer: self,
			})
		} else {
			None
		}
	}

	/// Create a new `LockingIterator` based on the next element in the iterator.
	///
	/// If a list contains a list of locks, then this method can be used to
	/// recurse into the next element of the list. To go back to the parent scope,
	/// use [`LockingIterator::exit`] on the new list.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [
	///     [Mutex::new(1), Mutex::new(2), Mutex::new(3)],
	///     [Mutex::new(4), Mutex::new(5), Mutex::new(6)],
	/// ];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let iter = ctx.iter(key);
	///
	/// let mut list = iter.recurse_last().unwrap();
	/// let mut sum = 0;
	/// while let Some(item) = list.lock_next() {
	///     sum += **item;
	/// }
	///
	/// assert_eq!(sum, 15);
	/// ```
	pub fn recurse_last(self) -> Option<LockingIterator<'c, <I as IntoIterator>::IntoIter, O>> {
		if let Some(iterator) = self.iterator.last() {
			Some(LockingIterator {
				key: self.key,
				iterator: iterator.into_iter(),
				outer: self.outer,
			})
		} else {
			None
		}
	}
}

impl<'c, L: Iterator<Item = &'c T>, T: 'c, O> LockingIterator<'c, L, O> {
	/// Create a new `LockingIterator` based on the next element in the iterator.
	///
	/// If a list contains a list of locks, then this method can be used to
	/// recurse into the next element of the list. To go back to the parent scope,
	/// use [`LockingIterator::exit`] on the new list.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [
	///     (Mutex::new(true), Mutex::new(1)),
	///     (Mutex::new(false), Mutex::new(2)),
	///     (Mutex::new(true), Mutex::new(3)),
	/// ];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key);
	///
	/// let mut sum = 0;
	/// while let Some(tuple) = iter.recurse_next_tuple() {
	///     let (should_count, mut tuple) = tuple.lock_0();
	///     if **should_count {
	///         let num = tuple.lock_mut_1();
	///         sum += **num;
	///     }
	///     iter = tuple.exit();
	/// }
	///
	/// assert_eq!(sum, 4);
	/// ```
	pub fn recurse_next_tuple(mut self) -> Option<LockingTuple<'c, T, T, Self>> {
		if let Some(tuple) = self.iterator.next() {
			Some(LockingTuple {
				key: self.key,
				_lockable: PhantomData,
				tuple,
				outer: self,
			})
		} else {
			None
		}
	}

	/// Create a new `LockingIterator` based on the next element in the iterator.
	///
	/// If a list contains a list of locks, then this method can be used to
	/// recurse into the next element of the list. To go back to the parent scope,
	/// use [`LockingIterator::exit`] on the new list.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [
	///     (Mutex::new(true), Mutex::new(1)),
	///     (Mutex::new(false), Mutex::new(2)),
	///     (Mutex::new(true), Mutex::new(3)),
	/// ];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let iter = ctx.iter(key);
	///
	/// let tuple = iter.recurse_last_tuple().unwrap();
	/// let (should_count, mut tuple) = tuple.lock_0();
	/// if **should_count {
	///     let num = tuple.lock_mut_1();
	///     assert_eq!(**num, 3);
	/// } else {
	///     panic!();
	/// }
	/// ```
	pub fn recurse_last_tuple(self) -> Option<LockingTuple<'c, T, T, O>> {
		if let Some(tuple) = self.iterator.last() {
			Some(LockingTuple {
				key: self.key,
				_lockable: PhantomData,
				tuple,
				outer: self.outer,
			})
		} else {
			None
		}
	}
}

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Lockable, O> LockingIterator<'c, L, O> {
	/// Advances the iterator, locking the next element and returning a guard to
	/// the inner data.
	///
	/// Returns `None` when iteration is finished. Individual iterator
	/// implementations may choose to resume iteration, and so calling `next()`
	/// again may or may not eventually start returning `Some(Item)` again at some
	/// point.
	///
	/// # 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);
	/// ```
	pub fn lock_next(&mut self) -> Option<ContextGuard<'c, <I as Lockable>::Guard<'c>, ThreadKey>> {
		if let Some(lock) = self.iterator.next() {
			unsafe {
				lock.raw_write();
				let guard = lock.guard();

				Some(ContextGuard {
					_key: self.key,
					guard,
				})
			}
		} else {
			None
		}
	}

	/// Consumes the iterator, returning the last element, without locking any
	/// other elements.
	///
	/// This method will evaluate the iterator until it returns `None`. While
	/// doing so, it keeps track of the current element. After `None` is returned,
	/// `lock_last()` will then lock the last element it saw and return the
	/// lock's data.
	///
	/// # Panics
	///
	/// This function might panic if the iterator is infinite.
	///
	/// # 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 iter = ctx.iter(key);
	///
	/// let last = iter.lock_last().unwrap();
	/// assert_eq!(**last, 8);
	/// ```
	pub fn lock_last(self) -> Option<ContextGuard<'c, <I as Lockable>::Guard<'c>, ThreadKey>> {
		self.iterator.last().map(|lock| unsafe {
			lock.raw_write();
			let guard = lock.guard();

			ContextGuard {
				_key: self.key,
				guard,
			}
		})
	}
}

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Lockable, O>
	LockingIterator<'c, Peekable<L>, O>
{
	/// Attempts to lock the next element and returning a guard to
	/// the inner data.
	///
	/// # Errors
	///
	/// Returns `Err(TryLockNextError::FinishedIteration)` when iteration is
	/// finished. Individual iterator implementations may choose to resume
	/// iteration, and so calling `next()` again may or may not eventually start
	/// returning `Some(Item)` again at some point.
	///
	/// Returns `Err(TryLockNextError::WouldBlock)` if the next lock in the
	/// iterator is already locked. This will not advance the iterator.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	/// use happylock::context::iterator::TryLockNextError;
	///
	/// 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).peekable();
	///
	/// let mut sum = 0;
	/// loop {
	///     match iter.try_lock_next() {
	///         Ok(item) => sum += **item,
	///         Err(TryLockNextError::WouldBlock) => continue,
	///         Err(TryLockNextError::FinishedIteration) => break,
	///     }
	/// }
	///
	/// assert_eq!(sum, 12);
	/// ```
	pub fn try_lock_next(
		&mut self,
	) -> Result<ContextGuard<'c, <I as Lockable>::Guard<'c>, ThreadKey>, TryLockNextError> {
		if let Some(lock) = self.iterator.peek().copied() {
			unsafe {
				if lock.raw_try_write() {
					// safety: we just saw that there is a valid value
					let lock = self.iterator.next().unwrap_unchecked();
					let guard = lock.guard();

					Ok(ContextGuard {
						_key: self.key,
						guard,
					})
				} else {
					Err(TryLockNextError::WouldBlock)
				}
			}
		} else {
			Err(TryLockNextError::FinishedIteration)
		}
	}
}

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Sharable, O> LockingIterator<'c, L, O> {
	/// Advances the iterator, acquiring a shared lock to the next element and
	/// returning a guard to the inner data.
	///
	/// Returns `None` when iteration is finished. Individual iterator
	/// implementations may choose to resume iteration, and so calling `next()`
	/// again may or may not eventually start returning `Some(Item)` again at some
	/// point.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{RwLock, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [RwLock::new(1), RwLock::new(3), RwLock::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.read_next() {
	///     sum += **item;
	/// }
	///
	/// assert_eq!(sum, 12);
	/// ```
	pub fn read_next(
		&mut self,
	) -> Option<ContextGuard<'c, <I as Sharable>::ReadGuard<'c>, ThreadKey>> {
		if let Some(lock) = self.iterator.next() {
			unsafe {
				lock.raw_read();
				let guard = lock.read_guard();

				Some(ContextGuard {
					_key: self.key,
					guard,
				})
			}
		} else {
			None
		}
	}

	/// Consumes the iterator, returning the last element with readonly access,
	/// without locking any other elements.
	///
	/// This method will evaluate the iterator until it returns `None`. While
	/// doing so, it keeps track of the current element. After `None` is returned,
	/// `lock_last()` will then lock the last element it saw and return the
	/// lock's data.
	///
	/// # Panics
	///
	/// This function might panic if the iterator is infinite.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{RwLock, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [RwLock::new(1), RwLock::new(3), RwLock::new(8)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let iter = ctx.iter(key);
	///
	/// let last = iter.read_last().unwrap();
	/// assert_eq!(**last, 8);
	/// ```
	pub fn read_last(self) -> Option<ContextGuard<'c, <I as Sharable>::ReadGuard<'c>, ThreadKey>> {
		self.iterator.last().map(|lock| unsafe {
			lock.raw_read();
			let guard = lock.read_guard();

			ContextGuard {
				_key: self.key,
				guard,
			}
		})
	}
}

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Sharable, O>
	LockingIterator<'c, Peekable<L>, O>
{
	/// Attempts to acquire a shared lock the next element and returning a guard
	/// to the inner data.
	///
	/// # Errors
	///
	/// Returns `Err(TryLockNextError::FinishedIteration)` when iteration is
	/// finished. Individual iterator implementations may choose to resume
	/// iteration, and so calling `next()` again may or may not eventually start
	/// returning `Some(Item)` again at some point.
	///
	/// Returns `Err(TryLockNextError::WouldBlock)` if the next lock in the
	/// iterator is already locked. This will not advance the iterator.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{RwLock, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	/// use happylock::context::iterator::TryLockNextError;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [RwLock::new(1), RwLock::new(3), RwLock::new(8)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key).peekable();
	///
	/// let mut sum = 0;
	/// loop {
	///     match iter.try_read_next() {
	///         Ok(item) => sum += **item,
	///         Err(TryLockNextError::WouldBlock) => continue,
	///         Err(TryLockNextError::FinishedIteration) => break,
	///     }
	/// }
	///
	/// assert_eq!(sum, 12);
	/// ```
	pub fn try_read_next(
		&mut self,
	) -> Result<ContextGuard<'c, <I as Sharable>::ReadGuard<'c>, ThreadKey>, TryLockNextError> {
		if let Some(lock) = self.iterator.peek().copied() {
			unsafe {
				if lock.raw_try_read() {
					// safety: we just saw that there is a valid value
					let lock = self.iterator.next().unwrap_unchecked();
					let guard = lock.read_guard();

					Ok(ContextGuard {
						_key: self.key,
						guard,
					})
				} else {
					Err(TryLockNextError::WouldBlock)
				}
			}
		} else {
			Err(TryLockNextError::FinishedIteration)
		}
	}
}

impl<'l, L: Iterator, O> LockingIterator<'l, L, O> {
	/// Advances the iterator, without locking the next element in the iterator.
	///
	/// Returns `false` when iteration is finished. Individual iterator
	/// implementations may choose to resume iteration, and so calling
	/// `skip_next()` again may or may not eventually start returning `true` again
	/// at some point.
	///
	/// # 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);
	///
	/// iter.skip_next();
	/// assert!(iter.lock_next().is_some_and(|v| **v == 3));
	/// ```
	pub fn skip_next(&mut self) -> bool {
		self.iterator.next().is_some()
	}

	/// Advances the iterator, skipping `n` elements without locking.
	///
	/// See [`Iterator::skip`] for more information.
	///
	/// # 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);
	///
	/// iter.skip_mut(2);
	/// assert!(iter.lock_next().is_some_and(|v| **v == 8));
	/// ```
	pub fn skip_mut(&mut self, n: usize) {
		for _ in 0..n {
			self.iterator.next();
		}
	}

	/// Returns the bounds on the remaining length of the iterator.
	///
	/// See [`Iterator::size_hint`] for more information.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [Mutex::new(1), Mutex::new(2), Mutex::new(3)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key);
	///
	/// assert_eq!((3, Some(3)), iter.size_hint());
	/// let _ = iter.skip_next();
	/// assert_eq!((2, Some(2)), iter.size_hint());
	/// ```
	#[must_use]
	pub fn size_hint(&self) -> (usize, Option<usize>) {
		self.iterator.size_hint()
	}

	/// Creates a new [`LockingIterator`] that skips the first `n` elements.
	///
	/// Unlike `skip_next` or `skip_mut`, this method does not modify the iterator
	/// in place. Instead, it returns a new iterator which skips the first `n`
	/// elements.
	///
	/// See [`Iterator::skip`] for more information.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [Mutex::new(1), Mutex::new(2), Mutex::new(3)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key).skip(2);
	///
	/// assert!(iter.lock_next().is_some_and(|v| **v == 3));
	/// assert!(iter.lock_next().is_none());
	/// ```
	#[must_use]
	pub fn skip(self, n: usize) -> LockingIterator<'l, Skip<L>, O> {
		self.with_iterator(|i| i.skip(n))
	}

	/// Creates a new [`LockingIterator`] that yields only the first `n` elements,
	/// or fewer if the iterator ends sooner.
	///
	/// See [`Iterator::take`] for more information.
	///
	/// # Example
	///
	/// ```
	/// use happylock::{Mutex, ThreadKey};
	/// use happylock::collection::OwnedLockCollection;
	///
	/// let key = ThreadKey::get().unwrap();
	/// let data = [Mutex::new(1), Mutex::new(2), Mutex::new(3)];
	/// let locks = OwnedLockCollection::new(data);
	/// let mut ctx = locks.context();
	/// let mut iter = ctx.iter(key).take(2);
	///
	/// assert!(iter.lock_next().is_some_and(|v| **v == 1));
	/// assert!(iter.lock_next().is_some_and(|v| **v == 2));
	/// assert!(iter.lock_next().is_none());
	/// ```
	#[must_use]
	pub fn take(self, n: usize) -> LockingIterator<'l, Take<L>, O> {
		self.with_iterator(|i| i.take(n))
	}

	/// Creates a new [`LockingIterator`] that ends after the first `None`
	///
	/// See [`Iterator::fuse`] for more information
	#[must_use]
	pub fn fuse(self) -> LockingIterator<'l, Fuse<L>, O> {
		self.with_iterator(Iterator::fuse)
	}

	/// Creates a new [`LockingIterator`] which has access to the
	/// [`try_lock_next`] and/or [`try_read_next`] methods.
	///
	/// See [`Iterator::peekable`] for more information
	///
	/// [`try_lock_next`]: `LockingIterator::try_lock_next`
	/// [`try_read_next`]: `LockingIterator::try_read_next`
	#[must_use]
	pub fn peekable(self) -> LockingIterator<'l, Peekable<L>, O> {
		self.with_iterator(Iterator::peekable)
	}
}