#include <chrono>
#include <memory>
#include <cmath>

#include "rclcpp/rclcpp.hpp"
#include "advrobotics_lab3_interfaces/msg/joints.hpp"
#include "advrobotics_lab3_interfaces/msg/gripper.hpp"
#include "advrobotics_lab3_interfaces/srv/invkin.hpp"

using namespace std::chrono_literals;

class BuildTower : public rclcpp::Node
{
public:
    BuildTower() : Node("build_tower")
    {
        joint_cmd_pub_ = create_publisher<advrobotics_lab3_interfaces::msg::Joints>("joint_cmd", 10);
        gripper_cmd_pub_ = create_publisher<advrobotics_lab3_interfaces::msg::Gripper>("gripper_cmd", 10);

        joint_pos_sub_ = create_subscription<advrobotics_lab3_interfaces::msg::Joints>(
            "joint_position", 10,
            std::bind(&BuildTower::jointPosCallback, this, std::placeholders::_1));

        ik_client_ = create_client<advrobotics_lab3_interfaces::srv::Invkin>("invkin");

        while (!ik_client_->wait_for_service(1s))
            RCLCPP_INFO(get_logger(), "Waiting for IK service...");

        executeTower();
    }

private:
    /* ---------------- CALLBACKS ---------------- */

    void jointPosCallback(const advrobotics_lab3_interfaces::msg::Joints::SharedPtr msg)
    {
        current_q1_ = msg->q1;
        current_q2_ = msg->q2;
        current_q3_ = msg->q3;
    }

    /* ---------------- MOTION PRIMITIVES ---------------- */

    void moveToCartesian(double x, double y, double z)
    {
        auto request = std::make_shared<advrobotics_lab3_interfaces::srv::Invkin::Request>();
        request->x = x;
        request->y = y;
        request->z = z;

        auto future = ik_client_->async_send_request(request);
        rclcpp::spin_until_future_complete(shared_from_this(), future);

        auto response = future.get();

        advrobotics_lab3_interfaces::msg::Joints cmd;
        cmd.q1 = response->q1 * 180.0 / M_PI;
        cmd.q2 = response->q2 * 180.0 / M_PI;
        cmd.q3 = response->q3 * 180.0 / M_PI;

        joint_cmd_pub_->publish(cmd);
        waitUntilReached(cmd);
    }

    void waitUntilReached(const advrobotics_lab3_interfaces::msg::Joints &target)
    {
        while (rclcpp::ok())
        {
            if (std::fabs(current_q1_ - target.q1) < 2.0 &&
                std::fabs(current_q2_ - target.q2) < 2.0 &&
                std::fabs(current_q3_ - target.q3) < 2.0)
                break;

            rclcpp::sleep_for(100ms);
        }
    }

    void setGripper(double angle)
    {
        advrobotics_lab3_interfaces::msg::Gripper g;
        g.gripper = angle;
        gripper_cmd_pub_->publish(g);
        rclcpp::sleep_for(1s);
    }

    /* ---------------- MAIN LOGIC ---------------- */

    void executeTower()
    {
        constexpr double pick_x = 0.10;
        constexpr double pick_y = 0.10;
        constexpr double pick_z = 0.02;

        constexpr double place_x = 0.15;
        constexpr double place_y = 0.0;
        constexpr double base_z  = 0.02;
        constexpr double block_height = 0.03;

        constexpr double approach_z = 0.10;
        constexpr double gripper_open = 30.0;
        constexpr double gripper_closed = -10.0;

        setGripper(gripper_open);

        for (int i = 0; i < 3; ++i)
        {
            double place_z = base_z + i * block_height;

            moveToCartesian(pick_x, pick_y, approach_z);
            moveToCartesian(pick_x, pick_y, pick_z);
            setGripper(gripper_closed);

            moveToCartesian(pick_x, pick_y, approach_z);
            moveToCartesian(place_x, place_y, approach_z);
            moveToCartesian(place_x, place_y, place_z);

            setGripper(gripper_open);
            moveToCartesian(place_x, place_y, approach_z);
        }

        RCLCPP_INFO(get_logger(), "Tower completed.");
    }

private:
    rclcpp::Publisher<advrobotics_lab3_interfaces::msg::Joints>::SharedPtr joint_cmd_pub_;
    rclcpp::Publisher<advrobotics_lab3_interfaces::msg::Gripper>::SharedPtr gripper_cmd_pub_;
    rclcpp::Subscription<advrobotics_lab3_interfaces::msg::Joints>::SharedPtr joint_pos_sub_;
    rclcpp::Client<advrobotics_lab3_interfaces::srv::Invkin>::SharedPtr ik_client_;

    double current_q1_{0.0}, current_q2_{0.0}, current_q3_{0.0};
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<BuildTower>());
    rclcpp::shutdown();
    return 0;
}
