#include <ros/ros.h>
#include <sensor_msgs/LaserScan.h>
#include <nav_msgs/OccupancyGrid.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <geometry_msgs/Twist.h>
#include <unitree_legged_msgs/HighCmd.h>
#include <unitree_legged_msgs/HighState.h>
#include "unitree_legged_sdk/unitree_legged_sdk.h"

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <thread>
#include <mutex>
#include <json/json.h>

using namespace UNITREE_LEGGED_SDK;

class PCSlamwareReceiver
{
private:
    ros::NodeHandle nh_;
    ros::NodeHandle private_nh_;
    
    // ROS Publishers - 发布接收到的数据
    ros::Publisher scan_pub_;
    ros::Publisher map_pub_;
    ros::Publisher odom_pub_;
    ros::Publisher pose_pub_;
    
    // Network
    int socket_fd_;
    std::string go1_ip_;
    int data_port_;
    bool running_;
    std::thread receive_thread_;
    
public:
    PCSlamwareReceiver() : private_nh_("~"), running_(false), socket_fd_(-1)
    {
        // 获取参数
        private_nh_.param<std::string>("go1_ip", go1_ip_, "192.168.123.15");
        private_nh_.param<int>("data_port", data_port_, 8888);
        
        // 创建发布者 - 发布到标准ROS话题供RViz使用
        scan_pub_ = nh_.advertise<sensor_msgs::LaserScan>("scan", 10);
        map_pub_ = nh_.advertise<nav_msgs::OccupancyGrid>("map", 1);
        odom_pub_ = nh_.advertise<nav_msgs::Odometry>("odom", 10);
        pose_pub_ = nh_.advertise<geometry_msgs::PoseWithCovarianceStamped>("robot_pose", 10);
        
        ROS_INFO("PC Slamware Receiver initialized");
        ROS_INFO("Will connect to GO1 at %s:%d", go1_ip_.c_str(), data_port_);
        ROS_INFO("Publishing to topics:");
        ROS_INFO("  - /scan (for RViz laser display)");
        ROS_INFO("  - /map (for RViz map display)");
        ROS_INFO("  - /odom (for odometry)");
        ROS_INFO("  - /robot_pose (for robot pose)");
    }
    
    bool connectToGO1()
    {
        socket_fd_ = socket(AF_INET, SOCK_STREAM, 0);
        if (socket_fd_ < 0) {
            ROS_ERROR("Failed to create socket");
            return false;
        }
        
        struct sockaddr_in server_addr;
        server_addr.sin_family = AF_INET;
        server_addr.sin_port = htons(data_port_);
        inet_pton(AF_INET, go1_ip_.c_str(), &server_addr.sin_addr);
        
        if (connect(socket_fd_, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
            ROS_ERROR("Failed to connect to GO1 at %s:%d", go1_ip_.c_str(), data_port_);
            close(socket_fd_);
            socket_fd_ = -1;
            return false;
        }
        
        ROS_INFO("✓ Connected to GO1 successfully!");
        return true;
    }
    
    void start()
    {
        if (!connectToGO1()) {
            ROS_ERROR("Cannot start - connection failed");
            return;
        }
        
        running_ = true;
        receive_thread_ = std::thread(&PCSlamwareReceiver::receiveLoop, this);
        ROS_INFO("Started receiving data from GO1");
    }
    
    void receiveLoop()
    {
        std::string buffer;
        char recv_buffer[4096];
        
        while (running_ && ros::ok()) {
            int bytes_received = recv(socket_fd_, recv_buffer, sizeof(recv_buffer) - 1, 0);
            if (bytes_received <= 0) {
                if (running_) {
                    ROS_WARN("Connection lost, attempting to reconnect...");
                    close(socket_fd_);
                    if (connectToGO1()) {
                        continue;
                    } else {
                        ROS_ERROR("Reconnection failed");
                        break;
                    }
                }
                break;
            }
            
            recv_buffer[bytes_received] = '\0';
            buffer += recv_buffer;
            
            // 处理完整的JSON消息（以换行符分隔）
            size_t pos;
            while ((pos = buffer.find('\n')) != std::string::npos) {
                std::string json_message = buffer.substr(0, pos);
                buffer.erase(0, pos + 1);
                
                if (!json_message.empty()) {
                    processMessage(json_message);
                }
            }
        }
    }
    
    void processMessage(const std::string& json_str)
    {
        try {
            Json::Value root;
            Json::Reader reader;
            
            if (!reader.parse(json_str, root)) {
                ROS_WARN_THROTTLE(5, "Failed to parse JSON message");
                return;
            }
            
            std::string type = root["type"].asString();
            
            if (type == "scan") {
                processScanData(root);
            } else if (type == "map") {
                processMapData(root);
            } else if (type == "odom") {
                processOdomData(root);
            } else if (type == "pose") {
                processPoseData(root);
            }
            
        } catch (const std::exception& e) {
            ROS_WARN_THROTTLE(5, "Error processing message: %s", e.what());
        }
    }
    
    void processScanData(const Json::Value& data)
    {
        sensor_msgs::LaserScan scan_msg;
        
        // 设置消息头
        scan_msg.header.stamp = ros::Time::now();  // 使用当前时间
        scan_msg.header.frame_id = data["frame_id"].asString();
        
        // 设置扫描参数
        scan_msg.angle_min = data["angle_min"].asFloat();
        scan_msg.angle_max = data["angle_max"].asFloat();
        scan_msg.angle_increment = data["angle_increment"].asFloat();
        scan_msg.time_increment = data["time_increment"].asFloat();
        scan_msg.scan_time = data["scan_time"].asFloat();
        scan_msg.range_min = data["range_min"].asFloat();
        scan_msg.range_max = data["range_max"].asFloat();
        
        // 转换距离数据
        const Json::Value& ranges = data["ranges"];
        scan_msg.ranges.resize(ranges.size());
        for (int i = 0; i < static_cast<int>(ranges.size()); ++i) {
            float range = ranges[i].asFloat();
            // 将无穷大值转换为NaN（RViz更好处理）
            if (range > 1e6) {
                scan_msg.ranges[i] = std::numeric_limits<float>::quiet_NaN();
            } else {
                scan_msg.ranges[i] = range;
            }
        }
        
        // 转换强度数据
        const Json::Value& intensities = data["intensities"];
        scan_msg.intensities.resize(intensities.size());
        for (int i = 0; i < static_cast<int>(intensities.size()); ++i) {
            scan_msg.intensities[i] = intensities[i].asFloat();
        }
        
        // 发布消息
        scan_pub_.publish(scan_msg);
        
        ROS_INFO_THROTTLE(2, "✓ Published scan data: %zu points, frame: %s", 
                         scan_msg.ranges.size(), scan_msg.header.frame_id.c_str());
    }
    
    void processMapData(const Json::Value& data)
    {
        nav_msgs::OccupancyGrid map_msg;
        
        // 设置消息头
        map_msg.header.stamp = ros::Time::now();
        map_msg.header.frame_id = data["frame_id"].asString();
        
        // 设置地图信息
        map_msg.info.width = data["width"].asUInt();
        map_msg.info.height = data["height"].asUInt();
        map_msg.info.resolution = data["resolution"].asFloat();
        
        // 设置原点
        const Json::Value& origin = data["origin"];
        map_msg.info.origin.position.x = origin["position"]["x"].asFloat();
        map_msg.info.origin.position.y = origin["position"]["y"].asFloat();
        map_msg.info.origin.position.z = origin["position"]["z"].asFloat();
        map_msg.info.origin.orientation.x = origin["orientation"]["x"].asFloat();
        map_msg.info.origin.orientation.y = origin["orientation"]["y"].asFloat();
        map_msg.info.origin.orientation.z = origin["orientation"]["z"].asFloat();
        map_msg.info.origin.orientation.w = origin["orientation"]["w"].asFloat();
        
        // 转换地图数据
        const Json::Value& map_data = data["data"];
        map_msg.data.resize(map_data.size());
        for (int i = 0; i < static_cast<int>(map_data.size()); ++i) {
            map_msg.data[i] = map_data[i].asInt();
        }
        
        // 发布消息
        map_pub_.publish(map_msg);
        
        ROS_INFO_THROTTLE(5, "✓ Published map: %dx%d (res: %.3fm)", 
                         map_msg.info.width, map_msg.info.height, map_msg.info.resolution);
    }
    
    void processOdomData(const Json::Value& data)
    {
        nav_msgs::Odometry odom_msg;
        
        // 设置消息头
        odom_msg.header.stamp = ros::Time::now();
        odom_msg.header.frame_id = data["frame_id"].asString();
        odom_msg.child_frame_id = data["child_frame_id"].asString();
        
        // 设置位置和姿态
        odom_msg.pose.pose.position.x = data["pose"]["position"]["x"].asFloat();
        odom_msg.pose.pose.position.y = data["pose"]["position"]["y"].asFloat();
        odom_msg.pose.pose.position.z = data["pose"]["position"]["z"].asFloat();
        odom_msg.pose.pose.orientation.x = data["pose"]["orientation"]["x"].asFloat();
        odom_msg.pose.pose.orientation.y = data["pose"]["orientation"]["y"].asFloat();
        odom_msg.pose.pose.orientation.z = data["pose"]["orientation"]["z"].asFloat();
        odom_msg.pose.pose.orientation.w = data["pose"]["orientation"]["w"].asFloat();
        
        // 设置速度
        odom_msg.twist.twist.linear.x = data["twist"]["linear"]["x"].asFloat();
        odom_msg.twist.twist.linear.y = data["twist"]["linear"]["y"].asFloat();
        odom_msg.twist.twist.linear.z = data["twist"]["linear"]["z"].asFloat();
        odom_msg.twist.twist.angular.x = data["twist"]["angular"]["x"].asFloat();
        odom_msg.twist.twist.angular.y = data["twist"]["angular"]["y"].asFloat();
        odom_msg.twist.twist.angular.z = data["twist"]["angular"]["z"].asFloat();
        
        // 发布消息
        odom_pub_.publish(odom_msg);
        
        ROS_INFO_THROTTLE(5, "✓ Published odometry");
    }
    
    void processPoseData(const Json::Value& data)
    {
        geometry_msgs::PoseWithCovarianceStamped pose_msg;
        
        // 设置消息头
        pose_msg.header.stamp = ros::Time::now();
        pose_msg.header.frame_id = data.get("frame_id", "map").asString();
        
        // 检查数据完整性
        if (!data.isMember("pose") || !data["pose"].isMember("position")) {
            ROS_WARN("Incomplete pose data received");
            return;
        }
        
        // 设置位置和姿态
        const Json::Value& pose = data["pose"];
        const Json::Value& position = pose["position"];
        const Json::Value& orientation = pose["orientation"];
        
        pose_msg.pose.pose.position.x = position.get("x", 0.0).asFloat();
        pose_msg.pose.pose.position.y = position.get("y", 0.0).asFloat();
        pose_msg.pose.pose.position.z = position.get("z", 0.0).asFloat();
        
        pose_msg.pose.pose.orientation.x = orientation.get("x", 0.0).asFloat();
        pose_msg.pose.pose.orientation.y = orientation.get("y", 0.0).asFloat();
        pose_msg.pose.pose.orientation.z = orientation.get("z", 0.0).asFloat();
        pose_msg.pose.pose.orientation.w = orientation.get("w", 1.0).asFloat();
        
        // 设置协方差矩阵（表示位置的不确定性）
        for (int i = 0; i < 36; ++i) {
            pose_msg.pose.covariance[i] = 0.0;
        }
        // 设置对角线元素
        pose_msg.pose.covariance[0] = 0.1;  // x
        pose_msg.pose.covariance[7] = 0.1;  // y
        pose_msg.pose.covariance[35] = 0.1; // yaw
        
        // 发布消息
        pose_pub_.publish(pose_msg);
        
        ROS_INFO_THROTTLE(2, "✓ Published robot pose: (%.2f, %.2f)", 
                 pose_msg.pose.pose.position.x, pose_msg.pose.pose.position.y);
    }
};

int main(int argc, char** argv)
{
    ros::init(argc, argv, "pc_slamware_receiver_node");
    
    PCSlamwareReceiver receiver;
    
    receiver.start();
    
    ros::spin();
    
    return 0;
} 