#ifndef ROBOT_H
#define ROBOT_H

#include <kdl/chain.hpp>
#include <kdl/chainfksolverpos_recursive.hpp>
#include <kdl/chainiksolverpos_lma.hpp>
#include <kdl/chainiksolverpos_nr.hpp>
#include <kdl/chainiksolverpos_nr_jl.hpp>
#include <kdl/chainiksolvervel_pinv.hpp>
#include <kdl/frames.hpp>
#include <kdl/jntarray.hpp>
#include <kdl/tree.hpp>

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

#include "utils.h"

class Robot
{
private:
    KDL::Chain kinematicChain;
    KDL::ChainFkSolverPos_recursive *fkSolver;
    KDL::ChainIkSolverVel_pinv *ikVelSolver;
    KDL::ChainIkSolverPos_NR_JL *ikSolverNR;
    KDL::ChainIkSolverPos_LMA *ikSolverLMA;
    bool m_initialized;
    KDL::JntArray jointLowerLimits;
    KDL::JntArray jointUpperLimits;
    std::vector<std::string> activeJointNames;
    bool m_hasJointLimits;

    // 记录 joint 名称到 child link UUID 的映射，便于前端回写对象姿态。
    std::unordered_map<std::string, std::string> jointChildLinkUuidMap;

    // 内部工具方法。
    double quaternionAngleDifference(const double q1[4], const double q2[4]);
    int calculateAutoSteps(const double startPose[7], const double endPose[7],
                           int minSteps, int maxSteps,
                           double positionResolution, double orientationResolution);
    bool parseJointChildLinkUuidsFromUrdf(const std::string &urdfString);
    // 自动收集 KDL Tree 中没有子节点的末端 link 名称。
    std::vector<std::string> collectLeafSegmentNames(const KDL::Tree &tree) const;
    // 自动选择可用于正逆运动学的 6 轴串联链，避免固定依赖 base/tool0 命名。
    bool selectKinematicChain(const KDL::Tree &tree, KDL::Chain &selectedChain) const;
    // 从 URDF 中提取当前运动学链的关节上下限。
    bool parseJointLimitsFromUrdf(const std::string &urdfString);
    // 检查关节数组是否处于 URDF 定义的上下限范围内。
    bool validateJointLimits(const double joints[6], const std::string &context) const;
    // 检查关节向量是否处于 URDF 定义的上下限范围内。
    bool validateJointLimits(const std::vector<double> &joints, const std::string &context) const;

public:
    /**
     * @brief 默认构造函数。
     */
    Robot();

    /**
     * @brief 使用 URDF 字符串直接初始化机器人。
     * @param urdfString URDF 格式字符串。
     */
    explicit Robot(const std::string &urdfString);

    /**
     * @brief 释放运动学求解器资源。
     */
    ~Robot();

    int numberOfJoints;

    /**
     * @brief 根据 URDF 字符串初始化机器人运动学链。
     * @param urdfString URDF 格式字符串。
     * @return 初始化成功返回 true，否则返回 false。
     */
    bool initRobot(const std::string &urdfString);

    /**
     * @brief 根据关节序号获取 child link UUID。
     * @param jointIndex 关节序号，从 1 开始。
     * @return 找到时返回对应 UUID，未找到返回空字符串。
     */
    std::string getJointUuidByIndex(int jointIndex) const;

    /**
     * @brief 根据关节序号生成标准关节名称。
     * @param jointIndex 关节序号，从 1 开始。
     * @return 如 `joint_1`、`joint_2` 这样的关节名称。
     */
    static std::string getJointNameByIndex(int jointIndex);

    /**
     * @brief 获取当前已解析到的全部关节序号列表。
     * @return 已排序且去重后的关节序号集合。
     */
    std::vector<int> getAvailableJointIndices() const;

    /**
     * @brief 兼容多种命名格式查找关节 UUID。
     * @param jointIndex 关节序号，从 1 开始。
     * @return 找到时返回 UUID，未找到返回空字符串。
     */
    std::string findJointUuidByAnyFormat(int jointIndex) const;

    /**
     * @brief 对姿态序列做逆解，自动进行轨迹插补。
     */
    json inversePoseStr(const std::string &pose_str, const std::string &q_init_str);

    /**
     * @brief 对两点姿态按指定步数插补后做逆解。
     */
    json inversePoseStr2PSteps(const std::string &pose_str, const std::string &q_init_str, const std::int32_t steps);

    /**
     * @brief 直接对输入姿态点做逆解，不做相邻点插补。
     */
    json inversePoseStrNoDifference(const std::string &pose_str, const std::string &q_init_str);

    /**
     * @brief 解析姿态字符串，格式为 `x,y,z,qx,qy,qz,qw;...`。
     */
    std::vector<std::vector<double>> parsePoseString(const std::string &pose_str);

    /**
     * @brief 解析关节列表字符串，格式为 `j1,j2,j3,j4,j5,j6;...`。
     */
    std::vector<std::vector<double>> parseJointListString(const std::string &pose_str);

    std::unordered_map<std::string, std::string> getJointChildLinkUuidMap() const { return jointChildLinkUuidMap; }
    std::string getJointChildLinkUuid(const std::string &jointName) const;
    std::vector<double> parseJointString(const std::string &joint_str);

    /**
     * @brief 计算所有关节位姿，并按前端需要的 objStates 格式组织数据。
     */
    json kinematicsForwardAllJointsList(std::vector<double> joints);

    /**
     * @brief 批量计算多帧关节正解，输出 OPERATION 结构。
     */
    json handleKinematicsForwardAllJoints(const std::string &joints_str);

    /**
     * @brief 兼容旧格式的全部关节正解输出。
     */
    json handleKinematicsForwardAllJoints_objStates(const std::string &joints_str);

    /**
     * @brief 判断机器人是否已初始化。
     * @return 已初始化返回 true。
     */
    bool isInitialized() const { return m_initialized; }

    /**
     * @brief 使用 NR 求解器计算逆运动学。
     * @param pose 目标位姿，格式为 `[x, y, z, qx, qy, qz, qw]`。
     * @param iniJ 初始关节角，格式为 `[j1, ..., j6]`。
     * @param resultJoints 输出关节角结果。
     * @return 求解成功返回 true，否则返回 false。
     */
    bool calculateIK_NR(const double pose[7], const double iniJ[6], double resultJoints[6]);

    /**
     * @brief 使用 LMA 求解器计算逆运动学。
     * @param pose 目标位姿，格式为 `[x, y, z, qx, qy, qz, qw]`。
     * @param iniJ 初始关节角，格式为 `[j1, ..., j6]`。
     * @param resultJoints 输出关节角结果。
     * @param eps 迭代收敛精度。
     * @param maxiter 最大迭代次数。
     * @param eps_joints 关节收敛阈值。
     * @return 求解成功返回 true，否则返回 false。
     */
    bool calculateIK_LMA(const double pose[7], const double iniJ[6], double resultJoints[6],
                         double eps = 1e-8, int maxiter = 3000, double eps_joints = 1e-12);

    /**
     * @brief 计算 TCP 的正运动学结果。
     * @param joints 输入关节角，格式为 `[j1, ..., j6]`。
     * @param tcpPose 输出 TCP 位姿，格式为 `[x, y, z, qx, qy, qz, qw]`。
     * @return 计算成功返回 true，否则返回 false。
     */
    bool calculateFK_TCP(const double joints[6], double tcpPose[7]);

    /**
     * @brief 计算每个关节节点的正运动学结果。
     * @param joints 输入关节角，格式为 `[j1, ..., j6]`。
     * @param jointPoses 输出全部关节位姿，共 6 组，每组 7 个值。
     * @return 计算成功返回 true，否则返回 false。
     */
    bool calculateFK_AllJointsforwardKinematics(const double joints[6], double jointPoses[42]);
    bool forwardKinematics(const double joints[6], double jointPoses[42]);

    /**
     * @brief 获取当前 KDL 运动学链。
     */
    const KDL::Chain &getKinematicChain() const { return kinematicChain; }

    /**
     * @brief 获取机器人关节数量。
     */
    int getNumberOfJoints() const { return kinematicChain.getNrOfJoints(); }

    /**
     * @brief 生成姿态轨迹，输出为固定数组。
     * @param pose1 起点位姿。
     * @param pose2 终点位姿。
     * @param outputPoses 输出轨迹数组。
     * @param npoint 轨迹点数量，传 0 表示自动计算。
     * @return 实际生成的轨迹点数量，失败返回 -1。
     */
    int trajectoryPlanning(const double pose1[7], const double pose2[7],
                           double outputPoses[][7], int npoint = 0);

    /**
     * @brief 生成姿态轨迹，输出为二维向量。
     * @param pose1 起点位姿。
     * @param pose2 终点位姿。
     * @param outputPoses 输出轨迹点集合。
     * @param npoint 轨迹点数量，传 0 表示自动计算。
     * @return 实际生成的轨迹点数量，失败返回 -1。
     */
    int trajectoryPlanning(const double pose1[7], const double pose2[7],
                           std::vector<std::vector<double>> &outputPoses,
                           int npoint = 0);

    /**
     * @brief 使用 `std::vector<double>` 形式输入位姿生成轨迹。
     * @param pose1 起点位姿。
     * @param pose2 终点位姿。
     * @param outputPoses 输出轨迹点集合。
     * @param npoint 轨迹点数量，传 0 表示自动计算。
     * @return 实际生成的轨迹点数量，失败返回 -1。
     */
    int trajectoryPlanning(const std::vector<double> &pose1, const std::vector<double> &pose2,
                           std::vector<std::vector<double>> &outputPoses, int npoint = 0);

    int calculateAutoSteps(const std::vector<double> &startPose, const std::vector<double> &endPose,
                           int minSteps, int maxSteps,
                           double positionResolution, double orientationResolution);
    double quaternionAngleDifference(const std::vector<double> &q1, const std::vector<double> &q2);

    /**
     * @brief 使用 LMA 求解器计算逆解，输入为数组。
     */
    std::vector<double> inverse(const double pose[7], const double iniJ[6],
                                double eps = 1e-8, int maxiter = 3000,
                                double eps_joints = 1e-12);

    /**
     * @brief 使用 LMA 求解器计算逆解，输入为向量。
     */
    std::vector<double> inverse(const std::vector<double> &pose,
                                const std::vector<double> &iniJ,
                                double eps = 1e-8, int maxiter = 3000,
                                double eps_joints = 1e-12);

    /**
     * @brief 使用默认零位初值计算逆解，输入为数组。
     */
    std::vector<double> inverse(const double pose[7],
                                double eps = 1e-8, int maxiter = 3000,
                                double eps_joints = 1e-12);

    /**
     * @brief 使用默认零位初值计算逆解，输入为向量。
     */
    std::vector<double> inverse(const std::vector<double> &pose,
                                double eps = 1e-8, int maxiter = 3000,
                                double eps_joints = 1e-12);
};

#endif // ROBOT_H