#nullable enable

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

namespace ShrinkNetwork
{
    [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
    public sealed class ShrinkNetworkMessageRegistryAttribute : Attribute
    {
        public ShrinkNetworkMessageRegistryAttribute(params Type[] messageTypes)
        {
            MessageTypes = messageTypes ?? Array.Empty<Type>();
        }

        public Type[] MessageTypes { get; }
    }

    [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
    public sealed class ShrinkNetworkStaticSubscriberRegistryAttribute : Attribute
    {
        public ShrinkNetworkStaticSubscriberRegistryAttribute(params Type[] subscriberTypes)
        {
            SubscriberTypes = subscriberTypes ?? Array.Empty<Type>();
        }

        public Type[] SubscriberTypes { get; }
    }

    internal static class ShrinkNetworkGeneratedRegistry
    {
        public static IReadOnlyList<Type> GetAttributedMessageTypes()
            => GetAssemblyRegisteredTypes<ShrinkNetworkMessageRegistryAttribute>(attribute => attribute.MessageTypes);

        public static IReadOnlyList<Type> GetStaticSubscriberTypes()
            => GetAssemblyRegisteredTypes<ShrinkNetworkStaticSubscriberRegistryAttribute>(attribute => attribute.SubscriberTypes);

        public static void RegisterAll(ShrinkNetworkService service)
        {
            if (service == null)
                throw new ArgumentNullException(nameof(service));

            ShrinkNetworkRegHelper.RegisterAttributedMessages(service, GetAttributedMessageTypes());
            ShrinkNetworkRegHelper.RegisterStaticHandlers(service, GetStaticSubscriberTypes());
        }

        private static IReadOnlyList<Type> GetAssemblyRegisteredTypes<TAttribute>(Func<TAttribute, Type[]> selector)
            where TAttribute : Attribute
        {
            var types = new List<Type>();
            var seen = new HashSet<Type>();
            foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies())
            {
                object[] attributes;
                try
                {
                    attributes = assembly.GetCustomAttributes(typeof(TAttribute), false);
                }
                catch
                {
                    continue;
                }

                foreach (var attribute in attributes.OfType<TAttribute>())
                {
                    foreach (var registeredType in selector(attribute) ?? Array.Empty<Type>())
                    {
                        if (registeredType == null || !seen.Add(registeredType))
                            continue;

                        types.Add(registeredType);
                    }
                }
            }

            return types;
        }
    }
}