//
// Created by lgv on 2025/8/15.
//
#ifdef MAX_ITER
#undef MAX_ITER
#endif

#include <opencv2/opencv.hpp>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <mutex>

#include "../src/devices/robot/abstract_robot.h"
#include "../src/device_manager/include/device_manager.h"
#include "cmvr/msgs/geometry.pb.h"

namespace py = pybind11;
using namespace cmvr::device;

class PyRobotWrapper {
public:
    // 构造时只需要传入 config_path 和 robot_name（只在第一次初始化有效）
    PyRobotWrapper(const std::string& config_path, const std::string& robot_name) {
        std::lock_guard<std::mutex> lock(init_mutex); // 线程安全
        if (!robot_) {
            const XmlNode config(config_path);
            if (!config.hasChild("DeviceManager")) {
                throw std::runtime_error("Device Manager node not found");
            }
            auto dmgr_cfg = config.getChild("DeviceManager");
            dmgr_ = &DeviceManager::getInstance(dmgr_cfg);
            robot_ = dmgr_->getDevice<AbstractRobot>(robot_name);
        }
    }

    void moveJ(const std::string &side, const std::vector<double> &q, double speed=0.8) {
        if (q.size() != 7)
            throw std::runtime_error("Expected 7 joint values");

        const auto& joint_names = getJointNames(side);

        std::vector<JointPoint> cmd = {{"WAIST_Y", 0}, {"WAIST_P", 0}};
        for (size_t i = 0; i < 7; ++i)
            cmd.push_back({joint_names[i], q[i]});

        robot_->moveJ(cmd, speed);
    }

    void torqueOn() {
        checkRobotInit();
        robot_->torqueOn();
    }

    void torqueOn(const std::string& joint_name) {
        checkRobotInit();
        robot_->torqueOn(joint_name);
    }

    void torqueOff() {
        checkRobotInit();
        robot_->torqueOff();
    }

    void torqueOff(const std::string& joint_name) {
        checkRobotInit();
        robot_->torqueOff(joint_name);
    }

    void calibrateZeroQ(const std::string &joint_name) {
        checkRobotInit();
        robot_->calibrateZeroQ(joint_name);
    }

    std::vector<double> getJointQ(const std::string& side) const {
        checkRobotInit();

        auto joint_qs = robot_->getJointQ();  // unordered_map<std::string,double>
        const auto& joint_names = getJointNames(side);

        std::vector<double> values;
        for (const auto& name : joint_names) {
            values.push_back(joint_qs.at(name));
        }

        return values;
    }

    void move(const std::string &base_link, const std::string &ee_link,
           double x, double y, double z, double rx, double ry, double rz,
           double vel, double acc) {
        checkRobotInit();

        cmvr::msgs::Pose3d pose;
        pose.mutable_position()->set_x(x);
        pose.mutable_position()->set_y(y);
        pose.mutable_position()->set_z(z);

        pose.mutable_euler()->set_rx(rx);
        pose.mutable_euler()->set_ry(ry);
        pose.mutable_euler()->set_rz(rz);

        robot_->moveJ(base_link, ee_link, pose,vel,acc);
    }

    std::vector<double> ik(const std::string& side,
                           const std::string &base_link,
                           const std::string &ee_link,
                           double x, double y, double z,
                           double rx, double ry, double rz)
    {
        checkRobotInit();

        cmvr::msgs::Pose3d pose;
        pose.mutable_position()->set_x(x);
        pose.mutable_position()->set_y(y);
        pose.mutable_position()->set_z(z);

        pose.mutable_euler()->set_rx(rx);
        pose.mutable_euler()->set_ry(ry);
        pose.mutable_euler()->set_rz(rz);

        auto js_cmd = robot_->ik(base_link, ee_link, pose);

        if (js_cmd.size() < 14) {
            throw std::runtime_error("IK solution size is smaller than expected");
        }

        if (side == "left") {
            return std::vector<double>(js_cmd.begin(), js_cmd.begin() + 7);
        } else if (side == "right") {
            return std::vector<double>(js_cmd.begin() + 7, js_cmd.begin() + 14);
        } else {
            throw std::runtime_error("Side must be 'left' or 'right'");
        }
    }


    std::tuple<double, double, double, double, double, double> fk(const std::string &base_link, const std::string &ee_link) {
        checkRobotInit();
        auto pose = robot_->fk(base_link, ee_link);

        // 返回一个元组，包含 (x, y, z, rx, ry, rz)
        return std::make_tuple(pose.position().x(), pose.position().y(),  pose.position().z(),  pose.euler().rx(), pose.euler().ry(), pose.euler().ry());
    }

private:
    static DeviceManager* dmgr_;
    static std::shared_ptr<AbstractRobot> robot_;
    static std::mutex init_mutex;

    static const std::vector<std::string> LEFT_JOINTS;
    static const std::vector<std::string> RIGHT_JOINTS;

    void checkRobotInit() const {
        if (!robot_)
            throw std::runtime_error("Robot not initialized");
    }

    const std::vector<std::string>& getJointNames(const std::string& side) const {
        if (side == "left") return LEFT_JOINTS;
        if (side == "right") return RIGHT_JOINTS;
        throw std::runtime_error("Side must be 'left' or 'right'");
    }
};

// 静态成员初始化
DeviceManager* PyRobotWrapper::dmgr_ = nullptr;
std::shared_ptr<AbstractRobot> PyRobotWrapper::robot_ = nullptr;
std::mutex PyRobotWrapper::init_mutex;

const std::vector<std::string> PyRobotWrapper::LEFT_JOINTS = {"L_SHOULDER_P", "L_SHOULDER_R", "L_SHOULDER_Y",
                                                              "L_ELBOW_R", "L_WRIST_P", "L_WRIST_Y", "L_WRIST_R"};
const std::vector<std::string> PyRobotWrapper::RIGHT_JOINTS = {"R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
                                                               "R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"};

// Python 绑定
PYBIND11_MODULE(robot_wrapper, m) {
    py::class_<PyRobotWrapper>(m, "Robot")
        .def(py::init<const std::string&, const std::string&>(),
             py::arg("config_path"), py::arg("robot_name"),
             "Initialize the robot wrapper with config and robot name")
        .def("moveJ", &PyRobotWrapper::moveJ, py::arg("side"), py::arg("q"), py::arg("speed")=0.8,
             "Move robot joints in joint space with optional speed")
        .def("torqueOn", py::overload_cast<>(&PyRobotWrapper::torqueOn))
        .def("torqueOn", py::overload_cast<const std::string&>(&PyRobotWrapper::torqueOn))
        .def("torqueOff", py::overload_cast<>(&PyRobotWrapper::torqueOff))
        .def("torqueOff", py::overload_cast<const std::string&>(&PyRobotWrapper::torqueOff))
        .def("calibrateZeroQ", &PyRobotWrapper::calibrateZeroQ)
        .def("getJointQ", &PyRobotWrapper::getJointQ, py::arg("side"),
             "Get joint positions for the specified side ('left' or 'right')")
        .def("ik", &PyRobotWrapper::ik,
             py::arg("side"), py::arg("base_link"), py::arg("ee_link"),
             py::arg("x"), py::arg("y"), py::arg("z"),
             py::arg("rx"), py::arg("ry"), py::arg("rz"),
             "Compute inverse kinematics and return 7 joint values for the given side")
        .def("moveJ", &PyRobotWrapper::move, "Move robot in joint space",
            py::arg("base_link"), py::arg("ee_link"),
            py::arg("x"), py::arg("y"), py::arg("z"),
            py::arg("rx"), py::arg("ry"), py::arg("rz"),
            py::arg("vel") = 0.5, py::arg("acc") = 0.1)
        .def("fk", &PyRobotWrapper::fk, "Compute forward kinematics",
             py::arg("base_link"), py::arg("ee_link"));
}