using System;
using Unity.MLAgents.Actuators;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Actuating
{
    [AddComponentMenu("/")] // hide component from menu
    public class ProxyMlActuatorComponent : ActuatorComponent
    {
        [field: NonSerialized] public ProxyMlActuator ProxyActuator { get; private set; }
        [SerializeField] public DirectControl originalControl;
        
        public override IActuator[] CreateActuators()
        {
            if (originalControl == null)
            {
                Debug.LogError($"{nameof(originalControl)} is not set");
                return Array.Empty<IActuator>();
            }
            ProxyActuator = new ProxyMlActuator(originalControl);
            return new IActuator[] { ProxyActuator };
        }

        public override ActionSpec ActionSpec => originalControl.ActionSpec;
    }
    
    public class ProxyMlActuator : IActuator
    {
        private readonly DirectControl _control;
        public string Name { get; }

        public ProxyMlActuator(DirectControl control)
        {
            _control = control;
            Name = _control.ActuatorName;
        }

        public void OnActionReceived(ActionBuffers actionBuffers)
        {
            _control.OnActionReceived(actionBuffers);
        }

        public void WriteDiscreteActionMask(IDiscreteActionMask actionMask)
        {
            _control.WriteDiscreteActionMask(actionMask);
        }

        public void Heuristic(in ActionBuffers actionBuffersOut)
        {
            if (!_control.IsHeuristicsEnabled)
            {
                return;
            }

            _control.Heuristic(actionBuffersOut);
        }

        public void ResetData()
        {
            _control.ResetData();
        }

        public ActionSpec ActionSpec => _control.ActionSpec;
    }
}