// FlexiMotion // https://kybernetik.com.au/flexi-motion // Copyright 2023-2025 Kybernetik //

using FlexiMotion.Jobs;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using Unity.Jobs;
using UnityEngine;

namespace FlexiMotion.Modifiers
{
    /// <summary>
    /// A <see cref="FlexiMotionModifier"/> which alters the simulation to be in an object's local-space instead of in
    /// world-space.
    /// </summary>
    /// https://kybernetik.com.au/flexi-motion/api/FlexiMotion.Modifiers/SimulationRootModifier
    /// 
    [AddComponentMenu(FMStrings.ModifiersMenuPrefix + "Simulation Root Modifier")]
    [HelpURL(FMStrings.DocsURLs.ModifiersAPIDocumentation + "/" + nameof(SimulationRootModifier))]
    public class SimulationRootModifier : FlexiMotionModifier,
        JobModifierGroup.IOnReEnable,
        JobModifierGroup.IPreUpdate,
        JobModifierGroup.IPostUpdate
    {
        /************************************************************************************************************************/

        /// <summary><see cref="PreUpdate"/> before anything else and <see cref="PostUpdate"/> after.</summary>
        public const int
            PreExecutionTime = -10000,
            PostExecutionTime = 10000;

        /************************************************************************************************************************/

        private bool _HasRoot;

        [SerializeField, Tooltip(FMStrings.Tooltips.SimulationRoot)]
        private Transform _SimulationRoot;

        /// <summary>
        /// Inertia is simulated in this object's local-space. If not set, simulation will be done in world-space.
        /// </summary>
        /// <remarks>
        /// If you want to teleport the character without affecting the simulation, set this to the object that will be
        /// teleported before doing so, then revert it to its previous value afterwards.
        /// </remarks>
        public Transform SimulationRoot
        {
            get => _SimulationRoot;
            set
            {
                if (_SimulationRoot == value)
                    return;

                if (_SimulationRoot != null && Target != null && Target.Runtime != null)
                {
                    ScheduleLocalToWorldJob(_SimulationRoot).Complete();
                }

                _SimulationRoot = value;

                if (_SimulationRoot != null)
                {
                    _HasRoot = true;
                    if (Target != null && Target.Runtime != null)
                    {
                        ScheduleWorldToLocalJob(_SimulationRoot).Complete();
                    }
                }
                else
                {
                    _HasRoot = false;
                }
            }
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        protected override void OnValidate()
        {
            base.OnValidate();

            if (_SimulationRoot == null)
                _SimulationRoot = transform;
        }

        /************************************************************************************************************************/

        /// <summary>Initializes this modifier.</summary>
        protected virtual void Awake()
        {
            _HasRoot = _SimulationRoot != null;
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        protected override void Enable(JobModifierGroup modifiers)
        {
            base.Enable(modifiers);

#if UNITY_EDITOR
            modifiers.SimulationRoot = SimulationRoot;
#endif

            modifiers.OnReEnable.Add(PostExecutionTime, this);
            modifiers.PreUpdate.Add(PreExecutionTime, this);
            modifiers.PostUpdate.Add(PostExecutionTime, this);
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        protected override bool Disable(JobModifierGroup modifiers)
        {
            base.Disable(modifiers);

#if UNITY_EDITOR
            modifiers.SimulationRoot = null;
#endif

            var wasEnabled = false;
            wasEnabled |= modifiers.OnReEnable.Remove(PostExecutionTime);
            wasEnabled |= modifiers.PreUpdate.Remove(PreExecutionTime);
            wasEnabled |= modifiers.PostUpdate.Remove(PostExecutionTime);
            return wasEnabled;
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        public virtual JobHandle OnReInitialize(JobHandle dependsOn = default)
        {
            switch (Target.Definition.ActionOnDisable)
            {
                case ActionOnDisable.Recenter:
                case ActionOnDisable.Reset:
                    return PostUpdate(dependsOn);

                default:
                    return dependsOn;
            }
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        public virtual JobHandle PreUpdate(JobHandle dependsOn = default)
        {
            if (_HasRoot)
                dependsOn = ScheduleLocalToWorldJob(_SimulationRoot, dependsOn);

            return dependsOn;
        }

        /************************************************************************************************************************/

        /// <inheritdoc/>
        public virtual JobHandle PostUpdate(JobHandle dependsOn = default)
        {
            if (_HasRoot)
                dependsOn = ScheduleWorldToLocalJob(_SimulationRoot, dependsOn);

            return dependsOn;
        }

        /************************************************************************************************************************/

        /// <summary>Schedule a job to convert the simulation to local-space.</summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private JobHandle ScheduleWorldToLocalJob(Transform root, JobHandle dependsOn = default)
            => new TransformPositionsJob
            {
                transform = root.worldToLocalMatrix,
                positions = Target.Runtime.RuntimeSprings.TailPositions,
                previousPositions = Target.Runtime.RuntimeSprings.TailPreviousPositions,
            }.Schedule(dependsOn);

        /// <summary>Schedule a job to convert the simulation to world-space.</summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private JobHandle ScheduleLocalToWorldJob(Transform root, JobHandle dependsOn = default)
            => new TransformPositionsJob
            {
                transform = root.localToWorldMatrix,
                positions = Target.Runtime.RuntimeSprings.TailPositions,
                previousPositions = Target.Runtime.RuntimeSprings.TailPreviousPositions,
            }.Schedule(dependsOn);

        /************************************************************************************************************************/

        /// <inheritdoc/>
        public override IEnumerable<Transform> GetUsedBones()
        {
            yield return _SimulationRoot;
        }

        /************************************************************************************************************************/
    }
}

