SPT-Server-Build/Libraries/SPTarkov.DI/DependencyInjectionHandler.cs

using System.Reflection;
using SPTarkov.DI.Annotations;

namespace SPTarkov.DI;

public class DependencyInjectionHandler(IServiceCollection serviceCollection)
{
    private static List<Type>? _allLoadedTypes;
    private static List<ConstructorInfo>? _allConstructors;

    private readonly Dictionary<string, Type> _injectedTypeNames = new();

    private readonly Dictionary<string, object> _injectedValues = new();
    private readonly Lock _injectedValuesLock = new();

    private bool _oneTimeUseFlag;

    public void AddInjectableTypesFromAssembly(Assembly assembly)
    {
        AddInjectableTypesFromTypeList(assembly.GetTypes());
    }

    public void AddInjectableTypesFromAssemblies(IEnumerable<Assembly> assemblies)
    {
        foreach (var assembly in assemblies)
        {
            AddInjectableTypesFromAssembly(assembly);
        }
    }

    public void AddInjectableTypesFromTypeAssembly(Type type)
    {
        AddInjectableTypesFromAssembly(type.Assembly);
    }

    public void AddInjectableTypesFromTypeList(IEnumerable<Type> types)
    {
        var typesToInject = types.Where(type =>
            Attribute.IsDefined(type, typeof(Injectable))
            && !_injectedTypeNames.ContainsKey($"{type.Namespace}.{type.Name}")
        );
        if (typesToInject.Any())
        {
            foreach (var type in typesToInject)
            {
                _injectedTypeNames.Add($"{type.Namespace}.{type.Name}", type);
            }
        }
    }

    public void InjectAll()
    {
        if (_oneTimeUseFlag)
        {
            throw new Exception(
                "Invalid usage of DependencyInjectionHandler, this is a one time use service!"
            );
        }
        _oneTimeUseFlag = true;
        var typeRefValues = _injectedTypeNames.Values.Select(t => new TypeRefContainer(
            ((Injectable[])Attribute.GetCustomAttributes(t, typeof(Injectable)))[0],
            t,
            t
        ));
        // All the components that have a type override, we need to find them and remove them before injecting everything
        var componentsToRemove = typeRefValues
            .Where(tr => tr.InjectableAttribute.TypeOverride != null)
            .Select(tr =>
                string.IsNullOrEmpty(tr.InjectableAttribute.TypeOverride!.FullName)
                    ? $"{tr.InjectableAttribute.TypeOverride.Namespace}.{tr.InjectableAttribute.TypeOverride.Name}"
                    : tr.InjectableAttribute.TypeOverride.FullName!
            )
            .ToHashSet();
        // All the components without the removed overrides
        var cleanedComponents = typeRefValues.Where(tr =>
        {
            var name = string.IsNullOrEmpty(tr.Type.FullName)
                ? $"{tr.Type.Namespace}.{tr.Type.Name}"
                : tr.Type.FullName!;
            return !componentsToRemove.Contains(name);
        });
        // All the components sorted and ready to be inserted into the DI container
        var sortedInjectableTypes = cleanedComponents.OrderBy(tRef =>
            tRef.InjectableAttribute.TypePriority
        );

        foreach (var typeRefToInject in sortedInjectableTypes)
        {
            var nodes = new Queue<TypeRefContainer>();
            nodes.Enqueue(typeRefToInject);
            foreach (
                var implementedInterface in typeRefToInject
                    .Type.GetInterfaces()
                    .Where(t => !t.Namespace.StartsWith("System"))
            )
            {
                nodes.Enqueue(
                    new TypeRefContainer(
                        typeRefToInject.InjectableAttribute,
                        typeRefToInject.Type,
                        implementedInterface
                    )
                );
            }

            while (nodes.Any())
            {
                var node = nodes.Dequeue();
                if (node.Type.BaseType != null && node.Type.BaseType != typeof(object))
                {
                    nodes.Enqueue(
                        new TypeRefContainer(
                            node.InjectableAttribute,
                            typeRefToInject.Type,
                            node.Type.BaseType
                        )
                    );
                }

                if (node.Type.IsGenericType)
                {
                    RegisterGenericComponents(node);
                }
                else
                {
                    RegisterComponent(
                        node.InjectableAttribute.InjectionType,
                        node.Type,
                        node.ParentType
                    );
                }
            }
        }
    }

    private void RegisterGenericComponents(TypeRefContainer typeRef)
    {
        try
        {
            _allLoadedTypes ??= AppDomain
                .CurrentDomain.GetAssemblies()
                .SelectMany(t => t.GetTypes())
                .ToList();
        }
        catch (ReflectionTypeLoadException ex)
        {
            Console.WriteLine($"COULD NOT LOAD TYPE: {ex}");
        }

        _allConstructors ??= _allLoadedTypes.SelectMany(t => t.GetConstructors()).ToList();

        var typeName = $"{typeRef.Type.Namespace}.{typeRef.Type.Name}";
        try
        {
            var matchedConstructors = _allConstructors.Where(c =>
                c.GetParameters()
                    .Any(p =>
                        p.ParameterType.IsGenericType
                        && p.ParameterType.GetGenericTypeDefinition().FullName == typeName
                    )
            );

            var constructorInfos = matchedConstructors.ToList();
            if (constructorInfos.Count == 0)
            {
                return;
            }

            foreach (var matchedConstructor in constructorInfos)
            {
                var constructorParams = matchedConstructor.GetParameters();
                foreach (
                    var parameterInfo in constructorParams.Where(x =>
                        IsMatchingGenericType(x, typeName)
                    )
                )
                {
                    var parameters = parameterInfo.ParameterType.GetGenericArguments();
                    var typedGeneric = typeRef.ParentType.MakeGenericType(parameters);
                    RegisterComponent(
                        typeRef.InjectableAttribute.InjectionType,
                        parameterInfo.ParameterType,
                        typedGeneric
                    );
                }
            }
        }
        catch (Exception e)
        {
            Console.WriteLine(e);
            throw;
        }
    }

    private static bool IsMatchingGenericType(ParameterInfo paramInfo, string typeName)
    {
        return paramInfo.ParameterType.IsGenericType
            && paramInfo.ParameterType.GetGenericTypeDefinition().FullName == typeName;
    }

    private void RegisterComponent(
        InjectionType injectionType,
        Type registrableInterface,
        Type implementationType
    )
    {
        switch (injectionType)
        {
            case InjectionType.Singleton:
                HandleSingletonRegistration(registrableInterface, implementationType);
                break;
            case InjectionType.Transient:
                serviceCollection.AddTransient(registrableInterface, implementationType);
                break;
            case InjectionType.Scoped:
                serviceCollection.AddScoped(registrableInterface, implementationType);
                break;
            default:
                throw new ArgumentOutOfRangeException(
                    nameof(injectionType),
                    $"Unknown injection type on {implementationType.Namespace}.{implementationType.Name}"
                );
        }
    }

    private void HandleSingletonRegistration(Type registrableInterface, Type implementationType)
    {
        var serviceKey = $"{implementationType.Namespace}.{implementationType.Name}";
        if (registrableInterface != implementationType)
        {
            serviceCollection.AddSingleton(
                registrableInterface,
                (serviceProvider) =>
                {
                    object service;
                    lock (_injectedValuesLock)
                    {
                        if (!_injectedValues.TryGetValue(serviceKey, out service))
                        {
                            service = serviceProvider.GetService(implementationType);
                            _injectedValues.Add(serviceKey, service);
                        }
                    }

                    return service;
                }
            );
        }
        else
        {
            serviceCollection.AddSingleton(registrableInterface, implementationType);
        }
    }

    private class TypeRefContainer
    {
        public Injectable InjectableAttribute { get; }
        public Type Type { get; }
        public Type ParentType { get; }

        public TypeRefContainer(Injectable injectable, Type parentType, Type type)
        {
            InjectableAttribute = injectable;
            Type = type;
            ParentType = parentType;
        }
    }
}