using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;

#if UNITY_EDITOR
using UnityEditor;
#endif

[Serializable]
public class EnumArray<TEnum, TData> : IEnumerable<TData>, ISerializationCallbackReceiver where TEnum : Enum
{
    public static readonly int s_Length;

    [SerializeField] private TData[] _values;
    public int Length => s_Length;

    static EnumArray()
    {
        var names = Enum.GetNames(typeof(TEnum));
        s_Length = names.Length;
        
#if UNITY_EDITOR
        EnumArrayDrawer.Names[typeof(TEnum)] = names;

        if (Enum.GetUnderlyingType(typeof(TEnum)) != typeof(int))
            Debug.LogWarning($"EnumArray should not be used with Enum type {typeof(TEnum)} as it doesn't derive from int");
        
        var values = Enum.GetValues(typeof(TEnum)) as int[];
        Array.Sort(values, (i0, i1) => i0 - i1);
        if (values[0] != 0) Debug.LogWarning($"EnumArray should not be used with Enum type {typeof(TEnum)} as it doesn't start from 0");
        for (int i = 1; i < values.Length; i++)
        {
            if (values[i] != values[i - 1] + 1)
            {
                Debug.LogWarning($"EnumArray should not be used with Enum type {typeof(TEnum)} as it contains a gap in enums");
                break;
            }
        }
#endif
    }

    public EnumArray()
    {
        _values = new TData[s_Length];
    }

    public ref TData this[TEnum e]
    {
        get => ref _values[e.GetHashCode()];
    }

    public ref TData this[int i]
    {
        get => ref _values[i];
    }
    
    public IEnumerator<TData> GetEnumerator()
    {
        return ((IEnumerable<TData>) _values).GetEnumerator();
    }
    
    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }

    public void OnAfterDeserialize()
    {
        // required check in case array was serialized when Enumerator class had a different number of elements
        if (_values.Length != s_Length)
            Array.Resize(ref _values, s_Length);
    }

    public void OnBeforeSerialize() { }
}

#if UNITY_EDITOR
[CustomPropertyDrawer(typeof(EnumArray<,>), true)]
public class EnumArrayDrawer : PropertyDrawer
{
    private const float Padding = 4;
    private const float Spacing = 2;

    public static readonly Dictionary<Type, string[]> Names = new();

    private Type _type;

    public override float GetPropertyHeight(SerializedProperty property, GUIContent label)
    {
        if (_type == null)
            _type = fieldInfo.FieldType.GenericTypeArguments[0];

        property.Next(true); // this will be the array

        // header + top/bottom padding
        float height = EditorGUIUtility.singleLineHeight + EditorGUIUtility.standardVerticalSpacing;
        if (property.isExpanded)
        {
            int arraySize = property.arraySize;
            
            // spacing between array elements
            height += arraySize * Spacing + 2 * Padding; 

            // get each array element height
            for (int i = 0; i < arraySize; i++)
                height += EditorGUI.GetPropertyHeight(property.GetArrayElementAtIndex(i));
        }

        return height;
    }

    public override void OnGUI(Rect position, SerializedProperty property, GUIContent label)
    {
        //var enumArrayProperty = property.
        
        var header = position;
        header.height = EditorGUIUtility.singleLineHeight;
        header.width -= EditorGUIUtility.standardVerticalSpacing * 3;

        bool expanded = EditorGUI.BeginFoldoutHeaderGroup(header, property.isExpanded, label);
        property.isExpanded = expanded;
        if (property.hasMultipleDifferentValues == false)
        {
            // only use header as a property if array has same values
            EditorGUI.BeginProperty(header, label, property);
            EditorGUI.EndProperty();
        }
        
        property.Next(true); // this will be the array
        
        if (expanded)
        {
            if (Event.current.type == EventType.Repaint)
            {
                Rect background = position;
            
                background.height -= EditorGUIUtility.singleLineHeight + 2 * EditorGUIUtility.standardVerticalSpacing;
                background.y += EditorGUIUtility.singleLineHeight + 2 * EditorGUIUtility.standardVerticalSpacing;
                EditorStyles.helpBox.Draw(background, false, false, false, false);
            }
            
            position.y += EditorGUIUtility.singleLineHeight + EditorGUIUtility.standardVerticalSpacing * 4;

            var labels = Names[_type];

            for (int i = 0; i < property.arraySize; i++)
            {
                Rect propertyRect = position;
                var serializedElement = property.GetArrayElementAtIndex(i);
                propertyRect.height = EditorGUI.GetPropertyHeight(serializedElement);
                label.text = labels[i];
                
                if (propertyRect.Contains(Event.current.mousePosition))
                {
                    Color col = EditorGUIUtility.isProSkin
                        ? new Color32(80, 80, 80, 255)
                        : new Color32 (210, 210, 210, 255);
                    EditorGUI.DrawRect(propertyRect, col);
                }

                propertyRect.x += 6 * EditorGUIUtility.standardVerticalSpacing;
                propertyRect.width -= 9 * EditorGUIUtility.standardVerticalSpacing;
                EditorGUI.PropertyField(propertyRect, serializedElement, label, true);
                
                position.y += propertyRect.height + Spacing;
            }
        }
        
        EditorGUI.EndFoldoutHeaderGroup();
    }
}
#endif