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

using System;
using System.Collections.Generic;

namespace FlowScope.Container
{
    internal sealed class ContainerRegistration
    {
        private enum Lifetime
        {
            Instance,
            Singleton,
            Scoped,
            Transient
        }

        private readonly Container _owner;
        private readonly Func<Container, object> _factory;
        private readonly Lifetime _lifetime;
        private readonly Dictionary<Container, object> _scopedInstances;
        private object _instance;
        private bool _created;

        private ContainerRegistration(
            Container owner,
            Func<Container, object> factory,
            object instance,
            bool created,
            Lifetime lifetime)
        {
            _owner = owner;
            _factory = factory;
            _instance = instance;
            _created = created;
            _lifetime = lifetime;
            _scopedInstances = lifetime == Lifetime.Scoped
                ? new Dictionary<Container, object>()
                : null;
        }

        public static ContainerRegistration ForInstance(object instance)
        {
            return new ContainerRegistration(null, null, instance, true, Lifetime.Instance);
        }

        public static ContainerRegistration ForSingletonFactory(Container owner, Func<Container, object> factory)
        {
            return new ContainerRegistration(owner, factory, null, false, Lifetime.Singleton);
        }

        public static ContainerRegistration ForScopedFactory(Container owner, Func<Container, object> factory)
        {
            return new ContainerRegistration(owner, factory, null, false, Lifetime.Scoped);
        }

        public static ContainerRegistration ForTransientFactory(Container owner, Func<Container, object> factory)
        {
            return new ContainerRegistration(owner, factory, null, false, Lifetime.Transient);
        }

        public object Resolve(Container requestScope)
        {
            switch (_lifetime)
            {
                case Lifetime.Instance:
                    return _instance;
                case Lifetime.Singleton:
                    return ResolveSingleton();
                case Lifetime.Scoped:
                    return ResolveScoped(requestScope);
                case Lifetime.Transient:
                    return CreateOwnedInstance(requestScope);
                default:
                    throw new InvalidOperationException($"Unsupported container lifetime {_lifetime}.");
            }
        }

        public void ForgetScope(Container scope)
        {
            _scopedInstances?.Remove(scope);
        }

        private object ResolveSingleton()
        {
            if (!_created)
            {
                _instance = CreateOwnedInstance(_owner);
                _created = true;
            }

            return _instance;
        }

        private object ResolveScoped(Container requestScope)
        {
            if (_scopedInstances.TryGetValue(requestScope, out var instance))
            {
                return instance;
            }

            instance = CreateOwnedInstance(requestScope);
            _scopedInstances.Add(requestScope, instance);
            return instance;
        }

        private object CreateOwnedInstance(Container owner)
        {
            var instance = _factory(owner);

            if (instance is IDisposable disposable)
            {
                owner.TrackOwnedDisposable(disposable);
            }

            return instance;
        }
    }
}