#include "arm_driver.hpp"
#include "rm_driver_kinematics.hpp"
#include <cstring>
#include <Eigen/Dense>

using namespace Eigen;

#define MIN_MOVING_TIME_STEP 80

RmArmDriverAndKinematics::RmArmDriverAndKinematics()
{
    rmService_ = nullptr;
    initFlag_ = RM_DRIVER_KINEMATICS_UNINIT;
}

RmArmDriverAndKinematics::~RmArmDriverAndKinematics()
{
    if (rmService_ != nullptr) {
        rmService_->rm_destroy();
        delete rmService_;
    }
}

RmArmDriverAndKinematics::RmArmStatusCallbackFun RmArmDriverAndKinematics::armStatusCallback_ = nullptr;
RmArmDriverAndKinematics::ApiLogFunc RmArmDriverAndKinematics::customApiLog_ = nullptr;
rm_robot_handle *RmArmDriverAndKinematics::leftRobotHandle_;
rm_robot_handle *RmArmDriverAndKinematics::rightRobotHandle_;
RM_Service *RmArmDriverAndKinematics::rmService_ = nullptr;

void RmArmDriverAndKinematics::CallbackRmRealtimeArmJointState(rm_realtime_arm_joint_state_t data) 
{
    ArmStatusCallbackData callbackData;
    callbackData.errCode = 0;
    if (armStatusCallback_ == nullptr) {
        return;
    }
    if (strcmp(data.arm_ip, LEFT_ARM_IP_ADDRESS) == 0) {
        callbackData.armId = LEFT_ARM;
    }
    else if (strcmp(data.arm_ip, RIGHT_ARM_IP_ADDRESS) == 0) {
        callbackData.armId = RIGHT_ARM;
    } else {
        api_log("Unknown arm ip: %s", data.arm_ip);
        return;
    }

    callbackData.errCode = data.errCode;
    for (size_t i = 0; i < USED_ARM_DOF; i++) {
        callbackData.jointCurrent[i] = data.joint_status.joint_current[i];
        callbackData.jointEnFlag[i] = data.joint_status.joint_en_flag[i];
        callbackData.jointAngle[i] = data.joint_status.joint_position[i];
        callbackData.jointSpeed[i] = data.joint_status.joint_speed[i];
        callbackData.jointTemp[i] = data.joint_status.joint_temperature[i];
        callbackData.jointVoltage[i] = data.joint_status.joint_voltage[i];
        callbackData.jointTorque[i] = data.force_sensor.force[i];
        callbackData.jointErrCode[i] = data.joint_status.joint_err_code[i];
    }

    armStatusCallback_(&callbackData);
}

void RmArmDriverAndKinematics::api_log(const char* format, ...)
{
    va_list args;
    va_start(args, format);// 初始化可变参数列表
    customApiLog_(format, args);// 调用函数指针，传递 va_list
    va_end(args);// 清理可变参数列表
}

int RmArmDriverAndKinematics::InitHandle(rm_robot_handle **robotHandle_, const char *robot_ip_address, int robot_port, rm_realtime_push_config_t *config)
{
    *robotHandle_ = rmService_->rm_create_robot_arm(robot_ip_address, robot_port);
    if (*robotHandle_ == NULL) {
        api_log("Failed to create arm(%s) handle.", robot_ip_address);
        return UNKNOWN_ERR;
    }
    if ((*robotHandle_)->id <= 0) {
        api_log("Failed to create arm(%s) handle. %d", robot_ip_address, (*robotHandle_)->id);
        return UNKNOWN_ERR;
    }

    int result = rmService_-> rm_set_realtime_push(*robotHandle_, *config);
    if (result != 0) {
        api_log("Failed to set realtime push for arm(%s).", robot_ip_address);
        return UNKNOWN_ERR;
    }
    result = rm_set_arm_max_line_speed(*robotHandle_, 0.01f);
    if (result != 0) {
        api_log("Failed to set max line speed for arm(%s).", robot_ip_address);
        return UNKNOWN_ERR;
    }

    rm_robot_info_t robot_info;
    int info_result = rmService_->rm_get_robot_info(*robotHandle_, &robot_info);
    if (info_result != 0) {
        api_log("Failed to get robot info(%s)", robot_ip_address);
        return UNKNOWN_ERR;
    }
    int dof = robot_info.arm_dof;
    if (dof != USED_ARM_DOF) {
        api_log("Invalid degree of freedom(%s)", robot_ip_address);
        return UNKNOWN_ERR;
    }
    result = rm_set_joint_max_pos(*robotHandle_, 2, 112.0f);
    if (result != 0) {
        api_log("Failed to set joint max pos for arm(%s).", robot_ip_address);
        return UNKNOWN_ERR;
    }
    result = rm_set_joint_min_pos(*robotHandle_, 2, -112.0f);
    if (result != 0) {
        api_log("Failed to set joint max pos for arm(%s).", robot_ip_address);
        return UNKNOWN_ERR;
    }
    return 0;
}

int RmArmDriverAndKinematics::Init()
{
    if (initFlag_ == RM_DRIVER_KINEMATICS_INIT_OK) {
        return OK;
    }
    if (customApiLog_ == nullptr || armStatusCallback_ == nullptr) {
        return UNKNOWN_ERR;
    }
    rmService_ = new RM_Service();
    if (rmService_ == nullptr) {
        return UNKNOWN_ERR;
    }
    int result = rmService_->rm_init(RM_TRIPLE_MODE_E);
    if (result != 0) {
        rmService_->rm_destroy();
        delete rmService_;
        return UNKNOWN_ERR;
    }

    rmService_->rm_set_log_call_back(customApiLog_, 3);
    rmService_->rm_realtime_arm_state_call_back(CallbackRmRealtimeArmJointState);

    int robot_port = ARM_IP_PORT;
    rm_realtime_push_config_t config = {20, true, 8089, 0, "192.168.2.3"};
    result = InitHandle(&leftRobotHandle_, LEFT_ARM_IP_ADDRESS, robot_port, &config);
    if (result != 0) {
        rmService_->rm_destroy();
        delete rmService_;
        return UNKNOWN_ERR;
    }
    config = {20, true, 8089, 0, "192.168.2.19"};
    result = InitHandle(&rightRobotHandle_, RIGHT_ARM_IP_ADDRESS, robot_port, &config);
    if (result != 0) {
        rmService_->rm_destroy();
        delete rmService_;
        return UNKNOWN_ERR;
    }

    rm_robot_arm_model_e Mode = RM_MODEL_RM_65_E;
    rm_force_type_e Type = RM_MODEL_RM_B_E;
    rm_algo_init_sys_data(Mode, Type);
    
    initFlag_ = RM_DRIVER_KINEMATICS_INIT_OK;
    return OK;
}

void RmArmDriverAndKinematics::getAngleForCheckingInner(NodataTrajectory *traj, int timeStep, float *aimAngles)
{
    return;
}

void RmArmDriverAndKinematics::getAngleForCheckingOffset(NodataTrajectory *traj, int timeStepOffset, float *aimAngles)
{
    return;
}

int RmArmDriverAndKinematics::ArmMove(const float points[USED_ARM_DOF], int armId)
{
    return OK;
}

int RmArmDriverAndKinematics::updateTrajectory(NodataTrajectory *trajectory, float *aimAngles)
{
    if (trajectory->goalType == GOAL_TYPE_MOVING){
        ++trajectory->currentStep;
        if (trajectory->currentStep >= MIN_MOVING_TIME_STEP) {
            return TRAJECTORY_COMPLETE;
        }
    } else {
        rm_robot_handle *robotHandle = nullptr;
        if (trajectory->armId == LEFT_ARM) {
            robotHandle = leftRobotHandle_;
        } else if (trajectory->armId == RIGHT_ARM) {
            robotHandle = rightRobotHandle_;
        } else {
            api_log("Invalid armId: %d", trajectory->armId);
            return UNKNOWN_ERR;
        }
        int result = 0;
        if (trajectory->goalType == GOAL_TYPE_POSE_STEPPING) {
            if (trajectory->end[POSE_POSITION_X] != 0.0f) {
                result |= rm_set_pos_step(robotHandle, RM_X_DIR_E, trajectory->end[POSE_POSITION_X], 40, 0);
            }
            if (trajectory->end[POSE_POSITION_Y] != 0.0f) {
                result |= rm_set_pos_step(robotHandle, RM_Y_DIR_E, trajectory->end[POSE_POSITION_Y], 40, 0);
            }
            if (trajectory->end[POSE_POSITION_Z] != 0.0f) {
                result |= rm_set_pos_step(robotHandle, RM_Z_DIR_E, trajectory->end[POSE_POSITION_Z], 40, 0);
            }
        } else if (trajectory->goalType == GOAL_TYPE_STEPPING) {
            for (int i = 0; i < USED_ARM_DOF; ++i) {
                result |= rm_set_joint_step(robotHandle, i + 1, trajectory->end[i], 5, 0);
            }
        }

        if (result != 0) {
            return ARM_AIM_CANNOT_REACH;
        }
    }
    return TRAJECTORY_ONGOING;
}

int RmArmDriverAndKinematics::computeMovingTrajectory(NodataTrajectory *newTrajectory, const float *start, const float *end, int armId)
{
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->start[i] = start[i];
        newTrajectory->end[i] = end[i];
    }
    newTrajectory->runState = RM_KINEMATICS_STATUS_INIT;
    newTrajectory->goalType = GOAL_TYPE_MOVING;
    newTrajectory->armId = armId;
    newTrajectory->currentStep = 0;

    rm_robot_handle *robotHandle = nullptr;
    if (armId == LEFT_ARM) {
        robotHandle = leftRobotHandle_;
    } else if (armId == RIGHT_ARM) {
        robotHandle = rightRobotHandle_;
    } else {
        api_log("Invalid armId: %d", armId);
        return UNKNOWN_ERR;
    }
    float joints[ARM_DOF];
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        joints[i] = end[i];
    }
    return rm_movej(robotHandle, joints, 20, 1, 0, 0);
}

int RmArmDriverAndKinematics::computeSteppingTrajectory(NodataTrajectory *newTrajectory, const float *start, const float *speed, int armId)
{
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->start[i] = start[i];
        newTrajectory->end[i] = speed[i];
    }
    newTrajectory->runState = RM_KINEMATICS_STATUS_INIT;
    newTrajectory->goalType = GOAL_TYPE_STEPPING;
    newTrajectory->armId = armId;
    newTrajectory->currentStep = 0;
    rm_robot_handle *robotHandle = nullptr;
    if (armId == LEFT_ARM) {
        robotHandle = leftRobotHandle_;
    } else if (armId == RIGHT_ARM) {
        robotHandle = rightRobotHandle_;
    } else {
        api_log("Invalid armId: %d", armId);
        return UNKNOWN_ERR;
    }
    float joints[ARM_DOF];
    int result = 0;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        result |= rm_set_joint_step(robotHandle, i + 1, speed[i], 5, 0);
    }
    return result;
}

int RmArmDriverAndKinematics::computePoseSteppingTrajectory(NodataTrajectory *newTrajectory, const float *start, const float *pose, int armId)
{
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->start[i] = start[i];
        newTrajectory->end[i] = pose[i];
    }
    newTrajectory->runState = RM_KINEMATICS_STATUS_INIT;
    newTrajectory->goalType = GOAL_TYPE_POSE_STEPPING;
    newTrajectory->armId = armId;
    newTrajectory->currentStep = 0;
    rm_robot_handle *robotHandle = nullptr;
    if (armId == LEFT_ARM) {
        robotHandle = leftRobotHandle_;
    } else if (armId == RIGHT_ARM) {
        robotHandle = rightRobotHandle_;
    } else {
        api_log("Invalid armId: %d", armId);
        return UNKNOWN_ERR;
    }

    int result = 0;
    if (pose[POSE_POSITION_X] != 0.0f) {
        result |= rm_set_pos_step(robotHandle, RM_X_DIR_E, pose[POSE_POSITION_X], 40, 0);
    }
    if (pose[POSE_POSITION_Y] != 0.0f) {
        result |= rm_set_pos_step(robotHandle, RM_Y_DIR_E, pose[POSE_POSITION_Y], 40, 0);
    }
    if (pose[POSE_POSITION_Z] != 0.0f) {
        result |= rm_set_pos_step(robotHandle, RM_Z_DIR_E, pose[POSE_POSITION_Z], 40, 0);
    }

    return result;
}

void RmArmDriverAndKinematics::stopTrajectory(NodataTrajectory *traj)
{
    rm_robot_handle *robotHandle = nullptr;
    if (traj->armId == LEFT_ARM) {
        robotHandle = leftRobotHandle_;
    } else if (traj->armId == RIGHT_ARM) {
        robotHandle = rightRobotHandle_;
    } else {
        api_log("Invalid armId: %d", traj->armId);
        return;
    }
    if (traj->goalType == GOAL_TYPE_MOVING) {
        rm_set_arm_stop(robotHandle);
        return;
    }
    else if (traj->goalType == GOAL_TYPE_STEPPING) {
        traj->goalType = GOAL_TYPE_MOVING;
        traj->currentStep = MIN_MOVING_TIME_STEP;
        return;
    }
    else if (traj->goalType == GOAL_TYPE_POSE_STEPPING) {
        traj->goalType = GOAL_TYPE_MOVING;
        traj->currentStep = MIN_MOVING_TIME_STEP;
        return;
    }
}

int RmArmDriverAndKinematics::continuePoseSteppingTrajectory(NodataTrajectory *newTrajectory, const float *nowAngles, const float *direct)
{
    return ARM_GOAL_CANCELLED;
}

int RmArmDriverAndKinematics::ForwardKinematics(float *jointAngles, float *endEffectorPose)
{
    if (!rmService_)
    {
        return UNKNOWN_ERR; // Error: RM_Service not initialized
    }

    // Call the RM_Service's ForwardKinematics method
    rm_pose_t pose = rmService_->rm_algo_forward_kinematics(NULL, jointAngles);
    endEffectorPose[POSE_POSITION_X] = pose.position.x;
    endEffectorPose[POSE_POSITION_Y] = pose.position.y;
    endEffectorPose[POSE_POSITION_Z] = pose.position.z;
    endEffectorPose[POSE_QUATERNION_X] = pose.quaternion.x;
    endEffectorPose[POSE_QUATERNION_Y] = pose.quaternion.y;
    endEffectorPose[POSE_QUATERNION_Z] = pose.quaternion.z;
    endEffectorPose[POSE_QUATERNION_W] = pose.quaternion.w;
    return 0;
}

int RmArmDriverAndKinematics::InverseKinematics(const float *endEffectorPose, float *jointAngles)
{
    if (!rmService_)
    {
        return UNKNOWN_ERR; // Error: RM_Service not initialized
    }

    float nowJoints_[ARM_DOF] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    for (int i = 0; i < USED_ARM_DOF; i++) {
        nowJoints_[i] = jointAngles[i];
    }
    rm_inverse_kinematics_params_t inverse_params;
    memcpy(inverse_params.q_in, nowJoints_, sizeof(nowJoints_));
    inverse_params.q_pose.position.x = endEffectorPose[POSE_POSITION_X];
    inverse_params.q_pose.position.y = endEffectorPose[POSE_POSITION_Y];
    inverse_params.q_pose.position.z = endEffectorPose[POSE_POSITION_Z];
    inverse_params.q_pose.quaternion.x = endEffectorPose[POSE_QUATERNION_X];
    inverse_params.q_pose.quaternion.y = endEffectorPose[POSE_QUATERNION_Y];
    inverse_params.q_pose.quaternion.z = endEffectorPose[POSE_QUATERNION_Z];
    inverse_params.q_pose.quaternion.w = endEffectorPose[POSE_QUATERNION_W];
    inverse_params.flag = 0;

    int result = rmService_->rm_algo_inverse_kinematics(NULL, inverse_params, nowJoints_);
    if (result != 0){
        return ARM_AIM_CANNOT_REACH; // Inverse kinematics failed
    }
    memcpy(jointAngles, nowJoints_, sizeof(float) * USED_ARM_DOF);
    return 0;
}


