mosaik 0.4.2

use {
	super::CandidatesMap,
	crate::{
		discovery::PeerEntry,
		primitives::{IntoIterOrSingle, Tag},
	},
	core::{
		fmt,
		pin::Pin,
		task::{Context, Poll},
	},
	futures::FutureExt,
	std::{collections::BTreeSet, sync::Arc},
	tokio::sync::watch,
	tokio_util::sync::ReusableBoxFuture,
};

/// Awaits changes to the caller's handler availability status.
///
/// This struct provides access to futures that can be used to await when
/// the caller becomes online (eligible handlers are available) or meets
/// other availability conditions.
pub struct When {
	/// Observer for the online status of the caller.
	///
	/// When the value is set to false there are not enough eligible
	/// handlers available to satisfy the caller's online conditions.
	pub(in crate::functions) online: watch::Receiver<bool>,

	/// Observer for the most recent snapshot of eligible handler peers.
	pub(in crate::functions) candidates: watch::Receiver<CandidatesMap>,
}

impl Clone for When {
	fn clone(&self) -> Self {
		Self::new(self.candidates.clone(), self.online.clone())
	}
}

impl When {
	/// Initialized by callers for each new caller instance.
	pub(in crate::functions) fn new(
		mut candidates: watch::Receiver<CandidatesMap>,
		mut online: watch::Receiver<bool>,
	) -> Self {
		candidates.mark_changed();
		online.mark_changed();
		Self { online, candidates }
	}
}

// Public API
impl When {
	/// Returns a future that resolves when the caller is ready to invoke
	/// the function — its configured online conditions are met (by default,
	/// at least one eligible handler is known).
	///
	/// Resolves immediately if the caller is already online.
	pub fn online(&self) -> impl Future<Output = ()> + Send + Sync + 'static {
		let mut online = self.online.clone();

		async move {
			if online.wait_for(|v| *v).await.is_err() {
				// if the watch channel is closed, consider the caller offline and
				// never resolve this future
				core::future::pending::<()>().await;
			}
		}
	}

	/// Returns whether the caller's online conditions are currently met.
	pub fn is_online(&self) -> bool {
		*self.online.borrow()
	}

	/// Returns a future that resolves when the caller is not ready to
	/// invoke the function.
	///
	/// Resolves immediately if the caller is already offline.
	pub fn offline(&self) -> impl Future<Output = ()> + Send + Sync + 'static {
		let mut online = self.online.clone();

		async move {
			// if the watch channel is closed, consider the caller offline and
			// always resolve this future
			let _ = online.wait_for(|v| !*v).await;
		}
	}

	/// Returns a future that resolves when at least one eligible handler is
	/// available for the function. This can be customized and combined with
	/// other conditions using the methods on the returned
	/// [`HandlerConditions`].
	pub fn available(&self) -> HandlerConditions {
		let mut candidates = self.candidates.clone();
		candidates.mark_changed();

		HandlerConditions {
			candidates: candidates.clone(),
			min_peers: 1,
			was_met: false,
			is_inverse: false,
			predicates: Vec::new(),
			changed_fut: ReusableBoxFuture::new(Box::pin(async move {
				let _ = candidates.changed().await;
			})),
		}
	}

	/// Returns a future that resolves when no eligible handlers are
	/// available for the function. This can be customized and combined with
	/// other conditions using the methods on the returned
	/// [`HandlerConditions`].
	///
	/// This is equivalent to calling `available().unmet()`.
	pub fn unavailable(&self) -> HandlerConditions {
		self.available().unmet()
	}
}

/// A future that resolves when the set of eligible handlers for a function
/// meets a certain condition.
///
/// This future can be polled multiple times even after it has resolved
/// once, and it will resolve again when the awaited condition transitions
/// again from not met to met.
///
/// In its initial state when instantiated and the condition is met
/// immediately, the future will resolve on the next poll, then reset to
/// awaiting state until the condition transitions from not met to met.
pub struct HandlerConditions {
	candidates: watch::Receiver<CandidatesMap>,
	min_peers: usize,
	predicates: Vec<Arc<PeerPredicate>>,
	was_met: bool,
	is_inverse: bool,
	changed_fut: ReusableBoxFuture<'static, ()>,
}

// Public API
impl HandlerConditions {
	/// Specifies that the future should resolve when there is at least the
	/// given number of eligible handlers.
	#[must_use]
	pub const fn minimum_of(mut self, min: usize) -> Self {
		self.min_peers = min;
		self
	}

	/// Specifies that the future should resolve when eligible handlers
	/// contain the given tags in their
	/// [`PeerEntry`](crate::discovery::PeerEntry).
	///
	/// When combined with `minimum_of`, the condition is met when there are
	/// at least that many handlers with the given tags.
	#[must_use]
	pub fn with_tags<V>(self, tags: impl IntoIterOrSingle<Tag, V>) -> Self {
		let tags: BTreeSet<Tag> = tags.iterator().into_iter().collect();
		self.with_predicate(move |peer: &PeerEntry| tags.is_subset(peer.tags()))
	}

	/// Specifies a custom predicate that must be met by eligible handlers
	/// for the condition to be considered met.
	#[must_use]
	pub fn with_predicate<F>(mut self, predicate: F) -> Self
	where
		F: Fn(&PeerEntry) -> bool + Send + Sync + 'static,
	{
		self.predicates.push(Arc::new(predicate));
		self
	}

	/// Checks the minimum number of eligible handlers that meet our
	/// predicates.
	pub fn is_condition_met(&self) -> bool {
		let matching_peers = self
			.candidates
			.borrow()
			.values()
			.filter(|candidate| {
				self.predicates.iter().all(|pred| pred(candidate.entry()))
			})
			.count();

		(matching_peers >= self.min_peers) != self.is_inverse
	}

	/// Inverts the condition, so that the future resolves when the
	/// condition is not met.
	#[must_use]
	pub const fn unmet(self) -> Self {
		let mut cloned = self;
		cloned.is_inverse = true;
		cloned
	}
}

impl HandlerConditions {
	/// The same check as [`is_condition_met`](Self::is_condition_met), but
	/// marks the evaluated snapshot as seen.
	///
	/// `borrow_and_update` reads the freshest snapshot and marks exactly
	/// that version seen in one step, so a `changed()` future armed from
	/// this receiver afterwards waits for a genuinely newer snapshot. A
	/// plain `borrow` leaves the receiver's version frozen, and every
	/// re-armed future then resolves instantly and spins the poll loop.
	fn evaluate(&mut self) -> bool {
		let Self {
			candidates,
			predicates,
			min_peers,
			is_inverse,
			..
		} = self;

		let matching_peers = candidates
			.borrow_and_update()
			.values()
			.filter(|candidate| predicates.iter().all(|pred| pred(candidate.entry())))
			.count();

		(matching_peers >= *min_peers) != *is_inverse
	}
}

impl Clone for HandlerConditions {
	fn clone(&self) -> Self {
		let mut candidates = self.candidates.clone();
		candidates.mark_changed();
		Self {
			candidates: candidates.clone(),
			min_peers: self.min_peers,
			predicates: self.predicates.clone(),
			was_met: false,
			is_inverse: self.is_inverse,
			changed_fut: ReusableBoxFuture::new(Box::pin(async move {
				let _ = candidates.changed().await;
			})),
		}
	}
}

impl fmt::Debug for HandlerConditions {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_struct("HandlerConditions")
			.field("min_peers", &self.min_peers)
			.field("predicates", &self.predicates.len())
			.field("is_condition_met", &self.is_condition_met())
			.finish_non_exhaustive()
	}
}

impl Future for HandlerConditions {
	type Output = ();

	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
		let this = self.get_mut();

		loop {
			// First check if the condition is currently met, marking the
			// snapshot seen so the future armed below waits for a newer one
			let condition_met = this.evaluate();

			if condition_met && !this.was_met {
				// Transition from not met -> met (or initially met)
				this.was_met = true;
				return Poll::Ready(());
			}

			// Update state tracking
			this.was_met = condition_met;

			// Poll the stored changed future to wait for updates
			match this.changed_fut.poll_unpin(cx) {
				Poll::Ready(()) => {
					// The watch was updated, set up a new changed future and loop
					let mut receiver = this.candidates.clone();
					this.changed_fut.set(Box::pin(async move {
						let _ = receiver.changed().await;
					}));
				}
				Poll::Pending => return Poll::Pending,
			}
		}
	}
}

impl PartialEq<bool> for HandlerConditions {
	fn eq(&self, other: &bool) -> bool {
		self.is_condition_met() == *other
	}
}

impl PartialEq<HandlerConditions> for bool {
	fn eq(&self, other: &HandlerConditions) -> bool {
		*self == other.is_condition_met()
	}
}

type PeerPredicate = dyn Fn(&PeerEntry) -> bool + Send + Sync;