using Fluctio.FluctioSim.Common.Configuration;
using Fluctio.FluctioSim.Core.MujocoExtensions;
using Fluctio.FluctioSim.Core.MujocoExtensions.States;
using Mujoco;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.TriggerActions
{
	[AddComponentMenu(Config.PrefixedName + "/Machine Learning/Trigger Actions/Reset MuJoCo Chain", Config.ComponentMenuOrder + 80)]
	[DefaultExecutionOrder(10)] // should be after MjScene, but before Randomizers
	public class MujocoChainReseter : TriggerAction
	{
		[SerializeField] public MjBody body;
		private ChainState _initialState = null;

		private void Reset()
		{
			body = GetComponent<MjBody>();
		}

		private void Start()
		{
			if (body == null)
			{
				Debug.LogWarning("No body specified, skipping reset");
				return;
			}
			_initialState = body.GetChainState();
		}

		public override void OnExecute()
		{
			if (_initialState == null)
			{
				return;
			}
			body.SetChainState(_initialState);
		}
	}
}