#include <chrono>
#include <memory>
#include <vector>

#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")
    {
        // Publishers
        joint_pub_ = this->create_publisher<advrobotics_lab3_interfaces::msg::Joints>("joint_cmd", 10);
        gripper_pub_ = this->create_publisher<advrobotics_lab3_interfaces::msg::Gripper>("gripper_cmd", 10);

        // IK client
        ik_client_ = this->create_client<advrobotics_lab3_interfaces::srv::Invkin>("invkin");

        RCLCPP_INFO(this->get_logger(), "Waiting for IK service...");
        while (!ik_client_->wait_for_service(1s)) {
            RCLCPP_WARN(this->get_logger(), "IK service not available, waiting...");
        }

        // Start the tower building sequence
        build_tower();
    }

private:
    void build_tower()
    {
        // Define cube positions and tower location (x,y,z)
        std::vector<std::vector<double>> cubes = {
            {11.0, 0.0, 0.0},
            {11.0, 8.0, 0.0},
            {11.0, -8.0, 0.0}
        };
        std::vector<double> tower = {9.0, 0.0, 0.0};
        double safe_height = 7;  // height above blocks to avoid collisions
        double cube_height = 2.4;  // height above blocks to avoid collisions

        // Open gripper at start
        advrobotics_lab3_interfaces::msg::Gripper gripper_msg;
        gripper_msg.gripper = -40.0; // open
        gripper_pub_->publish(gripper_msg);
        rclcpp::sleep_for(500ms);

        // Loop through all cubes
        for (size_t i = 0; i < cubes.size(); i++)
        {
            RCLCPP_INFO(this->get_logger(), "Picking cube %zu", i+1);

            // 1. Move above cube (safe height)
            move_to(cubes[i][0], cubes[i][1], safe_height);

            // 2. Move down to cube
            move_to(cubes[i][0], cubes[i][1], cubes[i][2]);

            // 3. Close gripper
            gripper_msg.gripper = 20; // closed
            gripper_pub_->publish(gripper_msg);
            rclcpp::sleep_for(500ms);

            // 4. Lift cube to safe height
            move_to(cubes[i][0], cubes[i][1], safe_height);
            rclcpp::sleep_for(500ms);

            // 5. Move above tower (safe height)
            move_to(tower[0], tower[1], safe_height); // stack cubes
            rclcpp::sleep_for(500ms);

            // 6. Lower to tower target
            if(i==0){
           	 move_to(tower[0], tower[1], cube_height*i);
            } 
            else{
           	 move_to(tower[0], tower[1], cube_height*(i-1));
            };

            // 7. Open gripper to release
            gripper_msg.gripper = -40.0; // open
            gripper_pub_->publish(gripper_msg);
            rclcpp::sleep_for(500ms);

            // 8. Lift gripper after releasing
            move_to(tower[0], tower[1], safe_height + (i+1)*2.0);
        }

        RCLCPP_INFO(this->get_logger(), "Tower building completed!");
    }

    void move_to(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);
        if (rclcpp::spin_until_future_complete(this->get_node_base_interface(), future) == rclcpp::FutureReturnCode::SUCCESS)
        {
            auto response = future.get();
            advrobotics_lab3_interfaces::msg::Joints joint_msg;
            joint_msg.q1 = response->q1 * 180.0 / M_PI; // radians to degrees
            joint_msg.q2 = response->q2 * 180.0 / M_PI;
            joint_msg.q3 = response->q3 * 180.0 / M_PI;
            joint_pub_->publish(joint_msg);

            // Wait for the robot to move (rough approximation)
            rclcpp::sleep_for(1s);
        }
        else
        {
            RCLCPP_ERROR(this->get_logger(), "Failed to call IK service for (%f,%f,%f)", x, y, z);
        }
    }

    rclcpp::Publisher<advrobotics_lab3_interfaces::msg::Joints>::SharedPtr joint_pub_;
    rclcpp::Publisher<advrobotics_lab3_interfaces::msg::Gripper>::SharedPtr gripper_pub_;
    rclcpp::Client<advrobotics_lab3_interfaces::srv::Invkin>::SharedPtr ik_client_;
};

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