using System;
using System.Collections;
using System.Collections.Generic;
using System.Threading.Tasks;
using LibMMD.Model;
using LibMMD.Motion;
using LibMMD.Util;
using UnityEngine;

namespace LibMMD.Unity3D.Animation
{
    // Bakes the kinematic (non-physics) bone FK evaluation of a VMD motion into plain
    // position/rotation samples, by sampling a throwaway Poser/MotionPlayer pair at a fixed
    // rate - no physics stepping involved. Physics-driven bones keep being computed live every
    // frame (see MmdAnimator); this only replaces the per-frame FK/IK bone-tree walk.
    //
    // Deliberately does not go through an AnimationClip/AnimationClipPlayable: AnimationClip.
    // SetCurve only works at runtime (in a Player build) on legacy clips, but
    // AnimationClipPlayable refuses legacy clips - there is no clip configuration that
    // satisfies both, so a runtime-built clip can never actually play back bone curves in a
    // build (it silently no-ops in the SetCurve case, or throws in the legacy case). Baked
    // samples are consumed directly by MmdAnimator.ApplyBakedPoseToPoser instead.
    internal static class MmdAnimationClipBaker
    {
        // Increment whenever pose sampling semantics change. This is part of the in-memory
        // cache key, preventing an old bake from masking IK/FK solver fixes until Unity exits.
        internal const int BakeAlgorithmVersion = 38;
        private const float DefaultSampleRate = 30f;

        public static BakedPose Bake(MmdModel model, MmdMotion motion, bool[] bakeableBoneMask,
            bool enableIk = true, float sampleRate = DefaultSampleRate)
        {
            return SampleFrames(model, motion, bakeableBoneMask, enableIk, sampleRate);
        }

        // Coroutine equivalent of Bake(): runs the per-frame IK/FK sampling (the expensive
        // part - O(bones * sampledFrames)) on a background thread via Task.Run instead of
        // blocking the calling (main) thread for the whole motion length. Drive with
        // `yield return BakeAsync(...)` from a MonoBehaviour coroutine; onComplete is invoked
        // with the finished result just before the coroutine ends.
        public static IEnumerator BakeAsync(MmdModel model, MmdMotion motion, bool[] bakeableBoneMask,
            System.Action<BakedPose> onComplete, bool enableIk = true, float sampleRate = DefaultSampleRate)
        {
            var task = Task.Run(() => SampleFrames(model, motion, bakeableBoneMask, enableIk, sampleRate));
            while (!task.IsCompleted)
            {
                yield return null;
            }
            onComplete(task.Result); // rethrows on the calling thread if SampleFrames faulted
        }

        // Per-bone (model.Bones.Length) samples of the baked FK pose, in the same model-space
        // Position/Rotation convention as BonePosePreCalculator.GetBonePoseImage (derived from
        // BoneImage.SkinningMatrix). Positions/Rotations entries are null for non-baked bones.
        // Plain data, no UnityEngine.Object involved - safe to build off the main thread and to
        // cache indefinitely without any Destroy()/lifecycle bookkeeping.
        internal sealed class BakedPose
        {
            public float[] Times;
            public float SampleRate;
            public Vector3[][] Positions;
            public Quaternion[][] Rotations;

            // Interpolates bone boneIndex's baked pose at time. Returns false (no output) for a
            // bone that wasn't baked (bakeableBoneMask[boneIndex] was false).
            public bool TrySample(int boneIndex, double time, out Vector3 position, out Quaternion rotation)
            {
                var positions = Positions[boneIndex];
                if (positions == null)
                {
                    position = default;
                    rotation = default;
                    return false;
                }
                var rotations = Rotations[boneIndex];
                var last = Times.Length - 1;
                // Samples are generated at a fixed rate. Derive the frame directly instead
                // of scanning Times from zero for every baked bone on every rendered frame.
                var samplePosition = (float) time * SampleRate;
                var i0 = Mathf.Clamp(Mathf.FloorToInt(samplePosition), 0, last);
                var i1 = i0 < last ? i0 + 1 : i0;
                var span = Times[i1] - Times[i0];
                var t = span > 1e-6f ? Mathf.Clamp01((float) (time - Times[i0]) / span) : 0f;
                position = Vector3.Lerp(positions[i0], positions[i1], t);
                rotation = Quaternion.Slerp(rotations[i0], rotations[i1], t);
                return true;
            }
        }

        private static BakedPose SampleFrames(MmdModel model, MmdMotion motion, bool[] bakeableBoneMask,
            bool enableIk, float sampleRate)
        {
            var motionLength = motion.Length / 30.0;
            var sampleCount = Mathf.Max(2, Mathf.FloorToInt((float) (motionLength * sampleRate)) + 1);

            var boneCount = model.Bones.Length;
            var positions = new Vector3[boneCount][];
            var rotations = new Quaternion[boneCount][];
            var bakedBoneIndexes = new List<int>(boneCount);
            for (var i = 0; i < boneCount; ++i)
            {
                if (!bakeableBoneMask[i]) continue;
                bakedBoneIndexes.Add(i);
                positions[i] = new Vector3[sampleCount];
                rotations[i] = new Quaternion[sampleCount];
            }

            var times = new float[sampleCount];
            // A Poser is mutable, so each worker owns an independent instance. Frames are
            // otherwise independent because PrePhysicsPosing resets all transient IK state.
            // Cap workers to avoid multiplying large per-model pose buffers on high-core CPUs.
            var parallelOptions = new ParallelOptions
            {
                MaxDegreeOfParallelism = Math.Max(1, Math.Min(Environment.ProcessorCount, 8))
            };
            Parallel.For(0, sampleCount, parallelOptions,
                () =>
                {
                    var workerPoser = new Poser(model) { EnableIk = enableIk };
                    return new BakeWorker(workerPoser, new MotionPlayer(motion, workerPoser));
                },
                (f, _, worker) =>
                {
                    var t = f / sampleRate;
                    if (t > motionLength) t = (float) motionLength;
                    times[f] = t;

                    worker.MotionPlayer.SeekTime(t);
                    worker.Poser.PrePhysicsPosing();

                    for (var j = 0; j < bakedBoneIndexes.Count; ++j)
                    {
                        var i = bakedBoneIndexes[j];
                        var image = worker.Poser.BoneImages[i];
                        positions[i][f] = image.SkinningMatrix.MultiplyPoint3x4(model.Bones[i].Position);
                        rotations[i][f] = image.SkinningMatrix.ExtractRotation();
                    }
                    return worker;
                },
                _ => { });

            return new BakedPose
            {
                Times = times,
                SampleRate = sampleRate,
                Positions = positions,
                Rotations = rotations
            };
        }

        private sealed class BakeWorker
        {
            internal readonly Poser Poser;
            internal readonly MotionPlayer MotionPlayer;

            internal BakeWorker(Poser poser, MotionPlayer motionPlayer)
            {
                Poser = poser;
                MotionPlayer = motionPlayer;
            }
        }
    }
}