#include "urdf_parser.hpp"
#include <fstream>
#include <sstream>
#include <algorithm>
#include <stdexcept>
#include <iostream>

namespace robot_model
{

std::shared_ptr<RobotModel> UrdfParser::parse_from_file(const std::string& urdf_file_path)
{
    // 读取URDF文件内容
    std::ifstream file(urdf_file_path);
    if (!file.is_open())
    {
        std::cerr << "Failed to open URDF file: " << urdf_file_path << std::endl;
        return nullptr;
    }

    std::stringstream ss;
    ss << file.rdbuf();
    std::string urdf_xml = ss.str();
    file.close();

    return parse_from_string(urdf_xml);
}

std::shared_ptr<RobotModel> UrdfParser::parse_from_string(const std::string& urdf_xml)
{
    urdf::Model model;
    if (!model.initString(urdf_xml))
    {
        std::cerr << "Failed to parse URDF XML string" << std::endl;
        return nullptr;
    }

    return build_robot_model(model);
}

std::shared_ptr<RobotModel> UrdfParser::build_robot_model(const urdf::Model& model)
{
    auto robot_model = std::make_shared<RobotModel>();
    
    // 我们需要访问RobotModel的私有成员来构建模型
    // 这里假设RobotModel类有一个友元类UrdfParser，或者提供了构建接口
    // 为了演示，我们直接修改RobotModel的私有成员（实际项目中应使用适当的接口）
    
    // 遍历所有关节和连杆
    for (const auto& joint_pair : model.joints_)
    {
        const auto& joint = joint_pair.second;
        auto parsed_joint = parse_joint(joint);
        if (!parsed_joint)
            continue;

        // 确定关节属于哪个部分
        std::string part = identify_joint_part(joint->name);

        // 手臂关节
        if (part == "arm")
        {
            robot_model->arm_.joints.push_back(parsed_joint);
        }
        // 腿部关节
        else if (part == "leg")
        {
            std::string leg_side = identify_leg_side(joint->name);
            
            // 查找或创建对应的腿
            bool leg_found = false;
            for (auto& leg : robot_model->legs_)
            {
                if (leg.side == leg_side)
                {
                    leg.joints.push_back(parsed_joint);
                    leg_found = true;
                    break;
                }
            }
            
            // 如果没有找到对应的腿，则创建新的
            if (!leg_found)
            {
                Leg new_leg;
                new_leg.side = leg_side;
                new_leg.joints.push_back(parsed_joint);
                robot_model->legs_.push_back(new_leg);
            }
        }
    }

    // 解析连杆
    for (const auto& link_pair : model.links_)
    {
        const auto& link = link_pair.second;
        auto parsed_link = parse_link(link);
        if (!parsed_link)
            continue;

        // 确定连杆属于哪个部分
        std::string part = identify_joint_part(link->name);

        // 手臂连杆
        if (part == "arm")
        {
            robot_model->arm_.links.push_back(parsed_link);
        }
        // 腿部连杆
        else if (part == "leg")
        {
            std::string leg_side = identify_leg_side(link->name);
            
            // 查找对应的腿
            for (auto& leg : robot_model->legs_)
            {
                if (leg.side == leg_side)
                {
                    leg.links.push_back(parsed_link);
                    break;
                }
            }
        }
    }

    // 确保腿部数量为4
    if (robot_model->legs_.size() != 4)
    {
        std::cerr << "Warning: Expected 4 legs, but found " << robot_model->legs_.size() << std::endl;
    }

    return robot_model;
}

std::shared_ptr<Link> UrdfParser::parse_link(const urdf::LinkConstSharedPtr& link)
{
    if (!link)
        return nullptr;

    // 获取连杆长度（简化处理，假设是盒子形状）
    double length = 0.1;  // 默认长度
    if (link->collision && link->collision->geometry && 
        link->collision->geometry->type == urdf::Geometry::BOX)
    {
        auto box = std::dynamic_pointer_cast<urdf::Box>(link->collision->geometry);
        if (box)
        {
            // 假设X轴为长度方向
            length = box->dim.x;
        }
    }

    // 获取质量
    double mass = 1.0;  // 默认质量
    if (link->inertial)
    {
        mass = link->inertial->mass;
    }

    return std::make_shared<Link>(link->name, length, mass);
}

std::shared_ptr<Joint> UrdfParser::parse_joint(const urdf::JointConstSharedPtr& joint)
{
    if (!joint)
        return nullptr;

    // 只处理旋转关节
    if (joint->type != urdf::Joint::REVOLUTE && joint->type != urdf::Joint::CONTINUOUS)
        return nullptr;

    // 获取关节角度限制
    double min_angle = -M_PI;
    double max_angle = M_PI;
    
    if (joint->type == urdf::Joint::REVOLUTE)
    {
        min_angle = joint->limits->lower;
        max_angle = joint->limits->upper;
    }

    // 获取力限制
    double effort_limit = 100.0;  // 默认力限制
    if (joint->limits)
    {
        effort_limit = joint->limits->effort;
    }

    return std::make_shared<Joint>(joint->name, min_angle, max_angle, effort_limit);
}

std::string UrdfParser::identify_joint_part(const std::string& name)
{
    // 转换为小写以便不区分大小写匹配
    std::string lower_name = name;
    std::transform(lower_name.begin(), lower_name.end(), lower_name.begin(), ::tolower);

    if (lower_name.find("arm") != std::string::npos || 
        lower_name.find("shoulder") != std::string::npos ||
        lower_name.find("elbow") != std::string::npos ||
        lower_name.find("wrist") != std::string::npos)
    {
        return "arm";
    }
    else if (lower_name.find("leg") != std::string::npos || 
             lower_name.find("hip") != std::string::npos ||
             lower_name.find("knee") != std::string::npos ||
             lower_name.find("thigh") != std::string::npos ||
             lower_name.find("shin") != std::string::npos)
    {
        return "leg";
    }

    return "unknown";
}

std::string UrdfParser::identify_leg_side(const std::string& name)
{
    // 转换为小写以便不区分大小写匹配
    std::string lower_name = name;
    std::transform(lower_name.begin(), lower_name.end(), lower_name.begin(), ::tolower);

    bool is_front = (lower_name.find("front") != std::string::npos || 
                     lower_name.find("fore") != std::string::npos);
    bool is_rear = (lower_name.find("rear") != std::string::npos || 
                    lower_name.find("hind") != std::string::npos);
    bool is_left = (lower_name.find("left") != std::string::npos);
    bool is_right = (lower_name.find("right") != std::string::npos);

    std::string position;
    if (is_front)
        position = "front_";
    else if (is_rear)
        position = "rear_";
    else
        position = "unknown_";

    std::string side;
    if (is_left)
        side = "left";
    else if (is_right)
        side = "right";
    else
        side = "unknown";

    return position + side;
}

}  // namespace robot_model
