use std::iter::{Enumerate, Fuse, Skip, Take};

use super::{IteratorGuard, LockingIterator};

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

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

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Lockable> LockingIterator<'c, L> {
	pub fn lock_next(
		&mut self,
	) -> Option<IteratorGuard<'c, <I as Lockable>::Guard<'c>, ThreadKey>> {
		if let Some(lock) = self.iterator.next() {
			unsafe {
				lock.raw_write();
				let guard = lock.guard();

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

	pub fn lock_last(self) -> Option<IteratorGuard<'c, <I as Lockable>::Guard<'c>, ThreadKey>> {
		self.iterator.last().map(|lock| unsafe {
			lock.raw_write();
			let guard = lock.guard();

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

impl<'c, L: 'c + Iterator<Item = &'c I>, I: 'c + RawLock + Sharable> LockingIterator<'c, L> {
	pub fn read_next(
		&mut self,
	) -> Option<IteratorGuard<'c, <I as Sharable>::ReadGuard<'c>, ThreadKey>> {
		if let Some(lock) = self.iterator.next() {
			unsafe {
				lock.raw_read();
				let guard = lock.read_guard();

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

	pub fn read_last(self) -> Option<IteratorGuard<'c, <I as Sharable>::ReadGuard<'c>, ThreadKey>> {
		self.iterator.last().map(|lock| unsafe {
			lock.raw_read();
			let guard = lock.read_guard();

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

impl<'l, L: Iterator> LockingIterator<'l, L> {
	pub fn skip_next(&mut self) -> Option<L::Item> {
		self.iterator.next()
	}

	pub fn skip_mut(&mut self, n: usize) {
		for _ in 0..n {
			self.iterator.next();
		}
	}

	#[must_use]
	pub fn size_hint(&self) -> (usize, Option<usize>) {
		self.iterator.size_hint()
	}

	#[must_use]
	pub fn enumerate(self) -> LockingIterator<'l, Enumerate<L>> {
		self.with_iterator(Iterator::enumerate)
	}

	#[must_use]
	pub fn skip(self, n: usize) -> LockingIterator<'l, Skip<L>> {
		self.with_iterator(|i| i.skip(n))
	}

	#[must_use]
	pub fn take(self, n: usize) -> LockingIterator<'l, Take<L>> {
		self.with_iterator(|i| i.take(n))
	}

	#[must_use]
	pub fn fuse(self) -> LockingIterator<'l, Fuse<L>> {
		self.with_iterator(Iterator::fuse)
	}
}