Editor/ControlWidgets/TypeSelectControlWidget.cs
using Editor;
using Editor.NodeEditor;
using ExtendedEditor.Attributes;
using ExtendedEditor.TypeLibraryFixes;
using Sandbox;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;

namespace ExtendedEditor.ControlWidgets;

public class TypeSelectControlWidget : ControlWidget
{
    public override bool IsControlButton => true;
    public override bool IsControlHovered => base.IsControlHovered || _menu.IsValid();
    public override bool SupportsMultiEdit => false;

    private Menu? _menu;

    private bool CanSetNone => IsValidType(null);

    private readonly bool _whitelistedOnly;
    private Type? _targetTypeAttributeType;
    private TypeSelectAttribute? _limiter;

    public TypeSelectControlWidget(SerializedProperty property) : this(property, true)
    {
    }

    public TypeSelectControlWidget(SerializedProperty property, bool whitelistedOnly = true) : base(property.Fix())
    {
        _whitelistedOnly = whitelistedOnly;

        Cursor = CursorShape.Finger;

        Layout = Layout.Row();
        Layout.Spacing = 2;

        UpdateLimiters();

        if(!CanSetNone)
        {
            var value = SerializedProperty.GetValue<Type>();
            if(value is null)
            {
                value = GetPossibleTypes().FirstOrDefault();
                if(value is not null)
                    SerializedProperty.SetValue(value);
            }
        }
    }

    [EditorEvent.Hotload]
    private void UpdateLimiters()
    {
        var attributes = SerializedProperty.GetAttributes();
        _targetTypeAttributeType = attributes?.OfType<TargetTypeAttribute>().FirstOrDefault()?.Type;
        _limiter = attributes?.OfType<TypeSelectAttribute>().FirstOrDefault();
        if(_limiter?.ValidatorName is not null)
        {
            _limiter = new TypeSelectAttribute(_limiter);
            _limiter.FindAndAppendValidatorMethod(SerializedProperty);
        }
    }

    protected override void PaintControl()
    {
        var value = SerializedProperty.GetValue<Type>();

        var color = IsControlHovered ? Theme.Blue : Theme.TextControl;
        var rect = LocalRect;

        rect = rect.Shrink(8, 0);

        var desc = value is not null ? TypeLibrary.GetType(value) : null;

        var title = "None";
        if(desc is not null)
        {
            title = desc.Title;
            if(value!.IsGenericType)
            {
                var parameters = value.GetGenericArguments();
                title += $" <{string.Join(", ", parameters.Select(x => x.Name))}>";
            }
        }

        Paint.SetPen(color);
        Paint.DrawText(rect, title, TextFlag.LeftCenter);

        Paint.SetPen(color);
        Paint.DrawIcon(rect, "Arrow_Drop_Down", 17, TextFlag.RightCenter);
    }

    protected override void OnMousePress(MouseEvent e)
    {
        base.OnMousePress(e);

        if(e.LeftMouseButton && !_menu.IsValid())
        {
            OpenMenu();
        }
    }

    public IEnumerable<Type> GetPossibleTypes() => GetPossibleTypes(IsValidType);

    private bool IsValidType(Type? type) => IsValidType(type, _targetTypeAttributeType, _limiter, _whitelistedOnly);


    public static IEnumerable<Type> GetPossibleTypes(IEnumerable<Attribute> attributes)
    {
        var targetTypeAttributeType = attributes.OfType<TargetTypeAttribute>().FirstOrDefault()?.Type;
        var limiter = attributes.OfType<TypeSelectAttribute>().FirstOrDefault();
        return GetPossibleTypes(t => IsValidType(t, targetTypeAttributeType, limiter, whitelistedOnly: true));
    }

    public static IEnumerable<Type> GetPossibleTypes(Func<Type?, bool> typeValidator)
    {
        var listedTypes = new HashSet<Type>();

        var allTypes = TypeResolver.SystemTypes.Concat(
            TypeLibrary.GetTypes().Select(x => x.TargetType)
        );

        foreach(var type in allTypes)
        {
            if(!typeValidator(type))
                continue;
            if(!listedTypes.Add(type))
                continue;
            yield return type;
        }
    }

    private static bool IsValidType(Type? type, Type? baseType, TypeSelectAttribute? limiter, bool whitelistedOnly)
    {
        if(type is not null)
        {
            if(type.IsAbstract && type.IsSealed) // is static
                return false;
            if(type.CustomAttributes.Any(x => x.AttributeType == typeof(CompilerGeneratedAttribute)))
                return false;
            if(type.Name.StartsWith('<') || type.Name.StartsWith('_'))
                return false;
            if(baseType is not null && !type.IsAssignableTo(baseType))
                return false;
            if(whitelistedOnly && TypeLibrary.GetType(type) is null)
                return false;
        }

        if(limiter is not null && !limiter.IsAllowed(type))
            return false;

        return true;
    }

    public Menu CreateMenu(Action<Type?>? action = null)
    {
        var types = GetPossibleTypes().Select(x => new TypeResolver.TypeOption(x)).ToArray();

        var menu = new ContextMenu(null);

        bool canSetNone = CanSetNone;

        menu.AddLineEdit("Filter",
            placeholder: "Filter Types..",
            autoFocus: true,
            onChange: s => PopulateTypeMenu(menu, types, action, s, canSetNone));

        menu.AboutToShow += () =>
        {
            PopulateTypeMenu(menu, types, action, canSetNone: canSetNone);
        };

        return menu;
    }

    private void OpenMenu()
    {
        _menu = CreateMenu(type =>
        {
            PropertyStartEdit();
            SerializedProperty.SetValue(type);
            PropertyFinishEdit();
            SignalValuesChanged();
        });

        _menu.DeleteOnClose = true;
        _menu.OpenAtCursor(true);
        _menu.MinimumWidth = ScreenRect.Width;
    }

    private static void PopulateTypeMenu(Menu menu, IEnumerable<TypeResolver.TypeOption> types, Action<Type?>? action, string? filter = null, bool canSetNone = false)
    {
        menu.RemoveMenus();
        menu.RemoveOptions();

        foreach(var widget in menu.Widgets.Skip(1))
        {
            menu.RemoveWidget(widget);
        }

        if(canSetNone)
        {
            menu.AddOption("None", "cancel", () => action?.Invoke(null));
        }

        const int maxFiltered = 10;

        var useFilter = !string.IsNullOrEmpty(filter);
        var truncated = 0;

        if(useFilter)
        {
            var filtered = types.Where(x => x.Type.Name.Contains(filter!, StringComparison.OrdinalIgnoreCase)).ToArray();

            if(filtered.Length > maxFiltered + 1)
            {
                truncated = filtered.Length - maxFiltered;
                types = filtered.Take(maxFiltered);
            }
            else
            {
                types = filtered;
            }
        }

        menu.AddOptions(types, x => x.Path, x => action?.Invoke(x.Type), flat: useFilter);

        if(truncated > 0)
        {
            menu.AddOption($"...and {truncated} more");
        }

        menu.AdjustSize();
        menu.Update();
    }
}