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

use crate::error::ServiceNotFoundError;
use arc_swap::{ArcSwap, Guard};

pub type Scope = HashMap<(TypeId, Option<TypeId>), std::sync::OnceLock<Arc<dyn Any + Send + Sync>>>;

pub trait ScopeDefault: ScopeProvider + Sized + 'static {
	fn initialize(scope: Arc<Scope>);
	fn update(scope: Arc<Scope>);
	fn set_as_default(self) {
		DEFAULT_SCOPE_PROVIDER.store(Arc::new(Arc::new(self)));
	}
}

pub trait ScopeProvider: Send + Sync {
	fn provide_scope(&self) -> Result<Arc<Scope>, ServiceNotFoundError>;
}

static DEFAULT_SCOPE_PROVIDER: LazyLock<ArcSwap<Arc<dyn ScopeProvider>>> =
	LazyLock::new(|| ArcSwap::new(Arc::new(Arc::new(GlobalScopeProvider))));

pub fn get_default_scope() -> Guard<Arc<Arc<dyn ScopeProvider>>> {
	DEFAULT_SCOPE_PROVIDER.load()
}

pub fn get_default_scope_full() -> Arc<Arc<dyn ScopeProvider>> {
	DEFAULT_SCOPE_PROVIDER.load_full()
}

pub fn set_default_scope<T: ScopeDefault + 'static>(scope: T) {
	DEFAULT_SCOPE_PROVIDER.store(Arc::new(Arc::new(scope)));
}

pub struct GlobalScopeProvider;

impl ScopeDefault for GlobalScopeProvider {
	fn initialize(scope: Arc<Scope>) {
		let _ = GLOBAL_SCOPE.set(ArcSwap::new(scope));
	}

	fn update(new_scope: Arc<Scope>) {
		let Some(scope) = GLOBAL_SCOPE.get() else {
			return;
		};
		scope.swap(new_scope);
	}
}

impl ScopeProvider for GlobalScopeProvider {
	fn provide_scope(&self) -> Result<Arc<Scope>, ServiceNotFoundError> {
		let Some(scope) = GLOBAL_SCOPE.get() else {
			return Err(ServiceNotFoundError);
		};
		Ok(scope.load_full().clone())
	}
}

static GLOBAL_SCOPE: OnceLock<ArcSwap<Scope>> = OnceLock::new();

pub struct ThreadLocalScope;

impl ScopeDefault for ThreadLocalScope {
	fn initialize(scope: Arc<Scope>) {
		let _ = TL_SCOPE.with(|s| s.set(ArcSwap::new(scope)));
	}

	fn update(new_scope: Arc<Scope>) {
		TL_SCOPE.with(|s| {
			let Some(scope) = s.get() else {
				return;
			};
			scope.swap(new_scope);
		});
	}
}

impl ScopeProvider for ThreadLocalScope {
	fn provide_scope(&self) -> Result<Arc<Scope>, ServiceNotFoundError> {
		TL_SCOPE.with(|s| {
			let Some(scope) = s.get() else {
				return Err(ServiceNotFoundError);
			};
			Ok(scope.load_full().clone())
		})
	}
}

thread_local! {
	static TL_SCOPE: OnceLock<ArcSwap<Scope>> = OnceLock::new();
}

#[cfg(feature = "tokio")]
pub mod tasks {
	use std::{intrinsics::unreachable, sync::Arc};

	use tokio::task::futures::TaskLocalFuture;

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

	tokio::task_local! {
		static TASK_SCOPE: Arc<Scope>;
	}

	pub struct TaskScopeProvider;

	impl super::ScopeDefault for TaskScopeProvider {
		fn initialize(_scope: Arc<Scope>) {
			unreachable!()
		}

		fn update(_scope: Arc<Scope>) {
			unreachable!()
		}
	}

	impl TaskScopeProvider {
		pub fn initialize_async<F>(scope: Scope, fut: F) -> TaskLocalFuture<Arc<Scope>, F>
		where
			F: Future,
		{
			TASK_SCOPE.scope(Arc::new(scope), fut)
		}

		pub fn initialize_sync<F, R>(scope: Scope, fun: F) -> R
		where
			F: FnOnce() -> R,
		{
			TASK_SCOPE.sync_scope(Arc::new(scope), fun)
		}

		pub fn update_async<F>(scope: Scope, fut: F) -> TaskLocalFuture<Arc<Scope>, F>
		where
			F: Future,
		{
			TASK_SCOPE.scope(Arc::new(scope), fut)
		}

		pub fn update_sync<F, R>(scope: Scope, fun: F) -> R
		where
			F: FnOnce() -> R,
		{
			TASK_SCOPE.sync_scope(Arc::new(scope), fun)
		}
	}

	impl ScopeProvider for TaskScopeProvider {
		fn provide_scope(&self) -> Result<Arc<Scope>, ServiceNotFoundError> {
			Ok(TASK_SCOPE.get())
		}
	}
}

#[cfg(feature = "tokio")]
pub use tasks::*;
