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

use super::{ConsumedIteratorGuard, IteratorGuard, LockingIterator};

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

impl<L, Key> LockingIterator<L, Key> {
	pub(crate) const fn new(key: Key, lockable: L) -> Self {
		Self { key, lockable }
	}

	fn with_iterator<M>(self, f: impl FnOnce(L) -> M) -> LockingIterator<M, Key> {
		LockingIterator {
			key: self.key,
			lockable: f(self.lockable),
		}
	}
}

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

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

	pub fn lock_last(self) -> Option<ConsumedIteratorGuard<<L as Lockable>::Guard<'a>, Key>> {
		self.lockable.last().map(|lock| unsafe {
			lock.raw_write();
			let guard = lock.guard();

			ConsumedIteratorGuard {
				key: self.key,
				guard,
			}
		})
	}
}

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

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

	pub fn read_last(self) -> Option<ConsumedIteratorGuard<<L as Sharable>::ReadGuard<'a>, Key>> {
		self.lockable.last().map(|lock| unsafe {
			lock.raw_read();
			let guard = lock.read_guard();

			ConsumedIteratorGuard {
				key: self.key,
				guard,
			}
		})
	}
}

impl<L: Iterator, Key> LockingIterator<L, Key> {
	pub fn size_hint(&self) -> (usize, Option<usize>) {
		self.lockable.size_hint()
	}

	pub fn enumerate(self) -> LockingIterator<Enumerate<L>, Key> {
		self.with_iterator(Iterator::enumerate)
	}

	pub fn skip(self, n: usize) -> LockingIterator<Skip<L>, Key> {
		self.with_iterator(|i| i.skip(n))
	}

	pub fn take(self, n: usize) -> LockingIterator<Take<L>, Key> {
		self.with_iterator(|i| i.take(n))
	}

	pub fn fuse(self) -> LockingIterator<Fuse<L>, Key> {
		self.with_iterator(Iterator::fuse)
	}
}