using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Animations;
using UnityEngine.Playables;
using UnityEngine.Serialization;
#if ENABLE_INPUT_SYSTEM
using UnityEngine.InputSystem;
#endif

/// <summary>
/// Owns ALL character animation (the Animator is left empty): plays the active <see cref="PlayerPose"/>'s
/// idle or optional walk/run/fall clip through a two-input mixer <see cref="PlayableGraph"/> that
/// crossfades every clip switch. One key cycles poses, sprinting resets to the default pose; each pose
/// also carries its own eased CharacterController capsule and per-pose camera y-offsets.
/// </summary>
[RequireComponent(typeof(PlayerCharacterController))]
[DisallowMultipleComponent]
// Run after PlayerCharacterController (default order 0) so the sprint / locomotion state polled below is
// this frame's, and the per-pose capsule / camera offset resolve without a one-frame seam.
[DefaultExecutionOrder(100)]
public class PlayerPoseController : MonoBehaviour
{
    [Header("Poses")]
    [Tooltip("Selectable poses. Index 0 is the default pose that sprinting resets to.")]
    [SerializeField] private List<PlayerPose> poses = new();

    [Tooltip("Index the selection resets to while sprinting in any state (the default/standard pose).")]
    [Min(0)]
    [SerializeField] private int defaultPoseIndex = 0;

    [Header("Input")]
#if ENABLE_INPUT_SYSTEM
    [Tooltip("Key that cycles to the next pose (wraps at the end of the list).")]
    [SerializeField] private Key cycleKey = Key.C;
#else
    [Tooltip("Key that cycles to the next pose (wraps at the end of the list).")]
    [SerializeField] private KeyCode cycleKey = KeyCode.C;
#endif

    [Header("Transitions")]
    [Tooltip("Seconds to crossfade from the previous pose to a newly selected one. 0 = instant snap.")]
    [Min(0f)]
    [SerializeField] private float poseBlendDuration = 0.2f;

    [Tooltip("Approach speed (local units/sec) for easing the per-pose camera-target y-offset toward its " +
             "target while cycling poses. 0 = instant.")]
    [Min(0f)]
    [FormerlySerializedAs("ceilingOffsetSpeed")]
    [SerializeField] private float cameraOffsetSpeed = 4f;

    [Tooltip("Blend speed (0..1 per second) for resizing the CharacterController capsule toward the active " +
             "pose's dimensions when the pose changes. 0 = instant snap; higher = snappier resize.")]
    [Min(0f)]
    [SerializeField] private float capsuleBlendSpeed = 8f;

    [Tooltip("Apply Mecanim foot IK to the played humanoid clips. Off (default) plays each pose exactly as " +
             "authored. On bends the knees to plant the feet on the IK goal, which makes a straight-knee idle " +
             "render visibly bent even when the clip's Lower Leg Stretch muscle is fully straight.")]
    [SerializeField] private bool applyFootIK = false;

    private PlayerCharacterController _player;
    private Animator _animator;
    private string _graphName; // cached so building the graph allocates no string from the Update path

    private int _currentIndex;

    /// <summary>Locomotion category that selects which clip the active pose plays.</summary>
    public enum Locomotion { Idle, Walk, Run, Fall }

    // Network drive: a remote avatar's controllers are disabled (NetworkTransform owns the transform),
    // so its movement flags are dead - pose index + locomotion come from the owner via SetNetworkState.
    // _currentLocomotion is the category resolved this frame, published by the owner for replication.
    private bool _networkControlled;
    private Locomotion _networkLoco;
    private Locomotion _currentLocomotion;

    // Playables: a 2-input AnimationMixerPlayable that overrides the Animator Controller. Input 1 holds
    // the current ("to") pose; input 0 holds the outgoing ("from") pose during a crossfade and is empty
    // once the blend completes. Blending both inputs over a per-frame rest baseline (see below) makes
    // poses interpolate instead of snapping; an empty/partial "to" clip still resolves to rest.
    private PlayableGraph _graph;
    private AnimationPlayableOutput _output;
    private AnimationMixerPlayable _mixer;
    private AnimationClipPlayable _fromPlayable;
    private AnimationClipPlayable _toPlayable;
    private AnimationClip _toClip;        // the clip currently targeted (held on input 1)
    private bool _posePlaying;
    private float _blend;                 // 0..1 weight of the "to" input; 1 = fully on the target pose
    private bool _blending;

    private float _cameraPoseOffset;      // eased per-pose camera-target offset routed to the controller's camera target
    private float _firstPersonCameraOffset; // eased per-pose first-person eye offset, read by the FirstPersonController (never written to the shared camera target)

    // Capsule blend: the CharacterController dimensions are eased toward the active pose's via a single
    // normalized blend (so the capsule stays a valid shape throughout) and written through the controller.
    // Snapped on the first frame so the player spawns at its pose size with no resize blend.
    private bool _capsuleStarted;
    private float _capFromCenterY, _capFromRadius, _capFromHeight;
    private float _capToCenterY, _capToRadius, _capToHeight;
    private float _capCenterY, _capRadius, _capHeight;
    private float _capBlend;

    // Rest-pose snapshot of the Animator subtree, taken before the Animator first evaluates. A clip only
    // writes the bones its curves cover, so un-keyed bones would stay frozen on the previous animation;
    // restoring rest first makes every pose define the full body. The controller's meshOffsetTarget is
    // dropped from the set (see CaptureRestPose) so the wall-lean / ground raise is never wiped.
    private Transform[] _restBones;
    private Vector3[] _restPos;
    private Quaternion[] _restRot;
    private Vector3[] _restScale;

    public int CurrentIndex => _currentIndex;

    /// <summary>Locomotion category resolved this frame; the owner publishes it for network sync.</summary>
    public Locomotion CurrentLocomotion => _currentLocomotion;

    /// <summary>
    /// Switches this avatar to network-driven animation: local input, the sprint-reset, the per-pose
    /// collision capsule and camera framing are all skipped; pose index and locomotion are supplied by
    /// <see cref="SetNetworkState"/>. Called on a remote avatar by <see cref="PlayerAnimationSync"/>.
    /// </summary>
    public void EnableNetworkDrive() => _networkControlled = true;

    /// <summary>Pushes the replicated pose index and locomotion category onto a network-driven avatar.</summary>
    public void SetNetworkState(int poseIndex, Locomotion loco)
    {
        _currentIndex = WrapIndex(poseIndex);
        _networkLoco = loco;
    }

    /// <summary>Eased per-pose y-offset for the first-person eye, applied on top of the camera target by the FirstPersonController.</summary>
    public float FirstPersonCameraYOffset => _firstPersonCameraOffset;

    /// <summary>Holds the avatar on its current pose frame: the Playable graph stops advancing (its bones stay
    /// exactly where they are) and all per-frame locomotion / rest-reset / capsule / camera work is skipped.
    /// Set per-client for every still-in-play hider during the resolution reveal so nobody can wiggle in place.
    /// Clearing it resumes normal playback. Independent of network drive — works on the owner and on remotes.</summary>
    public bool Frozen { get; set; }

    void Awake()
    {
        _player = GetComponent<PlayerCharacterController>();
        _animator = _player ? _player.GetComponentInChildren<Animator>() : null;
        _graphName = $"{name}-Pose";
        if (_animator)
            CaptureRestPose();
        else
            Debug.LogWarning("[PlayerPoseController] No Animator on the PlayerCharacterController; poses cannot play.", this);
    }

    void Update()
    {
        // Resolution freeze: keep the graph playing (it re-asserts the frozen frame every render, even
        // under Animator culling) but zero the MIXER's speed so clip time stops advancing; skip everything
        // below so no locomotion/rest-reset moves a bone. The per-input clip speeds are untouched; the
        // mixer multiplier is restored on the first un-frozen frame just below.
        if (Frozen)
        {
            if (_graph.IsValid())
            {
                if (!_graph.IsPlaying()) _graph.Play();
                if (_mixer.IsValid()) PlayableExtensions.SetSpeed(_mixer, 0d);
            }
            return;
        }
        if (_mixer.IsValid() && PlayableExtensions.GetSpeed(_mixer) == 0d)
            PlayableExtensions.SetSpeed(_mixer, 1d);

        // Local control only - a remote's pose index is pushed in via SetNetworkState.
        if (!_networkControlled)
        {
            if (InputUtil.GameplayKeyPressed(cycleKey))
                NextPose();

            // Sprinting in any state (ground, air, wall, ceiling) snaps the selection back to the default
            // pose. Held continuously, so the reset shows up immediately, not only once back on the floor.
            if (_player.IsSprinting)
                _currentIndex = WrapIndex(defaultPoseIndex);
        }

        PlayerPose pose = ActivePose();

        // Resolve the locomotion category and its clip: locally from the controller's movement state, on
        // a remote the replicated value. ClipFor falls back to idle for unset optional clips, so a remote
        // (same pose asset) resolves to exactly the clip the owner shows.
        _currentLocomotion = _networkControlled ? _networkLoco : ResolveLocalLocomotion();
        AnimationClip clip = ClipFor(pose, _currentLocomotion, out float speed);
        PlayPose(clip, speed);

        // Collision capsule and camera framing are local-owner concerns: a remote avatar's controller and
        // CharacterController are disabled and no camera tracks it, so a network-driven avatar stops here.
        if (_networkControlled)
            return;

        // Per-pose collision capsule so a low pose can crawl where the standing capsule is too tall;
        // unset dimensions (<= 0) keep the default capsule.
        bool hasCapsule = pose && pose.ControllerHeight > 0f && pose.ControllerRadius > 0f;
        float capCenterY = hasCapsule ? pose.ControllerCenterY : _player.BaseCapsuleCenterY;
        float capRadius = hasCapsule ? pose.ControllerRadius : _player.BaseCapsuleRadius;
        float capHeight = hasCapsule ? pose.ControllerHeight : _player.BaseCapsuleHeight;
        ApplyCapsule(capCenterY, capRadius, capHeight);

        // Camera target: a per-pose framing offset (applied in any state), eased so cycling poses glides
        // the camera, routed to the ThirdPersonCamera's follow target only (never the body).
        float cameraPoseTarget = pose ? pose.CameraTargetYOffset : 0f;
        _cameraPoseOffset = EaseCameraOffset(_cameraPoseOffset, cameraPoseTarget);
        _player.SetCameraTargetLocalYOffset(_cameraPoseOffset);

        // Separate per-pose first-person eye offset - NOT routed to the shared camera target (that would
        // also move the ThirdPersonCamera); the FirstPersonController reads it directly.
        float firstPersonTarget = pose ? pose.FirstPersonCameraTargetYOffset : 0f;
        _firstPersonCameraOffset = EaseCameraOffset(_firstPersonCameraOffset, firstPersonTarget);
    }

    private float EaseCameraOffset(float current, float target)
    {
        return cameraOffsetSpeed > 0f
            ? Mathf.MoveTowards(current, target, cameraOffsetSpeed * Time.deltaTime)
            : target;
    }

    void OnDisable()
    {
        // Hand the rig back to the Animator Controller and never leak a graph; rebuilt lazily on the
        // next idle frame. Also clear the camera offset and restore the default capsule we were holding.
        ReleasePose();
        _cameraPoseOffset = 0f;
        _firstPersonCameraOffset = 0f;
        if (_player)
        {
            _player.SetCameraTargetLocalYOffset(0f);
            _player.SetCapsule(_player.BaseCapsuleCenterY, _player.BaseCapsuleRadius, _player.BaseCapsuleHeight);
        }
        _capsuleStarted = false; // re-snap to the active pose's capsule if this is re-enabled
    }

    /// <summary>Advances to the next pose, wrapping at the end of the list.</summary>
    public void NextPose() => SetPoseIndex(_currentIndex + 1);

    /// <summary>Selects a pose by index, wrapping into range so cycling never goes out of bounds.</summary>
    public void SetPoseIndex(int index) => _currentIndex = WrapIndex(index);

    private int WrapIndex(int index)
    {
        int count = poses.Count;
        if (count == 0) return 0;
        return ((index % count) + count) % count;
    }

    private PlayerPose ActivePose()
    {
        if (poses.Count == 0) return null;
        return poses[Mathf.Clamp(_currentIndex, 0, poses.Count - 1)];
    }

    private Locomotion ResolveLocalLocomotion()
    {
        if (_player.IsMovingOnGround)
            return _player.IsSprinting ? Locomotion.Run : Locomotion.Walk;
        if (_player.IsFalling)
            return Locomotion.Fall;
        return Locomotion.Idle;
    }

    /// <summary>
    /// Returns the clip (and its <paramref name="speed"/>) for <paramref name="pose"/> in the given
    /// locomotion category. The walk/run/fall clips are optional: an unset one falls back to the pose's
    /// idle clip, so an unauthored category keeps the held-pose behaviour identically on the owner and on
    /// a network-driven remote avatar. A null pose hands the rig back (no clip).
    /// </summary>
    private AnimationClip ClipFor(PlayerPose pose, Locomotion loco, out float speed)
    {
        if (pose)
        {
            switch (loco)
            {
                case Locomotion.Walk when pose.WalkClip:
                    speed = pose.WalkClipSpeed;
                    return pose.WalkClip;
                case Locomotion.Run when pose.RunClip:
                    speed = pose.RunClipSpeed;
                    return pose.RunClip;
                case Locomotion.Fall when pose.FallClip:
                    speed = pose.FallClipSpeed;
                    return pose.FallClip;
            }
            speed = pose.ClipSpeed;
            return pose.Clip;
        }
        speed = 1f;
        return null;
    }

    /// <summary>
    /// Eases the CharacterController capsule toward the given target dimensions and writes them through
    /// the controller. A single normalized blend drives all three values together, so the capsule stays a
    /// valid shape (and its floor contact moves linearly) throughout the resize. The first call snaps, so
    /// the player spawns at its pose size without a resize blend; <see cref="capsuleBlendSpeed"/> 0 also snaps.
    /// </summary>
    private void ApplyCapsule(float centerY, float radius, float height)
    {
        if (!_capsuleStarted)
        {
            _capsuleStarted = true;
            _capFromCenterY = _capToCenterY = _capCenterY = centerY;
            _capFromRadius = _capToRadius = _capRadius = radius;
            _capFromHeight = _capToHeight = _capHeight = height;
            _capBlend = 1f;
        }
        else if (centerY != _capToCenterY || radius != _capToRadius || height != _capToHeight)
        {
            // New target (pose changed): ease from wherever the capsule currently is.
            _capFromCenterY = _capCenterY;
            _capFromRadius = _capRadius;
            _capFromHeight = _capHeight;
            _capToCenterY = centerY;
            _capToRadius = radius;
            _capToHeight = height;
            _capBlend = 0f;
        }

        _capBlend = capsuleBlendSpeed > 0f
            ? Mathf.MoveTowards(_capBlend, 1f, capsuleBlendSpeed * Time.deltaTime)
            : 1f;
        _capCenterY = Mathf.Lerp(_capFromCenterY, _capToCenterY, _capBlend);
        _capRadius = Mathf.Lerp(_capFromRadius, _capToRadius, _capBlend);
        _capHeight = Mathf.Lerp(_capFromHeight, _capToHeight, _capBlend);
        _player.SetCapsule(_capCenterY, _capRadius, _capHeight);
    }

    /// <summary>
    /// Plays and holds the given clip, overriding the Animator Controller. Switching to a different
    /// clip crossfades from the current pose over <see cref="poseBlendDuration"/> (the first pose, or a
    /// zero duration, appears instantly). Called every frame: it starts a transition on a clip change,
    /// advances any in-flight blend, applies <paramref name="speed"/> to the target playback, and
    /// re-snaps un-keyed bones to rest before the graph evaluates. A null clip (or no animator) releases
    /// back to the controller.
    /// </summary>
    private void PlayPose(AnimationClip clip, float speed)
    {
        if (!clip || !_animator)
        {
            ReleasePose();
            return;
        }

        EnsureGraph();

        if (!(_posePlaying && _toClip == clip))
            BeginTransition(clip);

        AdvanceBlend(Time.deltaTime);

        // Set every frame so live speed-slider tuning takes effect immediately; the outgoing pose on
        // input 0 keeps the speed it carried, fading out at its own pace.
        if (_toPlayable.IsValid())
            PlayableExtensions.SetSpeed(_toPlayable, speed);

        // Snap un-keyed bones to rest every frame so each mixer input blends relative to the same
        // baseline; the graph then overrides the keyed bones when it evaluates later this frame.
        ResetToRestPose();

        if (!_graph.IsPlaying())
            _graph.Play();

        _posePlaying = true;
    }

    private void EnsureGraph()
    {
        if (_graph.IsValid())
            return;
        _graph = PlayableGraph.Create(_graphName);
        _graph.SetTimeUpdateMode(DirectorUpdateMode.GameTime);
        _mixer = AnimationMixerPlayable.Create(_graph, 2);
        _output = AnimationPlayableOutput.Create(_graph, "Pose", _animator);
        _output.SetSourcePlayable(_mixer);
    }

    /// <summary>
    /// Starts a crossfade to <paramref name="clip"/>: the in-flight target becomes the outgoing pose on
    /// input 0 (carrying its playback time), the new clip goes on input 1, and the blend ramps up from
    /// 0. The first pose (or a zero duration) snaps in at full weight with no fade-from-nothing.
    /// </summary>
    private void BeginTransition(AnimationClip clip)
    {
        bool hadPose = _posePlaying && _toPlayable.IsValid();

        // Drop the previous outgoing pose (already faded out, or none) before reusing input 0.
        if (_fromPlayable.IsValid())
        {
            _mixer.DisconnectInput(0);
            _graph.DestroyPlayable(_fromPlayable);
            _fromPlayable = default;
        }

        if (hadPose)
        {
            // Move the current target from input 1 to input 0 so it becomes the outgoing pose without a
            // recreate (its animation time carries over, no jump).
            _mixer.DisconnectInput(1);
            _fromPlayable = _toPlayable;
            _mixer.ConnectInput(0, _fromPlayable, 0);
        }

        _toPlayable = AnimationClipPlayable.Create(_graph, clip);
        // Off by default: foot IK plants the foot on its goal and bends the knee to reach it, so a straight-knee
        // pose renders bent (SampleAnimation, which the authoring tools check against, never applies it). The
        // outgoing pose on input 0 is a former _toPlayable, so it already carries this setting.
        _toPlayable.SetApplyFootIK(applyFootIK);
        _mixer.ConnectInput(1, _toPlayable, 0);
        _toClip = clip;

        _blend = (hadPose && poseBlendDuration > 0f) ? 0f : 1f;
        _blending = _blend < 1f;
        ApplyBlendWeights();
    }

    private void AdvanceBlend(float dt)
    {
        if (!_blending)
            return;
        _blend = poseBlendDuration > 0f ? Mathf.MoveTowards(_blend, 1f, dt / poseBlendDuration) : 1f;
        if (_blend >= 1f)
        {
            _blend = 1f;
            _blending = false;
            if (_fromPlayable.IsValid())
            {
                _mixer.DisconnectInput(0);
                _graph.DestroyPlayable(_fromPlayable);
                _fromPlayable = default;
            }
        }
        ApplyBlendWeights();
    }

    private void ApplyBlendWeights()
    {
        if (!_mixer.IsValid())
            return;
        bool hasFrom = _fromPlayable.IsValid();
        float w = hasFrom ? Mathf.SmoothStep(0f, 1f, _blend) : 1f;
        _mixer.SetInputWeight(0, hasFrom ? 1f - w : 0f);
        _mixer.SetInputWeight(1, _toPlayable.IsValid() ? w : 0f);
    }

    /// <summary>
    /// Snapshots every transform a pose clip can address (the Animator's whole subtree) in its
    /// authored rest pose. Taken in Awake before the Animator first evaluates, so it captures the rig
    /// at rest (the paintman rig's rest pose is its T-pose).
    /// </summary>
    private void CaptureRestPose()
    {
        _restBones = _animator.transform.GetComponentsInChildren<Transform>(true);
        int n = _restBones.Length;
        _restPos = new Vector3[n];
        _restRot = new Quaternion[n];
        _restScale = new Vector3[n];
        for (int i = 0; i < n; i++)
        {
            Transform t = _restBones[i];
            _restPos[i] = t.localPosition;
            _restRot[i] = t.localRotation;
            _restScale[i] = t.localScale;
        }

        // Drop the model-offset transform from the reset set: the wall-lean / ground raise (SetMeshLean)
        // is its sole writer, and with the humanoid rig the Animator shares this node - resetting it each
        // frame was exactly what wiped the offsets. Nulling the slot makes ResetToRestPose skip it.
        Transform offsetTarget = _player ? _player.MeshOffsetTarget : null;
        if (offsetTarget)
        {
            int idx = System.Array.IndexOf(_restBones, offsetTarget);
            if (idx >= 0)
                _restBones[idx] = null;
        }
    }

    /// <summary>Restores the captured rest pose so an empty or partial pose clip fully defines the body.</summary>
    private void ResetToRestPose()
    {
        if (_restBones == null)
            return;
        for (int i = 0; i < _restBones.Length; i++)
        {
            Transform t = _restBones[i];
            if (!t)
                continue;
            t.localPosition = _restPos[i];
            t.localRotation = _restRot[i];
            t.localScale = _restScale[i];
        }
    }

    /// <summary>Destroys the graph so the Animator Controller (run/idle) drives the rig again.</summary>
    private void ReleasePose()
    {
        if (!_posePlaying)
            return;
        if (_graph.IsValid())
            _graph.Destroy();
        _graph = default;
        _output = default;
        _mixer = default;
        _fromPlayable = default;
        _toPlayable = default;
        _toClip = null;
        _blend = 0f;
        _blending = false;
        _posePlaying = false;
    }
}