#pragma once

#include "dds_interface.hpp"

/** Identity-calibrated DDS interface for the 8-DoF FireAnt MuJoCo model. */
class FireAntInterface : public DdsInterface {
public:
    explicit FireAntInterface(const std::string &robot_name)
        : DdsInterface(robot_name, 8) {
        for (int i = 0; i < dof_num_; ++i) {
            pos_offset_[i] = 0.0f;
            joint_dir_[i] = 1.0f;
            data_updated_[i] = false;
            joint_config_[i].dir = 1.0f;
            joint_config_[i].offset = 0.0f;
        }
    }

    void Start() override {
        time_stamp_ = GetTimestampMs();
    }

    void Stop() override {}

    void SetJointCommand(Eigen::Matrix<float, Eigen::Dynamic, 5> input) override {
        if (input.rows() != dof_num_ || input.cols() != 5) {
            throw std::runtime_error("FireAnt joint command must be 8x5");
        }
        joint_cmd_ = input;
        auto msg = drdds::msg::JointsDataCmd();
        for (int i = 0; i < dof_num_; ++i) {
            msg.data.joints_data[i].position = input(i, 1);
            msg.data.joints_data[i].velocity = input(i, 3);
            msg.data.joints_data[i].torque = input(i, 4);
            msg.data.joints_data[i].kp = (i % 2 == 0) ? input(i, 0) : 0.0f;
            msg.data.joints_data[i].kd = input(i, 2);
            msg.data.joints_data[i].control_word = kIndexMotorControl;
        }
        joint_cmd_pub_->publish(msg);
    }
};