Documentation
#[cfg(feature = "axum")]
pub mod axum;
pub mod error;
#[cfg(feature = "jav")]
pub mod jav;
pub mod scope;

use std::{
	any::{Any, TypeId},
	collections::HashMap,
	sync::{Arc, OnceLock, atomic::Ordering},
};

use funcext::FuncExt;
use saa::Semaphore;
use scopeguard::defer;
use sdd::{AtomicOwned, Guard, Owned, Tag};

use crate::{
	error::{ServiceAlreadyExistsError, ServiceNotFoundError},
	scope::{Scope, ScopeProvider, get_default_scope_full},
};

pub struct DI<T: ?Sized + Send + Sync + 'static>(pub Arc<T>);
pub struct DIK<T: ?Sized + Send + Sync + 'static, K: 'static>(
	pub Arc<T>,
	pub std::marker::PhantomData<K>,
);

#[derive(Clone)]
pub enum Service {
	Singleton(Arc<dyn Any + Send + Sync>),
	Transient(Arc<dyn Fn() -> Arc<dyn Any + Send + Sync> + Send + Sync>),
	Scoped(Arc<dyn Fn() -> Arc<dyn Any + Send + Sync> + Send + Sync>),
}

#[derive(Clone, Default)]
pub struct ServiceCollection(HashMap<(TypeId, Option<TypeId>), Service>);

impl ServiceCollection {
	pub fn new() -> Self { Self::default() }

	pub fn add_singleton<T: Send + Sync + 'static>(
		&mut self,
		service: T,
	) -> Result<(), ServiceAlreadyExistsError> {
		self.add_internal::<T>(Service::Singleton(Arc::new(Arc::new(service))), None)
	}

	pub fn add_singleton_keyed<T: Send + Sync + 'static>(
		&mut self,
		service: T,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError> {
		self.add_internal::<T>(Service::Singleton(Arc::new(Arc::new(service))), Some(key))
	}

	pub fn add_singleton_trait<Trait: ?Sized + Send + Sync + 'static>(
		&mut self,
		service: Arc<Trait>,
	) -> Result<(), ServiceAlreadyExistsError> {
		self.add_internal::<Trait>(Service::Singleton(Arc::new(service)), None)
	}

	pub fn add_singleton_trait_keyed<Trait: ?Sized + Send + Sync + 'static>(
		&mut self,
		service: Arc<Trait>,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError> {
		self.add_internal::<Trait>(Service::Singleton(Arc::new(service)), Some(key))
	}

	fn add_internal<T: ?Sized + Send + Sync + 'static>(
		&mut self,
		service: Service,
		key: Option<TypeId>,
	) -> Result<(), ServiceAlreadyExistsError> {
		let tid = (TypeId::of::<T>(), key);
		if self.0.contains_key(&tid) {
			return Err(ServiceAlreadyExistsError);
		}
		let _ = self.0.insert(tid, service);
		Ok(())
	}

	pub fn add_transient<T: Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> T + Send + Sync + 'static,
	{
		self.add_internal::<T>(
			Service::Transient(Arc::new(move || Arc::new(Arc::new(service_factory())))),
			None,
		)
	}
	pub fn add_transient_keyed<T: Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> T + Send + Sync + 'static,
	{
		self.add_internal::<T>(
			Service::Transient(Arc::new(move || Arc::new(Arc::new(service_factory())))),
			Some(key),
		)
	}

	pub fn add_transient_trait<Trait: ?Sized + Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> Arc<Trait> + Send + Sync + 'static,
	{
		self.add_internal::<Trait>(
			Service::Transient(Arc::new(move || Arc::new(service_factory()))),
			None,
		)
	}

	pub fn add_transient_trait_keyed<Trait: ?Sized + Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> Arc<Trait> + Send + Sync + 'static,
	{
		self.add_internal::<Trait>(
			Service::Transient(Arc::new(move || Arc::new(service_factory()))),
			Some(key),
		)
	}

	pub fn add_scoped<T: Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> T + Send + Sync + 'static,
	{
		self.add_internal::<T>(
			Service::Scoped(Arc::new(move || Arc::new(Arc::new(service_factory())))),
			None,
		)
	}
	pub fn add_scoped_keyed<T: Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> T + Send + Sync + 'static,
	{
		self.add_internal::<T>(
			Service::Scoped(Arc::new(move || Arc::new(Arc::new(service_factory())))),
			Some(key),
		)
	}

	pub fn add_scoped_trait<Trait: ?Sized + Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> Arc<Trait> + Send + Sync + 'static,
	{
		self.add_internal::<Trait>(
			Service::Scoped(Arc::new(move || Arc::new(service_factory()))),
			None,
		)
	}

	pub fn add_scoped_trait_keyed<Trait: ?Sized + Send + Sync + 'static, F>(
		&mut self,
		service_factory: F,
		key: TypeId,
	) -> Result<(), ServiceAlreadyExistsError>
	where
		F: Fn() -> Arc<Trait> + Send + Sync + 'static,
	{
		self.add_internal::<Trait>(
			Service::Scoped(Arc::new(move || Arc::new(service_factory()))),
			Some(key),
		)
	}
}

pub struct ServiceProvider {
	sc: AtomicOwned<ServiceCollection>,
	lock: Semaphore,
}

impl Default for ServiceProvider {
	fn default() -> Self {
		Self {
			sc: AtomicOwned::new(ServiceCollection::default()),
			lock: Semaphore::with_permits(1),
		}
	}
}

impl ServiceProvider {
	pub fn new(sc: ServiceCollection) -> Self {
		Self {
			sc: AtomicOwned::new(sc),
			lock: Semaphore::with_permits(1),
		}
	}

	fn mutate_internal<F>(&self, f: F)
	where
		F: FnOnce(&mut ServiceCollection),
	{
		defer!(self.lock.release(););
		let guard = Guard::new();
		let mut sc = self
			.sc
			.load(Ordering::Relaxed, &guard)
			.as_ref()
			.unwrap()
			.clone();
		f(&mut sc);
		self.sc.swap(
			(Some(Owned::new(sc)), self.sc.tag(Ordering::Relaxed)),
			Ordering::Relaxed,
		);
	}

	pub fn mutate<F>(&self, f: F)
	where
		F: FnOnce(&mut ServiceCollection),
	{
		self.lock.acquire_sync();
		self.mutate_internal(f);
	}

	pub async fn mutate_async<F>(&self, f: F)
	where
		F: FnOnce(&mut ServiceCollection),
	{
		self.lock.acquire_async().await;
		self.mutate_internal(f);
	}

	pub fn create_scope(&self) -> Arc<Scope> {
		let mut scope: Scope = HashMap::new();
		let guard = sdd::Guard::new();
		self.sc
			.load(Ordering::Relaxed, &guard)
			.as_ref()
			.unwrap()
			.0
			.iter()
			.for_each(|(k, v)| {
				let Service::Scoped(_) = v else {
					return;
				};
				let _ = scope.insert(*k, OnceLock::new());
			});
		Arc::new(scope)
	}

	fn di_internal<T: ?Sized + Send + Sync + 'static>(
		&self,
		key: Option<TypeId>,
		scope: Option<Arc<dyn ScopeProvider>>,
	) -> Result<Arc<T>, ServiceNotFoundError> {
		let tid = (TypeId::of::<T>(), key);
		let guard = sdd::Guard::new();
		let Some(s) = self
			.sc
			.load(Ordering::Relaxed, &guard)
			.as_ref()
			.unwrap()
			.0
			.get(&tid)
		else {
			return Err(ServiceNotFoundError);
		};
		match s {
			Service::Singleton(service) if let Ok(res) = service.clone().downcast::<Arc<T>>() => {
				Ok((*res).clone())
			},
			Service::Transient(factory) => factory.clone()().R(|i| {
				i.downcast::<Arc<T>>()
					.map(|m| (*m).clone())
					.map_err(|_| ServiceNotFoundError)
			}),
			Service::Scoped(factory) => scope
				.unwrap_or_else(|| get_default_scope_full())
				.provide_scope()?
				.get(&tid)
				.ok_or(ServiceNotFoundError)
				.map(|r| r.get_or_init(|| factory.clone()()).clone())
				.and_then(|i| {
					i.downcast::<Arc<T>>()
						.map(|m| (*m).clone())
						.map_err(|_| ServiceNotFoundError)
				}),
			_ => Err(ServiceNotFoundError),
		}
	}

	pub fn dir<T: ?Sized + Send + Sync + 'static>(&self) -> Result<Arc<T>, ServiceNotFoundError> {
		self.di_internal(None, None)
	}

	pub fn dikr<T: ?Sized + Send + Sync + 'static, K: 'static>(
		&self,
	) -> Result<Arc<T>, ServiceNotFoundError> {
		self.di_internal(Some(TypeId::of::<K>()), None)
	}

	pub fn disr<T: ?Sized + Send + Sync + 'static>(
		&self,
		scope: Arc<dyn ScopeProvider>,
	) -> Result<Arc<T>, ServiceNotFoundError> {
		self.di_internal(None, Some(scope))
	}
	pub fn diskr<T: ?Sized + Send + Sync + 'static, K: 'static>(
		&self,
		scope: Arc<dyn ScopeProvider>,
	) -> Result<Arc<T>, ServiceNotFoundError> {
		self.di_internal(Some(TypeId::of::<K>()), Some(scope))
	}

	pub fn di<T: ?Sized + Send + Sync + 'static>(&self) -> Arc<T> {
		self.di_internal(None, None)
			.expect("Failed to find service")
	}

	pub fn dik<T: ?Sized + Send + Sync + 'static, K: 'static>(&self) -> Arc<T> {
		self.di_internal(Some(TypeId::of::<K>()), None)
			.expect("Failed to find service")
	}

	pub fn dis<T: ?Sized + Send + Sync + 'static>(&self, scope: Arc<dyn ScopeProvider>) -> Arc<T> {
		self.di_internal(None, Some(scope))
			.expect("Failed to find service")
	}
	pub fn disk<T: ?Sized + Send + Sync + 'static, K: 'static>(
		&self,
		scope: Arc<dyn ScopeProvider>,
	) -> Arc<T> {
		self.di_internal(Some(TypeId::of::<K>()), Some(scope))
			.expect("Failed to find service")
	}
}

static SERVICE_PROVIDER: AtomicOwned<Arc<ServiceProvider>> = AtomicOwned::null();

pub fn spo() -> Option<Arc<ServiceProvider>> {
	let guard = Guard::new();
	SERVICE_PROVIDER
		.load(Ordering::SeqCst, &guard)
		.as_ref()
		.map(Arc::clone)
}

pub fn sp() -> Arc<ServiceProvider> { spo().expect("ServiceProvider not yet initialized") }

pub fn set_sp(sp: ServiceProvider) -> Option<Arc<ServiceProvider>> { set_sp_arc(Arc::new(sp)) }
pub fn set_sp_arc(sp: Arc<ServiceProvider>) -> Option<Arc<ServiceProvider>> {
	let (old, _) = SERVICE_PROVIDER.swap((Some(Owned::new(sp)), Tag::None), Ordering::SeqCst);
	old.map(|owned| (*owned).clone())
}

pub fn init_sp() { set_sp(ServiceProvider::default()); }