using System;
using LibMMD.Unity3D.BonePose;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Animations;
using UnityEngine.Playables;

namespace LibMMD.Unity3D.Animation
{
    // Owns the PlayableGraph that drives bone Transforms through the Animator
    // component. Callers push already-computed BonePoseImage[] (model-space,
    // pre-root-transform, same convention as BonePosePreCalculator.GetBonePoseImage)
    // and explicitly Evaluate() to apply them synchronously - there is no automatic
    // per-frame graph evaluation, so MotionPos-style immediate seeks stay synchronous.
    internal class MmdAnimationDriver : IDisposable
    {
        private PlayableGraph _graph;
        private AnimationScriptPlayable _playable;
        private NativeArray<TransformStreamHandle> _handles;
        private NativeArray<BonePoseTRS> _poses;
        private NativeArray<bool> _writeFlags;
        private bool _built;
        public void Build(Animator animator, Transform[] bones, bool[] skipMask = null)
        {
            Dispose();

            _graph = PlayableGraph.Create("MmdAnimationGraph");
            _graph.SetTimeUpdateMode(DirectorUpdateMode.Manual);

            _handles = new NativeArray<TransformStreamHandle>(bones.Length, Allocator.Persistent);
            _poses = new NativeArray<BonePoseTRS>(bones.Length, Allocator.Persistent);
            _writeFlags = new NativeArray<bool>(bones.Length, Allocator.Persistent);
            for (var i = 0; i < bones.Length; ++i)
            {
                // Unity physics reparents dynamic chains outside the Animator hierarchy.
                // Binding those transforms logs an error even when the animation job never
                // writes the resulting handle, so leave their handles unbound from the start.
                if (skipMask != null && i < skipMask.Length && skipMask[i]) continue;
                _handles[i] = animator.BindStreamTransform(bones[i]);
            }

            var job = new MmdPoseAnimationJob
            {
                Handles = _handles,
                Poses = _poses,
                WriteFlags = _writeFlags
            };
            _playable = AnimationScriptPlayable.Create(_graph, job);

            var output = AnimationPlayableOutput.Create(_graph, "MmdAnimationOutput", animator);
            output.SetSourcePlayable(_playable);
            _graph.Play();
            _built = true;
        }

        // poses: model-space Position/Rotation per bone (same convention as
        // BonePosePreCalculator.GetBonePoseImage / the file-storage playback path).
        // skipMask: true for bones a Unity Rigidbody/ConfigurableJoint currently
        // controls (PhysicsMode.Unity) - those Transforms are left untouched.
        public void SetBonePoses(BonePoseImage[] poses, Transform rootTransform, bool[] skipMask = null)
        {
            if (!_built)
            {
                throw new InvalidOperationException("MmdAnimationDriver.Build must be called first");
            }
            if (poses.Length != _poses.Length)
            {
                throw new ArgumentException("poses.Length does not match bone count the driver was built with");
            }
            for (var i = 0; i < poses.Length; ++i)
            {
                var skip = skipMask != null && skipMask[i];
                _writeFlags[i] = !skip;
                if (skip)
                {
                    continue;
                }
                _poses[i] = new BonePoseTRS
                {
                    Position = rootTransform.TransformPoint(poses[i].Position),
                    Rotation = rootTransform.rotation * poses[i].Rotation
                };
            }
        }

        public void Evaluate()
        {
            if (_built)
            {
                _graph.Evaluate(0f);
            }
        }

        public void Dispose()
        {
            if (_built)
            {
                _graph.Destroy();
                _built = false;
            }
            if (_handles.IsCreated)
            {
                _handles.Dispose();
            }
            if (_poses.IsCreated)
            {
                _poses.Dispose();
            }
            if (_writeFlags.IsCreated)
            {
                _writeFlags.Dispose();
            }
        }
    }
}