FlowScope/Packages/com.flowscope.gamecore/Runtime/Container/Container.cs

using System;
using System.Collections.Generic;
using System.Linq;

namespace FlowScope.Container
{
    public sealed class Container : IDisposable
    {
        [ThreadStatic]
        private static Stack<Type> _resolutionStack;

        private readonly Container _parent;
        private readonly Dictionary<Type, ContainerRegistration> _registrations = new Dictionary<Type, ContainerRegistration>();
        private readonly List<Container> _children = new List<Container>();
        private readonly List<IDisposable> _ownedDisposables = new List<IDisposable>();
        private bool _disposed;

        public Container()
        {
        }

        private Container(Container parent)
        {
            _parent = parent;
        }

        public void RegisterInstance<T>(T instance)
        {
            ThrowIfDisposed();

            if (instance == null)
            {
                throw new ArgumentNullException(nameof(instance));
            }

            RegisterLocal(typeof(T), ContainerRegistration.ForInstance(instance));
        }

        public void RegisterFactory<T>(Func<Container, T> factory)
        {
            RegisterSingletonFactory(factory);
        }

        public void RegisterSingletonFactory<T>(Func<Container, T> factory)
        {
            ThrowIfDisposed();

            if (factory == null)
            {
                throw new ArgumentNullException(nameof(factory));
            }

            RegisterLocal(typeof(T), ContainerRegistration.ForSingletonFactory(this, c => factory(c)));
        }

        public void RegisterScoped<T>(Func<Container, T> factory)
        {
            ThrowIfDisposed();

            if (factory == null)
            {
                throw new ArgumentNullException(nameof(factory));
            }

            RegisterLocal(typeof(T), ContainerRegistration.ForScopedFactory(this, c => factory(c)));
        }

        public void RegisterTransient<T>(Func<Container, T> factory)
        {
            ThrowIfDisposed();

            if (factory == null)
            {
                throw new ArgumentNullException(nameof(factory));
            }

            RegisterLocal(typeof(T), ContainerRegistration.ForTransientFactory(this, c => factory(c)));
        }

        public void RegisterType<TInterface, TImplementation>() where TImplementation : TInterface
        {
            ThrowIfDisposed();

            RegisterType(typeof(TInterface), typeof(TImplementation));
        }

        public void RegisterType<TImplementation>()
        {
            ThrowIfDisposed();

            RegisterType(typeof(TImplementation), typeof(TImplementation));
        }

        public void RegisterAssembly()
        {
            ThrowIfDisposed();

            GeneratedFactories.RegisterAssembly(this);
        }

        public T Resolve<T>()
        {
            ThrowIfDisposed();

            return (T)Resolve(typeof(T));
        }

        public bool TryResolve<T>(out T value)
        {
            ThrowIfDisposed();

            if (!TryFindRegistration(typeof(T), out var registration))
            {
                value = default;
                return false;
            }

            value = (T)ResolveRegistration(typeof(T), registration);
            return true;
        }

        public Container CreateScope()
        {
            ThrowIfDisposed();

            var child = new Container(this);
            _children.Add(child);
            return child;
        }

        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            _disposed = true;
            var exceptions = new List<Exception>();

            for (var i = _children.Count - 1; i >= 0; i--)
            {
                TryDispose(_children[i], exceptions);
            }

            for (var i = _ownedDisposables.Count - 1; i >= 0; i--)
            {
                TryDispose(_ownedDisposables[i], exceptions);
            }

            _children.Clear();
            _ownedDisposables.Clear();
            ForgetScopedInstancesFor(this);
            _registrations.Clear();

            if (_parent != null)
            {
                _parent._children.Remove(this);
            }

            if (exceptions.Count > 0)
            {
                throw new AggregateException("One or more container-owned instances failed to dispose.", exceptions);
            }
        }

        internal object Resolve(Type serviceType)
        {
            ThrowIfDisposed();

            if (!TryFindRegistration(serviceType, out var registration))
            {
                throw new InvalidOperationException($"No registration found for {FormatType(serviceType)}.");
            }

            return ResolveRegistration(serviceType, registration);
        }

        internal void RegisterFactory(Type serviceType, Func<Container, object> factory)
        {
            ThrowIfDisposed();

            if (serviceType == null)
            {
                throw new ArgumentNullException(nameof(serviceType));
            }

            if (factory == null)
            {
                throw new ArgumentNullException(nameof(factory));
            }

            RegisterLocal(serviceType, ContainerRegistration.ForTransientFactory(this, factory));
        }

        internal void TrackOwnedDisposable(IDisposable disposable)
        {
            _ownedDisposables.Add(disposable);
        }

        private void RegisterType(Type serviceType, Type implementationType)
        {
            if (!serviceType.IsAssignableFrom(implementationType))
            {
                throw new InvalidOperationException($"{FormatType(implementationType)} cannot be assigned to {FormatType(serviceType)}.");
            }

            if (!GeneratedFactories.TryGetFactory(implementationType, out var factory))
            {
                throw new InvalidOperationException(
                    $"No generated factory registered for {FormatType(implementationType)}. Use GeneratedFactories.Register, ReflectionFactoryBuilder.RegisterType, or RegisterFactory before resolving this type.");
            }

            RegisterFactory(serviceType, factory);
        }

        private void RegisterLocal(Type serviceType, ContainerRegistration registration)
        {
            if (_registrations.ContainsKey(serviceType))
            {
                throw new InvalidOperationException($"{FormatType(serviceType)} is already registered in this container scope.");
            }

            _registrations.Add(serviceType, registration);
        }

        private object ResolveRegistration(Type serviceType, ContainerRegistration registration)
        {
            _resolutionStack = _resolutionStack ?? new Stack<Type>();

            if (_resolutionStack.Contains(serviceType))
            {
                var chain = _resolutionStack.Reverse()
                    .Concat(new[] { serviceType })
                    .Select(t => t.Name);
                throw new InvalidOperationException($"Circular dependency detected: {string.Join(" -> ", chain)}");
            }

            _resolutionStack.Push(serviceType);
            try
            {
                return registration.Resolve(this);
            }
            finally
            {
                _resolutionStack.Pop();
                if (_resolutionStack.Count == 0)
                {
                    _resolutionStack = null;
                }
            }
        }

        private bool TryFindRegistration(Type serviceType, out ContainerRegistration registration)
        {
            if (_registrations.TryGetValue(serviceType, out registration))
            {
                return true;
            }

            if (_parent != null)
            {
                return _parent.TryFindRegistration(serviceType, out registration);
            }

            return false;
        }

        private void ForgetScopedInstancesFor(Container scope)
        {
            foreach (var registration in _registrations.Values)
            {
                registration.ForgetScope(scope);
            }

            _parent?.ForgetScopedInstancesFor(scope);
        }

        private void ThrowIfDisposed()
        {
            if (_disposed)
            {
                throw new ObjectDisposedException(nameof(Container));
            }
        }

        private static void TryDispose(IDisposable disposable, List<Exception> exceptions)
        {
            try
            {
                disposable.Dispose();
            }
            catch (Exception ex)
            {
                exceptions.Add(ex);
            }
        }

        private static string FormatType(Type type) => type.FullName ?? type.Name;
    }
}