encounter-system/Editor/SubclassSelectorDrawer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Jovian.EncounterSystem;
using UnityEditor;
using UnityEngine;

namespace Jovian.EncounterSystem.Editor {
    /// <summary>
    /// Custom drawer for any field marked <c>[SerializeReference, SubclassSelector]</c>. Renders a
    /// "pick a concrete type" dropdown populated by reflection over the declared base type, then
    /// draws the chosen instance's serialized children underneath. Works for single fields, arrays,
    /// and <see cref="List{T}"/> elements.
    /// </summary>
    /// <remarks>
    /// The type cache is a <c>static readonly</c> dictionary; it is implicitly cleared on domain
    /// reload because the static field is recreated with the new assembly image.
    /// </remarks>
    [CustomPropertyDrawer(typeof(SubclassSelectorAttribute))]
    public class SubclassSelectorDrawer : PropertyDrawer {
        private static readonly Dictionary<Type, Type[]> TypeCache = new();

        public override void OnGUI(Rect position, SerializedProperty property, GUIContent label) {
            if(property.propertyType != SerializedPropertyType.ManagedReference) {
                EditorGUI.LabelField(position, label.text, "[SubclassSelector] requires [SerializeReference]");
                return;
            }

            var baseType = ResolveBaseType(fieldInfo.FieldType);
            var concreteTypes = GetConcreteTypes(baseType);
            var currentType = GetCurrentType(property);
            var currentIndex = Array.IndexOf(concreteTypes, currentType);

            var names = new string[concreteTypes.Length + 1];
            names[0] = "<None>";
            for(int i = 0; i < concreteTypes.Length; i++) {
                names[i + 1] = concreteTypes[i].Name;
            }

            var displayLabel = ResolveDisplayLabel(property, label);
            var headerRect = new Rect(position.x, position.y, position.width, EditorGUIUtility.singleLineHeight);
            var labelRect = new Rect(headerRect.x, headerRect.y, EditorGUIUtility.labelWidth, headerRect.height);
            var popupRect = new Rect(headerRect.x + EditorGUIUtility.labelWidth, headerRect.y, headerRect.width - EditorGUIUtility.labelWidth, headerRect.height);

            EditorGUI.LabelField(labelRect, displayLabel);
            var newIndex = EditorGUI.Popup(popupRect, currentIndex + 1, names);
            if(newIndex != currentIndex + 1) {
                property.managedReferenceValue = newIndex == 0 ? null : Activator.CreateInstance(concreteTypes[newIndex - 1]);
                property.serializedObject.ApplyModifiedProperties();
            }

            if(property.managedReferenceValue == null) {
                return;
            }

            EditorGUI.indentLevel++;
            var y = position.y + EditorGUIUtility.singleLineHeight + EditorGUIUtility.standardVerticalSpacing;
            var iterator = property.Copy();
            var end = iterator.GetEndProperty();
            if(iterator.NextVisible(true)) {
                while(!SerializedProperty.EqualContents(iterator, end)) {
                    var h = EditorGUI.GetPropertyHeight(iterator, true);
                    var r = new Rect(position.x, y, position.width, h);
                    EditorGUI.PropertyField(r, iterator, true);
                    y += h + EditorGUIUtility.standardVerticalSpacing;
                    if(!iterator.NextVisible(false)) {
                        break;
                    }
                }
            }
            EditorGUI.indentLevel--;
        }

        public override float GetPropertyHeight(SerializedProperty property, GUIContent label) {
            var height = EditorGUIUtility.singleLineHeight;
            if(property.propertyType != SerializedPropertyType.ManagedReference || property.managedReferenceValue == null) {
                return height;
            }

            var iterator = property.Copy();
            var end = iterator.GetEndProperty();
            if(iterator.NextVisible(true)) {
                while(!SerializedProperty.EqualContents(iterator, end)) {
                    height += EditorGUI.GetPropertyHeight(iterator, true) + EditorGUIUtility.standardVerticalSpacing;
                    if(!iterator.NextVisible(false)) {
                        break;
                    }
                }
            }
            return height;
        }

        private static Type ResolveBaseType(Type fieldType) {
            if(fieldType.IsArray) {
                return fieldType.GetElementType();
            }
            if(fieldType.IsGenericType && fieldType.GetGenericTypeDefinition() == typeof(List<>)) {
                return fieldType.GetGenericArguments()[0];
            }
            return fieldType;
        }

        private static Type[] GetConcreteTypes(Type baseType) {
            if(TypeCache.TryGetValue(baseType, out var cached)) {
                return cached;
            }

            var types = AppDomain.CurrentDomain.GetAssemblies()
                .SelectMany(assembly => {
                    try {
                        return assembly.GetTypes();
                    }
                    catch {
                        return Array.Empty<Type>();
                    }
                })
                .Where(type => baseType.IsAssignableFrom(type)
                    && !type.IsAbstract
                    && !type.IsInterface
                    && !typeof(UnityEngine.Object).IsAssignableFrom(type)
                    && type.GetConstructor(Type.EmptyTypes) != null)
                .OrderBy(type => type.Name)
                .ToArray();
            TypeCache[baseType] = types;
            return types;
        }

        private static GUIContent ResolveDisplayLabel(SerializedProperty property, GUIContent fallback) {
            var value = property.managedReferenceValue;
            if(value is not IEncounter) {
                return fallback;
            }

            var definitionProp = value.GetType().GetProperty("EncounterDefinition");
            var definition = definitionProp?.GetValue(value);
            var nameField = definition?.GetType().GetField("name");
            var name = nameField?.GetValue(definition) as string;
            return string.IsNullOrEmpty(name) ? fallback : new GUIContent(name);
        }

        private static Type GetCurrentType(SerializedProperty property) {
            var full = property.managedReferenceFullTypename;
            if(string.IsNullOrEmpty(full)) {
                return null;
            }

            var space = full.IndexOf(' ');
            if(space < 0) {
                return null;
            }

            var assembly = full.Substring(0, space);
            var typeName = full.Substring(space + 1);
            return Type.GetType($"{typeName}, {assembly}");
        }
    }
}