BITKit/Src/Unity/Scripts/UX/Input/OnScreenControl.cs

using System;
using System.Collections.Generic;
using Unity.Collections;
using UnityEngine;
using UnityEngine.InputSystem;
using UnityEngine.InputSystem.Layouts;
using UnityEngine.InputSystem.LowLevel;
using UnityEngine.InputSystem.Utilities;
using UnityEngine.UIElements;

namespace BITKit.UX
{
     public abstract class OnScreenControl:VisualElement
    {
        
        public new class UxmlTraits : VisualElement.UxmlTraits
        {
            private readonly UxmlStringAttributeDescription m_ControlPathAttribute = new ()
            {
                name = "ControlPath"
            };

            public override void Init(VisualElement ve, IUxmlAttributes bag, CreationContext cc)
            {
                base.Init(ve, bag, cc);
                var table = (OnScreenControl)ve;
                table.ControlPath = m_ControlPathAttribute.GetValueFromBag(bag, cc);
            }
        }

        protected OnScreenControl()
        {
            RegisterCallback<GeometryChangedEvent>(OnGeometryChanged);
            RegisterCallback<AttachToPanelEvent>(OnAttachToPanel);
            RegisterCallback<DetachFromPanelEvent>(OnDetachFromPanel);
            IsEnabled.AddDisableElements(this);
            
            IsEnabled.AddListener(x =>
            {
                if (x)
                    OnEnable();
                else OnDisable();
            });
        }

        private void OnDetachFromPanel(DetachFromPanelEvent evt)
        {
            IsEnabled.AddDisableElements(32);
        }

        private void OnAttachToPanel(AttachToPanelEvent evt)
        {
            IsEnabled.RemoveDisableElements(32);
        }

        private void OnGeometryChanged(GeometryChangedEvent evt)
        {
            IsEnabled.SetDisableElements(this, !visible);
        }

        protected readonly ValidHandle IsEnabled = new();

        /// <summary>
        /// The control path (see <see cref="InputControlPath"/>) for the control that the on-screen
        /// control will feed input into.
        /// </summary>
        /// <remarks>
        /// A device will be created from the device layout referenced by the control path (see
        /// <see cref="InputControlPath.TryGetDeviceLayout"/>). The path is then used to look up
        /// <see cref="Control"/> on the device. The resulting control will be fed values from
        /// the on-screen control.
        ///
        /// Multiple on-screen controls sharing the same device layout will together create a single
        /// virtual device. If, for example, one component uses <c>"&lt;Gamepad&gt;/buttonSouth"</c>
        /// and another uses <c>"&lt;Gamepad&gt;/leftStick"</c> as the control path, a single
        /// <see cref="Gamepad"/> will be created and the first component will feed data to
        /// <see cref="Gamepad.buttonSouth"/> and the second component will feed data to
        /// <see cref="Gamepad.leftStick"/>.
        /// </remarks>
        /// <seealso cref="InputControlPath"/>
        public string ControlPath
        {
            get => ControlPathInternal;
            set
            {
                IsEnabled.SetElements(1, string.IsNullOrEmpty(value) is false);
                ControlPathInternal = value;
                IsEnabled.AddDisableElements("Force");
                IsEnabled.RemoveDisableElements("Force");
            }
        }

        /// <summary>
        /// The actual control that is fed input from the on-screen control.
        /// </summary>
        /// <remarks>
        /// This is only valid while the on-screen control is enabled. Otherwise, it is <c>null</c>. Also,
        /// if no <see cref="ControlPath"/> has been set, this will remain <c>null</c> even if the component is enabled.
        /// </remarks>
        public InputControl Control => _mControl;

        private InputControl _mControl;
        private OnScreenControl _mNextControlOnDevice;
        private InputEventPtr _mInputEventPtr;

        /// <summary>
        /// Accessor for the <see cref="ControlPath"/> of the component. Must be implemented by subclasses.
        /// </summary>
        /// <remarks>
        /// Moving the definition of how the control path is stored into subclasses allows them to
        /// apply their own <see cref="InputControlAttribute"/> attributes to them and thus set their
        /// own layout filters.
        /// </remarks>
        protected abstract string ControlPathInternal { get; set; }

        private void SetupInputControl()
        {
            Debug.Assert(_mControl == null, "InputControl already initialized");
            Debug.Assert(_mNextControlOnDevice == null, "Previous InputControl has not been properly uninitialized (m_NextControlOnDevice still set)");
            Debug.Assert(!_mInputEventPtr.valid, "Previous InputControl has not been properly uninitialized (m_InputEventPtr still set)");

            // Nothing to do if we don't have a control path.
            var path = ControlPathInternal;
            if (string.IsNullOrEmpty(path))
                return;

            // Determine what type of device to work with.
            var layoutName = InputControlPath.TryGetDeviceLayout(path);
            if (layoutName == null)
            {
                Debug.LogError(
                    $"Cannot determine device layout to use based on control path '{path}' used in {GetType().Name} component with {name}");
                return;
            }

            // Try to find existing on-screen device that matches.
            var internedLayoutName = new InternedString(layoutName);
            var deviceInfoIndex = -1;
            for (var i = 0; i < OnScreenDevices.Count; ++i)
            {
                ////FIXME: this does not take things such as different device usages into account
                if (OnScreenDevices[i].Device.layout != internedLayoutName) continue;
                deviceInfoIndex = i;
                break;
            }

            // If we don't have a matching one, create a new one.
            InputDevice device;
            if (deviceInfoIndex == -1)
            {
                // Try to create device.
                try
                {
                    device = InputSystem.AddDevice(layoutName);
                }
                catch (Exception exception)
                {
                    Debug.LogError(
                        $"Could not create device with layout '{layoutName}' used in '{GetType().Name}' component");
                    Debug.LogException(exception);
                    return;
                }
                InputSystem.AddDeviceUsage(device, "OnScreen");

                // Create event buffer.
                var buffer = StateEvent.From(device, out var eventPtr, Allocator.Persistent);

                // Add to list.
                deviceInfoIndex = OnScreenDevices.Count;
                OnScreenDevices.Add(new OnScreenDeviceInfo
                {
                    EventPtr = eventPtr,
                    Buffer = buffer,
                    Device = device,
                });
            }
            else
            {
                device = OnScreenDevices[deviceInfoIndex].Device;
            }

            // Try to find control on device.
            _mControl = InputControlPath.TryFindControl(device, path);
            if (_mControl == null)
            {
                Debug.LogError(
                    $"Cannot find control with path '{path}' on device of type '{layoutName}' referenced by component '{GetType().Name}' with {name}");

                // Remove the device, if we just created one.
                if (OnScreenDevices[deviceInfoIndex].FirstControl == null)
                {
                    OnScreenDevices[deviceInfoIndex].Destroy();
                    OnScreenDevices.RemoveAt(deviceInfoIndex);
                }

                return;
            }
            _mInputEventPtr = OnScreenDevices[deviceInfoIndex].EventPtr;

            // We have all we need. Permanently add us.
            OnScreenDevices[deviceInfoIndex] =
                OnScreenDevices[deviceInfoIndex].AddControl(this);
        }

        protected void SendValueToControl<TValue>(TValue value)
            where TValue : struct
        {
            if (_mControl == null)
                return;

            if (!(_mControl is InputControl<TValue> control))
                throw new ArgumentException(
                    $"The control path {ControlPath} yields a control of type {_mControl.GetType().Name} which is not an InputControl with value type {typeof(TValue).Name}", nameof(value));

            ////FIXME: this gives us a one-frame lag (use InputState.Change instead?)
            _mInputEventPtr.time = InputState.currentTime;
            control.WriteValueIntoEvent(value, _mInputEventPtr);
            InputSystem.QueueEvent(_mInputEventPtr);
        }

        protected void SentDefaultValueToControl()
        {
            if (_mControl == null)
                return;

            ////FIXME: this gives us a one-frame lag (use InputState.Change instead?)
            _mInputEventPtr.time = InputState.currentTime;
            _mControl.ResetToDefaultStateInEvent(_mInputEventPtr);
            InputSystem.QueueEvent(_mInputEventPtr);
        }

        protected virtual void OnEnable()
        {
            SetupInputControl();
        }

        protected virtual void OnDisable()
        {
            if (_mControl == null)
                return;

            var device = _mControl.device;
            for (var i = 0; i < OnScreenDevices.Count; ++i)
            {
                if (OnScreenDevices[i].Device != device)
                    continue;

                var deviceInfo = OnScreenDevices[i].RemoveControl(this);
                if (deviceInfo.FirstControl == null)
                {
                    // We're the last on-screen control on this device. Remove the device.
                    OnScreenDevices[i].Destroy();
                    OnScreenDevices.RemoveAt(i);
                }
                else
                {
                    OnScreenDevices[i] = deviceInfo;

                    // We're keeping the device , but we're disabling the on-screen representation
                    // for one of its controls. If the control isn't in default state, reset it
                    // to that now. This is what ensures that if, for example, OnScreenButton is
                    // disabled after OnPointerDown, we reset its button control to zero even
                    // though we will not see an OnPointerUp.
                    if (!_mControl.CheckStateIsAtDefault())
                        SentDefaultValueToControl();
                }

                _mControl = null;
                _mInputEventPtr = new InputEventPtr();
                Debug.Assert(_mNextControlOnDevice == null);

                break;
            }
        }

        private struct OnScreenDeviceInfo
        {
            public InputEventPtr EventPtr;
            public NativeArray<byte> Buffer;
            public InputDevice Device;
            public OnScreenControl FirstControl;

            public OnScreenDeviceInfo AddControl(OnScreenControl control)
            {
                control._mNextControlOnDevice = FirstControl;
                FirstControl = control;
                return this;
            }

            public OnScreenDeviceInfo RemoveControl(OnScreenControl control)
            {
                if (FirstControl == control)
                    FirstControl = control._mNextControlOnDevice;
                else
                {
                    for (OnScreenControl current = FirstControl._mNextControlOnDevice, previous = FirstControl;
                         current != null; previous = current, current = current._mNextControlOnDevice)
                    {
                        if (current != control)
                            continue;

                        previous._mNextControlOnDevice = current._mNextControlOnDevice;
                        break;
                    }
                }

                control._mNextControlOnDevice = null;
                return this;
            }

            public void Destroy()
            {
                if (Buffer.IsCreated)
                    Buffer.Dispose();
                if (Device != null)
                    InputSystem.RemoveDevice(Device);
                Device = null;
                Buffer = new NativeArray<byte>();
            }
        }

        private static readonly List<OnScreenDeviceInfo> OnScreenDevices=new ();

        internal string GetWarningMessage()
        {
            return $"{GetType()} needs to be attached as a child to a UI Canvas and have a RectTransform component to function properly.";
        }
    }
}