TheVVaS-Assets/RPGCoreCommon/Helpers/Editor/Drawers/SerializableDictionaryDrawer.cs

using System;
using System.Linq;
using System.Reflection;
using RPGCoreCommon.Helpers.CustomTypes;
using UnityEditor;
using UnityEditor.UIElements;
using UnityEngine;
using UnityEngine.UIElements;

namespace RPGCoreCommon.Helpers.Editor.Drawers
{
    [Serializable]
    [CustomPropertyDrawer(typeof(SerializableDictionary<,>), true)]
    public class SerializableDictionaryDrawer : PropertyDrawer
    {
        [SerializeField] private VisualTreeAsset _visualTreeAsset;
        
        private SerializedProperty _property;
        private SerializedProperty _tempKeyProperty;
        private SerializedProperty _keysProperty;
        private SerializedProperty _valuesProperty;

        private SerializableDictionaryAttribute _sdAttribute;
        private Type _keyType;
        private Type _valueType;

        public override VisualElement CreatePropertyGUI(SerializedProperty property)
        {
            _property = property;
            _sdAttribute = fieldInfo.GetCustomAttribute<SerializableDictionaryAttribute>() ?? new SerializableDictionaryAttribute();
            _tempKeyProperty = property.FindPropertyRelative("_tempKey");
            _keysProperty = property.FindPropertyRelative("_keys");
            _valuesProperty = property.FindPropertyRelative("_values");

            if (_keysProperty == null || _valuesProperty == null)
                return new HelpBox("Keys and Values in SerializableDictionary must be serializable!", HelpBoxMessageType.Error);
            
            var genericTypes = property.boxedValue.GetType().FindGenericDefinitionType(typeof(SerializableDictionary<,>)).GetGenericArguments();
            _keyType = genericTypes[0];
            _valueType = genericTypes[1];
            
            var root = _visualTreeAsset.CloneTree();
            root.Q<Label>().text = preferredLabel;
            root.DisableDefaultAligning();
            
            if (string.IsNullOrEmpty(preferredLabel))
            {
                var foldout = root.Q<Foldout>();
                foldout.Q<Toggle>().style.display = DisplayStyle.None;
                foldout.Q<Toggle>().value = true;
                foldout.Q(name: "unity-content").style.marginLeft = 0;
            }
            
            var list = root.Q<ListView>("sd-list");
            list.BindProperty(_valuesProperty);
            list.makeItem = CreateRow;
            list.bindItem = BindRow;
            list.onAdd = OnAdd;
            list.onRemove = OnRemove;
            list.itemIndexChanged += OnSwap;

            if (_sdAttribute is not null)
            {
                root.Q("sd-header").Q<Label>(className: "row-left").text = _sdAttribute.keyName;
                root.Q("sd-header").Q<Label>(className: "row-right").text = _sdAttribute.valueName;
                root.EnableInClassList("keys-disabled", !_sdAttribute.isKeyEditable);
            }

            var dragline = root.Q("sd-header").Q("unity-dragline-anchor"); 
            dragline.RegisterCallback<GeometryChangedEvent>(ev =>
            {
                var parentWidth = dragline.parent.worldBound.width;
                var leftWidth = dragline.style.left.value.value;
                leftWidth = Mathf.Clamp(leftWidth, 150f, parentWidth-180f);
                
                dragline.style.left = leftWidth;
                root.Query(className: "row-left").ForEach(element => element.style.width = leftWidth);
            });
            
            var tempKeyField = root.Q("sd-add-row").Q<PropertyField>(); 
            tempKeyField.BindProperty(_tempKeyProperty);
            tempKeyField.HideFoldoutAndLabel();
            
            // This additional key property should be always empty on beginning to not confuse user.
            // All because we are using "fake" serialized property to make unity draw field for key before even adding it to dictionary.
            _tempKeyProperty.boxedValue = _keyType.IsValueType ? Activator.CreateInstance(_keyType) : null;
            _property.serializedObject.ApplyModifiedProperties();
            _property.serializedObject.Update();
            
            return root;
        }

        private void OnSwap(int index1, int index2)
        {
            _keysProperty.MoveArrayElement(index1, index2);
            _property.serializedObject.ApplyModifiedProperties();
        }

        private void OnAdd(BaseListView list)
        {
            var tempKey = _tempKeyProperty.boxedValue;

            if (tempKey == null)
            {
                var helpBox = new HelpBox("Key is not selected!", HelpBoxMessageType.Error);
                helpBox.FadeOut(3000, 2000, true);
                list.parent.parent.Insert(0, helpBox);
                return;
            }
            
            var keyExists = Enumerable.Range(0, _keysProperty.arraySize)
                .Select(i => _keysProperty.GetArrayElementAtIndex(i))
                .Any(prop => prop.boxedValue.Equals(tempKey));
            
            if (keyExists)
            {
                var helpBox = new HelpBox("Key already exists!", HelpBoxMessageType.Error);
                helpBox.FadeOut(3000, 2000, true);
                list.parent.parent.Insert(0, helpBox);
                return;
            }

            // Reset temporary key to default
            _tempKeyProperty.boxedValue = _keyType.IsValueType ? Activator.CreateInstance(_keyType) : null;
            
            // New Key
            _keysProperty.arraySize++;
            _keysProperty.GetArrayElementAtIndex(_keysProperty.arraySize - 1).boxedValue = tempKey;
            
            // New Value
            _valuesProperty.arraySize++;
            _valuesProperty.GetArrayElementAtIndex(_valuesProperty.arraySize - 1).boxedValue =
                _valueType.IsValueType ? Activator.CreateInstance(_valueType) : null;
            
            // Update list and property
            _property.serializedObject.ApplyModifiedProperties();
            _property.serializedObject.Update();
            list.RefreshItems();
        }

        private void OnRemove(BaseListView list)
        {
            var ids = list.selectedIds.SortDescending().ToList();
            if (ids.Count < 0) return;
            
            ids.ForEach(i =>
            {
                _keysProperty.DeleteArrayElementAtIndex(i);
                _valuesProperty.DeleteArrayElementAtIndex(i);
            });
            
            // Update list and property
            _property.serializedObject.ApplyModifiedProperties();
            _property.serializedObject.Update();
            list.RefreshItems();
        }

        private VisualElement CreateRow()
        {
            var row = new VisualElement();
            row.AddToClassList("row");

            var keyField = new PropertyField();
            keyField.DisableDefaultAligning();
            keyField.AddToClassList("row-left");
            keyField.label = "";
            keyField.SetEnabled(_sdAttribute.isKeyEditable);
            row.Add(keyField);

            var valueField = new PropertyField();
            valueField.HideFoldoutAndLabel();
            valueField.SetContextAligningHere();
            valueField.AddToClassList("row-right");
            valueField.label = "";
            row.Add(valueField);
            
            return row;
        }

        private void BindRow(VisualElement row, int i)
        {
            if (i >= _keysProperty.arraySize) return;
            row.Q<PropertyField>(className: "row-left").BindProperty(_keysProperty.GetArrayElementAtIndex(i));
            row.Q<PropertyField>(className: "row-right").BindProperty(_valuesProperty.GetArrayElementAtIndex(i));
        }
    }
}