#include <iostream>
#include <chrono>
#include <memory>

#include <boost/asio.hpp>

#include "geometry_msgs/msg/transform.hpp"
#include "geometry_msgs/msg/transform_stamped.hpp"
#include "rclcpp/rclcpp.hpp"

#include "toid_odometry/msgs.h"
#include "toid_odometry/crc.h"

#include "toid_msgs/srv/send_double.hpp"
#include "std_srvs/srv/empty.hpp"

#include "tf2_ros/buffer.hpp"
#include "tf2_ros/transform_broadcaster.hpp"
#include "tf2_ros/transform_listener.hpp"
#include "tf2_ros/static_transform_broadcaster.hpp"
#include "tf2/LinearMath/Transform.hpp"

#include "tf2_geometry_msgs/tf2_geometry_msgs.hpp"
#include "tf2/convert.hpp"

#include "nav_msgs/msg/odometry.hpp"


namespace asio = boost::asio;

const std::array<double, 36> COVARIENCE_MATRIX{
    0.0001, .0, 0, 0, 0, 0,
    0, 0.0001, 0, 0, 0, 0,
    0, .0, 0.0001, 0, 0, 0,
    0, .0, 0, 0.0001, 0, 0,
    0, .0, 0, 0, 0.0001, 0,
    0, .0, 0, 0, 0, 0.0001,
}; 

class ToidOdomNode : public rclcpp::Node {
    using SendDoubleSrv = toid_msgs::srv::SendDouble;
    using ZeroSrv       = std_srvs::srv::Empty;
public:
    ToidOdomNode() : rclcpp::Node("OdometryNode"), serial_(io_ctx_) {
        using std::placeholders::_1;
        this->declare_parameter("base_frame", "base_footprint");
        this->declare_parameter("odometry_frame", "odom");
        this->declare_parameter("map_frame", "map");
        this->declare_parameter("serial_port", "/dev/ttyACM1");

        this->declare_parameter("initial_pose", "0 0 0");
        this->declare_parameter("mock_odom", false);

        odometry_frame_id_ = this->get_parameter("odometry_frame").as_string();
        base_frame_id_ = this->get_parameter("base_frame").as_string();
        map_frame_id_ = this->get_parameter("map_frame").as_string();
        serial_port_path_ = this->get_parameter("serial_port").as_string();

        tf_broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>(this);
        tf_static_broadcaster_ = std::make_shared<tf2_ros::StaticTransformBroadcaster>(this);

        tf_buffer_ = std::make_unique<tf2_ros::Buffer>(this->get_clock());
        tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

        geometry_msgs::msg::TransformStamped t;
        t.header.stamp = this->get_clock()->now();
        t.header.frame_id = map_frame_id_;
        t.child_frame_id = odometry_frame_id_;
        t.transform = this->make_transform(0,0,0);
        tf_static_broadcaster_->sendTransform(t);

        startpose_listener_ = this->create_subscription<geometry_msgs::msg::PoseWithCovarianceStamped>(
            "/initialpose", 1, std::bind(&ToidOdomNode::set_pose,this,_1)
        );

        if(!this->get_parameter("mock_odom").as_bool()) {
            timer_ = this->create_wall_timer(
                std::chrono::seconds(1/100), 
                [this]{this->publish_robot_state();});
            
            set_width_service_ = create_service<SendDoubleSrv> (
                "/set_width", std::bind(&ToidOdomNode::set_width, this, _1));

            set_ratio_service_ = create_service<SendDoubleSrv> (
                "/set_ratio", std::bind(&ToidOdomNode::set_ratio, this, _1));

            calib_service_ = create_service<ZeroSrv> (
                "/end_calib", std::bind(&ToidOdomNode::end_calib, this, _1));

            zero_service_ = create_service<ZeroSrv> (
                "/zero", std::bind(&ToidOdomNode::zero, this, _1));

            odometry_publisher_ = this->create_publisher<nav_msgs::msg::Odometry>(
                "/odom", 1);
        }

    }
private:

    void set_pose(const geometry_msgs::msg::PoseWithCovarianceStamped &pose) {

        geometry_msgs::msg::TransformStamped t;

        try {
            geometry_msgs::msg::TransformStamped base = tf_buffer_->lookupTransform(
                odometry_frame_id_, 
                base_frame_id_, 
                tf2::TimePointZero
            );

            t.transform.rotation = pose.pose.pose.orientation;
            auto[x,y,z] =pose.pose.pose.position;
            t.transform.translation.x = x;
            t.transform.translation.y = y;
            t.transform.translation.z = z;

            tf2::Transform map_target;
            tf2::fromMsg(t.transform, map_target);

            tf2::Transform odom_robot;
            tf2::fromMsg(base.transform, odom_robot);

            tf2::Transform map_odom;
            map_odom = map_target * odom_robot.inverse();
            t.transform = tf2::toMsg(map_odom);
        }
        catch (const tf2::TransformException &e) {
            RCLCPP_ERROR(this->get_logger(), "%s", e.what());
            return;
        }

        t.header.stamp = this->get_clock()->now();
        t.header.frame_id = map_frame_id_;
        t.child_frame_id = odometry_frame_id_;


        RCLCPP_INFO(this->get_logger(), "Setting new initial_pose");

        tf_static_broadcaster_->sendTransform(t);
    }

    void set_width(const std::shared_ptr<SendDoubleSrv::Request> req) {
        boost::system::error_code ec;
        const std::array<uint8_t, 2> cmd{TMSG_SET_WIDTH, TMSG_DELIM};
        tMsgWidth msg = { req->data, 0};
        msg.crc = crcFooter(msg);
        asio::write(this->serial_, asio::buffer(cmd), ec);
        asio::write(this->serial_, asio::buffer(&msg, sizeof(msg)), ec);
    }

    void set_ratio(const std::shared_ptr<SendDoubleSrv::Request> req) {
        boost::system::error_code ec;
        const std::array<uint8_t, 2> cmd{TMSG_SET_RATIO, TMSG_DELIM};
        tMsgRatio msg = { req->data, 0 };
        msg.crc = crcFooter(msg);
        asio::write(this->serial_, asio::buffer(cmd), ec);
        asio::write(this->serial_, asio::buffer(&msg, sizeof(msg)), ec);
    }

    void end_calib(const std::shared_ptr<ZeroSrv::Request>) {
        boost::system::error_code ec;
        const std::array<uint8_t, 2> cmd{TMSG_END_CALIB, TMSG_DELIM};
        asio::write(this->serial_, asio::buffer(cmd), ec);
        tRespEndCalib resp;
        asio::read(this->serial_, asio::buffer(&resp, sizeof(resp)) , ec);
        if(!ec.failed()) {
            RCLCPP_INFO(this->get_logger(), "Calculated ratio: %lf, %lf", resp.left_gain, resp.right_gain);
        }
    }

    void zero(const std::shared_ptr<ZeroSrv::Request>) {
        boost::system::error_code ec;
        const std::array<uint8_t, 2> cmd{TMSG_ZERO, TMSG_DELIM};
        asio::write(this->serial_, asio::buffer(cmd), ec);
    }

    void publish_robot_state() {
        tRespState state;
        if(!read_robot_state(state)) {
            return;
        }
        send_transform(state);
    }

    geometry_msgs::msg::Transform make_transform(double x, double y, double theta){
        geometry_msgs::msg::Transform t;
        tf2::Quaternion q;
        q.setRPY(0, 0, theta);
        const tf2::Transform tf(q, tf2::Vector3(x, y, 0));
        tf2::convert(tf, t);
        return t;
    }

    void send_transform(tRespState &state) {
        geometry_msgs::msg::TransformStamped t;
        nav_msgs::msg::Odometry o;

        t.header.stamp = this->get_clock()->now();
        t.header.frame_id = this->odometry_frame_id_;
        t.child_frame_id = this->base_frame_id_;

        o.header = t.header;
        o.child_frame_id = t.child_frame_id;

        t.transform = make_transform(state.x, state.y, state.theta);

        const auto& [x,y,z] = t.transform.translation;
        o.pose.pose.position.x = x;
        o.pose.pose.position.x = y;
        o.pose.pose.position.x = z;
        o.pose.pose.orientation = t.transform.rotation;
        o.pose.covariance = COVARIENCE_MATRIX;

        o.twist.twist.linear.x = state.vl;
        o.twist.twist.angular.z = state.vr;
        o.twist.covariance = COVARIENCE_MATRIX;

        odometry_publisher_->publish(o);

        tf_broadcaster_->sendTransform(t);
    }

    bool read_robot_state(tRespState &state) {
        const std::array<uint8_t, 2> cmd{TMSG_REQ_STATE, TMSG_DELIM};
        try {
            asio::write(this->serial_, asio::buffer(cmd));
            asio::read(this->serial_, asio::buffer(&state, sizeof(state)));
        } catch (const std::exception &e){
            RCLCPP_ERROR_SKIPFIRST_THROTTLE(
                this->get_logger(),
                *this->get_clock(), 
                5000, 
                "%s", 
                e.what()
            );
            boost::system::error_code ec;
            if(serial_.close(ec)) {}
            if(serial_.open(serial_port_path_,ec)) {
                RCLCPP_ERROR_THROTTLE(
                    this->get_logger(),
                    *this->get_clock(), 
                    5000, 
                    "%s", 
                    ec.what().c_str()
                );
            }
            return false;
        }
        return state.crc == crcFooter(state);
    }
private: 
    asio::io_context io_ctx_;
    asio::serial_port serial_;
    rclcpp::TimerBase::SharedPtr timer_;

    std::string odometry_frame_id_;
    std::string base_frame_id_;
    std::string map_frame_id_;
    std::string serial_port_path_;

    std::shared_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
    std::shared_ptr<tf2_ros::StaticTransformBroadcaster> tf_static_broadcaster_;
    std::unique_ptr<tf2_ros::Buffer> tf_buffer_;
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
    rclcpp::Service<SendDoubleSrv>::SharedPtr set_width_service_;
    rclcpp::Service<SendDoubleSrv>::SharedPtr set_ratio_service_;
    rclcpp::Service<ZeroSrv>::SharedPtr calib_service_;
    rclcpp::Service<ZeroSrv>::SharedPtr zero_service_;

    rclcpp::Subscription<geometry_msgs::msg::PoseWithCovarianceStamped>::SharedPtr startpose_listener_;
    rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr odometry_publisher_;
};

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