#pragma once

#include <string>
#include <unordered_map>
#include <vector>

#include <osg/ref_ptr>
#include <osg/Quat>
#include <osg/Vec3>
#include <osg/Matrix>
#include <osg/NodeVisitor>
#include <osgAnimation/Bone>
#include <osgAnimation/Skeleton>

/**
 * PoseManager
 * -----------
 * Owns the current pose as a map of bone name → local rotation + translation
 * offsets from the bind pose. Applies them to the osgAnimation bone tree
 * each frame during the update traversal.
 *
 * Usage:
 *   // At load time:
 *   mgr.init(skeleton, boneMap);
 *
 *   // From ImGui / bone gizmo:
 *   mgr.setBoneRotation("Head", osg::Quat(...));
 *
 *   // In update callback:
 *   mgr.applyPose();
 *
 *   // Reset:
 *   mgr.resetBone("Head");
 *   mgr.resetAll();
 */
class PoseManager {
public:
    struct BonePose {
        osg::Quat  rotation    = { 0, 0, 0, 1 };  // local rotation delta
        osg::Vec3  translation = { 0, 0, 0 };      // local translation delta
        bool       modified    = false;
    };

    PoseManager() = default;

    /// Called once after skeleton is loaded.
    void init(osgAnimation::Skeleton* skeleton,
              const std::unordered_map<std::string,
                  osg::ref_ptr<osgAnimation::Bone>>& boneMap,
              const std::unordered_map<std::string, osg::Vec3>& bindPositions = {});

    /// Sorted list of all bone names (for UI tree).
    const std::vector<std::string>& boneNames() const { return m_boneNames; }

    /// Bone hierarchy as parent name → child names (for tree display).
    const std::unordered_map<std::string, std::vector<std::string>>&
        boneChildren() const { return m_children; }

    std::string boneParent(const std::string& name) const;

    // ── Pose setters ──────────────────────────────────────────────────────────
    void setBoneRotation(const std::string& name, const osg::Quat& q);
    void setBoneTranslation(const std::string& name, const osg::Vec3& t);
    void resetBone(const std::string& name);
    void resetAll();

    // ── Pose getters ──────────────────────────────────────────────────────────
    const BonePose& getBonePose(const std::string& name) const;
    osg::Vec3  getBoneWorldPos(const std::string& name) const;

    // ── Per-frame update ──────────────────────────────────────────────────────
    /// Apply current pose to the osgAnimation bone tree.
    void applyPose();

    /// Run a skeleton update traversal to refresh world-space bone matrices.
    /// Call every frame so getBoneWorldPos() returns current positions.
    void updateSkeletonMatrices();

    bool isInitialized() const { return m_initialized; }

private:
    bool         m_initialized    = false;
    unsigned int m_traversalCount  = 1000;

    osg::ref_ptr<osgAnimation::Skeleton>  m_skeleton;
    std::unordered_map<std::string, osg::ref_ptr<osgAnimation::Bone>> m_boneMap;
    std::unordered_map<std::string, BonePose>       m_poses;
    std::unordered_map<std::string, osg::Matrix>    m_bindMatrices;       // local bind pose
    std::unordered_map<std::string, osg::Vec3>     m_bindWorldPositions;  // world pos from offset matrix
    std::unordered_map<std::string, std::string>    m_parents;      // child → parent
    std::unordered_map<std::string, std::vector<std::string>> m_children;
    std::vector<std::string>                        m_boneNames;    // sorted
    BonePose                                        m_defaultPose;  // returned for unknowns

    void buildHierarchy(osgAnimation::Bone* bone, const std::string& parentName);
};