mosaik 0.4.2

use {
	super::{
		super::{
			CallContext,
			Function,
			FunctionId,
			Functions,
			accept::CallReply,
			derived_function_id,
			effective_function_id,
		},
		Handler,
		Stats,
		registry::{DispatchError, Entry},
	},
	crate::{
		discovery::PeerInfo,
		primitives::{Bytes, Datum},
		tickets::TicketValidator,
	},
	core::marker::PhantomData,
	futures::FutureExt,
	std::sync::Arc,
	tokio::sync::Semaphore,
};

#[derive(Debug, thiserror::Error)]
pub enum BuilderError {
	/// A handler for the given function id already exists.
	///
	/// Each function id can only be served by one handler per node. Unlike
	/// stream producers, the existing handler is not returned — it is bound
	/// to the closure or [`Function`] instance it was registered with.
	#[error("A handler for this function id is already registered")]
	AlreadyExists,
}

/// Configuration options for a function handler.
pub struct HandlerConfig {
	/// The base function id this handler serves. When unset, it is derived
	/// from the request/response types (closure registration) or the
	/// [`Function`] implementation's signature.
	pub function_id: Option<FunctionId>,

	/// Predicate determining whether a caller peer is allowed to invoke
	/// this function.
	pub require: Box<dyn Fn(&PeerInfo) -> bool + Send + Sync>,

	/// Ticket validators that callers must satisfy (enforced by this
	/// handler). Their signatures fold into the effective wire id.
	pub caller_auth: Vec<Arc<dyn TicketValidator>>,

	/// Ticket validators that handlers must satisfy (enforced by callers).
	/// Not enforced locally — declared only so their signatures fold into
	/// the effective wire id, keeping it in sync with callers.
	pub handler_auth: Vec<Arc<dyn TicketValidator>>,

	/// Maximum number of concurrently executing calls.
	pub max_concurrent: usize,
}

/// Configurable builder for registering a function handler for a specific
/// request type `Req`, response type `Res` and application error type `E`.
pub struct Builder<'f, Req: Datum, Res: Datum, E: Datum = ()> {
	config: HandlerConfig,
	functions: &'f Functions,
	_marker: PhantomData<fn(&Req, &Res, &E)>,
}

/// Public API
impl<Req: Datum, Res: Datum, E: Datum> Builder<'_, Req, Res, E> {
	/// Sets the base function id this handler serves.
	///
	/// If not set, the base id is derived from the request/response type
	/// names (closure registration) or from
	/// [`Function::signature`] (trait registration). The signatures of any
	/// configured ticket validators are folded on top of the base id to
	/// produce the effective wire id.
	#[must_use]
	pub fn with_function_id(
		mut self,
		function_id: impl Into<FunctionId>,
	) -> Self {
		self.config.function_id = Some(function_id.into());
		self
	}

	/// Adds a peer eligibility requirement for incoming call connections.
	///
	/// The predicate receives a [`PeerInfo`] combining the caller's
	/// self-reported [`PeerEntry`](crate::discovery::PeerEntry) with
	/// locally-observed metrics like RTT.
	///
	/// When called multiple times, all predicates must pass — AND
	/// composition — so a call is accepted only if every requirement is
	/// satisfied.
	///
	/// The default requirement accepts all callers.
	#[must_use]
	pub fn require<F>(mut self, pred: F) -> Self
	where
		F: Fn(&PeerInfo) -> bool + Send + Sync + 'static,
	{
		let prev = self.config.require;
		self.config.require = Box::new(move |peer| prev(peer) && pred(peer));
		self
	}

	/// Adds a ticket validator that authorizes callers of this function.
	///
	/// This is counterparty sugar for
	/// [`require_caller_ticket`](Self::require_caller_ticket).
	#[must_use]
	pub fn require_ticket(self, validator: impl TicketValidator) -> Self {
		self.require_caller_ticket(validator)
	}

	/// Adds a ticket validator that callers of this function must satisfy.
	///
	/// Each caller attempting to invoke the function must present valid
	/// tickets that pass all configured caller validators. The validator's
	/// signature is folded into the effective wire id, so functions with
	/// different authorization rules have different identities on the
	/// network — callers must declare the same contract to find this
	/// handler. Can be called multiple times to require multiple types of
	/// tickets.
	#[must_use]
	pub fn require_caller_ticket(
		mut self,
		validator: impl TicketValidator,
	) -> Self {
		self.config.caller_auth.push(Arc::new(validator));
		self
	}

	/// Declares a ticket validator that handlers of this function must
	/// satisfy.
	///
	/// Handler tickets are enforced by callers, not by this handler — the
	/// declaration only folds the validator's signature into the effective
	/// wire id so it matches the id derived by callers enforcing the same
	/// contract. Can be called multiple times.
	#[must_use]
	pub fn require_handler_ticket(
		mut self,
		validator: impl TicketValidator,
	) -> Self {
		self.config.handler_auth.push(Arc::new(validator));
		self
	}

	/// Sets the maximum number of concurrently executing calls for this
	/// function. Calls arriving beyond this limit are rejected with a
	/// `NoCapacity` close reason and the caller fails over to another
	/// handler if one is available.
	///
	/// Defaults to unlimited if not set.
	#[must_use]
	pub const fn with_max_concurrent(mut self, max: usize) -> Self {
		self.config.max_concurrent = max;
		self
	}

	/// Registers the given closure as the function's implementation and
	/// advertises the function on the network.
	///
	/// The returned [`Handler`] owns the registration: dropping it
	/// deregisters the function, un-advertises it, and aborts in-flight
	/// calls.
	pub fn serve<F, Fut>(
		mut self,
		f: F,
	) -> Result<Handler<Req, Res, E>, BuilderError>
	where
		F: Fn(Req, CallContext) -> Fut + Send + Sync + 'static,
		Fut: Future<Output = Result<Res, E>> + Send + 'static,
	{
		// Closures have no stable, side-shared type name, so the default base
		// id is derived from the request/response types only.
		if self.config.function_id.is_none() {
			self.config.function_id = Some(derived_function_id::<Req, Res>());
		}

		self.serve_fn(FnAdapter {
			f,
			_marker: PhantomData,
		})
	}

	/// Registers a ready [`Function`] instance as the function's
	/// implementation and advertises the function on the network.
	///
	/// Any warmup or initialization belongs in the instance's own
	/// construction, before this call — the function only becomes
	/// discoverable once registered.
	///
	/// The returned [`Handler`] owns the registration: dropping it
	/// deregisters the function, un-advertises it, and aborts in-flight
	/// calls.
	pub fn serve_fn<F>(
		self,
		function: F,
	) -> Result<Handler<Req, Res, E>, BuilderError>
	where
		F: Function<Req = Req, Res = Res, Err = E>,
	{
		let base = self.config.function_id.unwrap_or_else(F::signature);
		let function_id = effective_function_id(
			base,
			&self.config.caller_auth,
			&self.config.handler_auth,
		);

		let function = Arc::new(function);
		let dispatch = Arc::new(move |payload: Bytes, ctx: CallContext| {
			Req::decode(&payload).map_or_else(
				|_| futures::future::ready(Err(DispatchError::DecodeRequest)).boxed(),
				|req| {
					let function = Arc::clone(&function);
					async move {
						match function.call(req, ctx).await {
							Ok(res) => res
								.encode()
								.map(CallReply::Ok)
								.map_err(|_| DispatchError::EncodeReply),
							Err(err) => err
								.encode()
								.map(CallReply::Err)
								.map_err(|_| DispatchError::EncodeReply),
						}
					}
					.boxed()
				},
			)
		});

		let registry = Arc::clone(&self.functions.registry);
		let cancel = self.functions.local.termination().child_token();
		let stats = Arc::new(Stats::default());
		let permits = Arc::new(Semaphore::new(
			self.config.max_concurrent.min(Semaphore::MAX_PERMITS),
		));

		registry
			.create(function_id, Entry {
				dispatch,
				require: self.config.require,
				caller_auth: self.config.caller_auth,
				permits,
				cancel: cancel.clone(),
				stats: Arc::clone(&stats),
			})
			.map_err(|()| BuilderError::AlreadyExists)?;

		Ok(Handler::new(function_id, stats, registry, cancel))
	}
}

impl<'f, Req: Datum, Res: Datum, E: Datum> Builder<'f, Req, Res, E> {
	pub(in crate::functions) fn new(functions: &'f Functions) -> Self {
		Self {
			functions,
			config: HandlerConfig {
				function_id: None,
				require: Box::new(|_| true),
				caller_auth: Vec::new(),
				handler_auth: Vec::new(),
				max_concurrent: usize::MAX,
			},
			_marker: PhantomData,
		}
	}
}

/// Adapts a plain closure into a [`Function`] implementation with no
/// startup logic.
struct FnAdapter<F, Req, Res, E> {
	f: F,
	_marker: PhantomData<fn(&Req, &Res, &E)>,
}

impl<F, Fut, Req, Res, E> Function for FnAdapter<F, Req, Res, E>
where
	F: Fn(Req, CallContext) -> Fut + Send + Sync + 'static,
	Fut: Future<Output = Result<Res, E>> + Send + 'static,
	Req: Datum,
	Res: Datum,
	E: Datum,
{
	type Err = E;
	type Req = Req;
	type Res = Res;

	fn call(
		&self,
		req: Req,
		ctx: CallContext,
	) -> impl Future<Output = Result<Res, E>> + Send {
		(self.f)(req, ctx)
	}
}