/*
 * Balance Chassis Module
 */

/* Includes ----------------------------------------------------------------- */
#include "module/balance_chassis.h"
#include "bsp/can.h"
#include "component/user_math.h"
#include "device/motor_lk.h"
#include "module/config.h"
#include <string.h>
#include <math.h>

/* Private typedef ---------------------------------------------------------- */
/* Private define ----------------------------------------------------------- */
#define CAN_CMD_ENABLE_ID      121  // 使能命令ID
#define CAN_CMD_JOINT_ID       122  // 关节控制命令ID  
#define CAN_CMD_WHEEL_LEFT_ID  0x128  // 左轮控制命令ID
#define CAN_CMD_WHEEL_RIGHT_ID 0x129  // 右轮控制命令ID

#define CAN_FEEDBACK_JOINT_BASE_ID 124  // 关节反馈基础ID (124-127)
#define CAN_FEEDBACK_WHEEL_LEFT_ID 130  // 左轮反馈ID
#define CAN_FEEDBACK_WHEEL_RIGHT_ID 131 // 右轮反馈ID

/* Private macro ------------------------------------------------------------ */
/* Private variables -------------------------------------------------------- */
static bool joint_command_received = false;
static bool wheel_command_received[2] = {false, false};

/* Private function --------------------------------------------------------- */

/**
 * @brief 检查并处理CAN控制命令
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK
 */
static int8_t Chassis_ProcessCANCommands(Chassis_t *chassis) {
    BSP_CAN_Message_t rx_msg;
    joint_command_received = false;
    wheel_command_received[0] = false;
    wheel_command_received[1] = false;
    
    // 检查ID 121 - 使能4个关节电机
    if (BSP_CAN_GetMessage(BSP_CAN_1, CAN_CMD_ENABLE_ID, &rx_msg, BSP_CAN_TIMEOUT_IMMEDIATE) == BSP_OK) {
        joint_command_received = true;
        for (int i = 0; i < 4; i++) {
            MOTOR_LZ_Enable(&chassis->param.joint_param[i]);
        }
    }
    
    // 检查ID 122 - 运控模式控制4个关节电机
    if (BSP_CAN_GetMessage(BSP_CAN_1, CAN_CMD_JOINT_ID, &rx_msg, BSP_CAN_TIMEOUT_IMMEDIATE) == BSP_OK) {
        joint_command_received = true;
        // 8字节数据分别是4个电机的力矩 (每个电机2字节，有符号整数，精度0.01 Nm)
        for (int i = 0; i < 4; i++) {
            int16_t torque_raw;
            memcpy(&torque_raw, &rx_msg.data[i * 2], sizeof(int16_t));
            float torque = (float)torque_raw / 100.0f; // 转换为浮点数力矩值
            
            // 使用运控模式控制电机，只设置力矩，其他参数为0
            MOTOR_LZ_MotionParam_t motion_param = {
                .target_angle = 0.0f,
                .target_velocity = 0.0f,
                .kp = 0.0f,
                .kd = 0.0f,
                .torque = torque
            };
            MOTOR_LZ_MotionControl(&chassis->param.joint_param[i], &motion_param);
        }
    }
    
    // 检查ID 128 - 左轮控制命令 (与电机发送格式一致)
    if (BSP_CAN_GetMessage(BSP_CAN_1, CAN_CMD_WHEEL_LEFT_ID, &rx_msg, BSP_CAN_TIMEOUT_IMMEDIATE) == BSP_OK) {
        wheel_command_received[0] = true;
        float left_out = (float)((int16_t)(rx_msg.data[4] | (rx_msg.data[5] << 8))) / 2048.0f;
        MOTOR_LK_SetOutput(&chassis->param.wheel_param[0], left_out);

    }
    
    // 检查ID 129 - 右轮控制命令 (与电机发送格式一致) 
    if (BSP_CAN_GetMessage(BSP_CAN_1, CAN_CMD_WHEEL_RIGHT_ID, &rx_msg, BSP_CAN_TIMEOUT_IMMEDIATE) == BSP_OK) {
        wheel_command_received[1] = true;
        float right_out = (float)((int16_t)(rx_msg.data[4] | (rx_msg.data[5] << 8))) / 2048.0f;
        MOTOR_LK_SetOutput(&chassis->param.wheel_param[1], right_out);
    }
    
    return DEVICE_OK;
}

/**
 * @brief 发送关节电机反馈数据
 * @param chassis 底盘实例  
 * @return 成功返回DEVICE_OK
 */
static int8_t Chassis_SendJointFeedback(Chassis_t *chassis) {
    // 发送4个关节电机的反馈数据，ID分别为124、125、126、127
    for (int i = 0; i < 4; i++) {
        MOTOR_LZ_t* motor = MOTOR_LZ_GetMotor(&chassis->param.joint_param[i]);
        if (motor != NULL) {
            BSP_CAN_StdDataFrame_t motor_frame = {
                .id = CAN_FEEDBACK_JOINT_BASE_ID + i,  // ID: 124, 125, 126, 127
                .dlc = 8,
                .data = {0}
            };
            
            // 数据重构：转矩电流(2字节) + 位置(3字节) + 速度(3字节) = 8字节
            // 转矩电流 - 转换为16位整数 (精度0.01 Nm)
            int16_t torque_int = (int16_t)(motor->lz_feedback.current_torque * 100);
            memcpy(&motor_frame.data[0], &torque_int, sizeof(int16_t));
            
            // 位置 - 转换为24位整数，使用3字节 (精度0.0001 rad)
            int32_t angle_int = (int32_t)(motor->lz_feedback.current_angle * 10000) & 0xFFFFFF;
            motor_frame.data[2] = (angle_int >> 16) & 0xFF;
            motor_frame.data[3] = (angle_int >> 8) & 0xFF;
            motor_frame.data[4] = angle_int & 0xFF;
            
            // 速度 - 转换为24位整数，使用3字节 (精度0.001 rad/s)
            int32_t velocity_int = (int32_t)(motor->lz_feedback.current_velocity * 1000) & 0xFFFFFF;
            motor_frame.data[5] = (velocity_int >> 16) & 0xFF;
            motor_frame.data[6] = (velocity_int >> 8) & 0xFF;
            motor_frame.data[7] = velocity_int & 0xFF;
            
            BSP_CAN_TransmitStdDataFrame(BSP_CAN_1, &motor_frame);
        }
    }
    
    return DEVICE_OK;
}

/**
 * @brief 发送轮子电机反馈数据
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK  
 */
static int8_t Chassis_SendWheelFeedback(Chassis_t *chassis) {
    // 发送2个轮子电机的反馈数据，ID分别为130、131
    for (int i = 0; i < 2; i++) {
        MOTOR_LK_t* motor = MOTOR_LK_GetMotor(&chassis->param.wheel_param[i]);
        if (motor != NULL) {
            BSP_CAN_StdDataFrame_t wheel_frame = {
                .id = CAN_FEEDBACK_WHEEL_LEFT_ID + i,  // ID: 130, 131
                .dlc = 8,
                .data = {0}
            };
            
            // 使用LK电机的标准反馈格式
            // 角度(2字节) + 速度(2字节) + 力矩电流(2字节) + 编码器(2字节)
            
            // 角度 - 转换为16位整数 (精度0.01度)
            int16_t angle_int = (int16_t)(motor->motor.feedback.rotor_abs_angle * 180.0f / M_PI * 100);
            wheel_frame.data[0] = (angle_int >> 8) & 0xFF;
            wheel_frame.data[1] = angle_int & 0xFF;
            
            // 速度 - 转换为16位整数 (精度1dps)  
            int16_t speed_int = (int16_t)(motor->motor.feedback.rotor_speed);
            wheel_frame.data[2] = (speed_int >> 8) & 0xFF;
            wheel_frame.data[3] = speed_int & 0xFF;
            
            // 力矩电流 - 转换为16位整数
            int16_t current_int = (int16_t)(motor->motor.feedback.torque_current);
            wheel_frame.data[4] = (current_int >> 8) & 0xFF;
            wheel_frame.data[5] = current_int & 0xFF;
            
            // 编码器值 - 直接使用角度值作为编码器
            uint16_t encoder = (uint16_t)(motor->motor.feedback.rotor_abs_angle / (2 * M_PI) * 65535);
            wheel_frame.data[6] = (encoder >> 8) & 0xFF;
            wheel_frame.data[7] = encoder & 0xFF;
            
            BSP_CAN_TransmitStdDataFrame(BSP_CAN_1, &wheel_frame);
        }
    }
    
    return DEVICE_OK;
}

/* Exported functions ------------------------------------------------------- */

/**
 * @brief 底盘初始化
 * @param chassis 底盘实例
 * @param param 底盘参数  
 * @return 成功返回DEVICE_OK
 */
int8_t Chassis_Init(Chassis_t *chassis, Chassis_Param_t *param) {
    if (chassis == NULL || param == NULL) {
        return DEVICE_ERR_NULL;
    }
    
    // 复制参数
    memcpy(&chassis->param, param, sizeof(Chassis_Param_t));
    
    // 初始化CAN
    BSP_CAN_Init();
    
    // 初始化电机驱动
    MOTOR_LZ_Init();
    
    // 注册关节电机
    for (int i = 0; i < 4; i++) {
        if (MOTOR_LZ_Register(&chassis->param.joint_param[i]) != DEVICE_OK) {
            return DEVICE_ERR;
        }
    }
    
    // 注册轮子电机
    for (int i = 0; i < 2; i++) {
        if (MOTOR_LK_Register(&chassis->param.wheel_param[i]) != DEVICE_OK) {
            return DEVICE_ERR;
        }
    }
    
    // 注册CAN接收ID
    BSP_CAN_RegisterId(BSP_CAN_1, CAN_CMD_ENABLE_ID, 0);      // 使能命令
    BSP_CAN_RegisterId(BSP_CAN_1, CAN_CMD_JOINT_ID, 0);       // 关节控制命令
    BSP_CAN_RegisterId(BSP_CAN_1, CAN_CMD_WHEEL_LEFT_ID, 0);  // 左轮控制命令
    BSP_CAN_RegisterId(BSP_CAN_1, CAN_CMD_WHEEL_RIGHT_ID, 0); // 右轮控制命令
    
    return DEVICE_OK;
}

/**
 * @brief 更新底盘数据
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK
 */
int8_t Chassis_Update(Chassis_t *chassis) {
    if (chassis == NULL) {
        return DEVICE_ERR_NULL;
    }
    
    // 更新所有电机数据
    MOTOR_LZ_UpdateAll();
    MOTOR_LK_UpdateAll();
    
    // 更新反馈数据
    for (int i = 0; i < 4; i++) {
        MOTOR_LZ_t* joint_motor = MOTOR_LZ_GetMotor(&chassis->param.joint_param[i]);
        if (joint_motor != NULL) {
            chassis->data.joint[i] = joint_motor->motor.feedback;
        }
    }
    
    for (int i = 0; i < 2; i++) {
        MOTOR_LK_t* wheel_motor = MOTOR_LK_GetMotor(&chassis->param.wheel_param[i]);
        if (wheel_motor != NULL) {
            chassis->data.wheel[i] = wheel_motor->motor.feedback;
        }
    }
    
    return DEVICE_OK;
}

/**
 * @brief 使能底盘
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK
 */
int8_t Chassis_Enable(Chassis_t *chassis) {
    if (chassis == NULL) {
        return DEVICE_ERR_NULL;
    }
    
    // 使能关节电机
    for (int i = 0; i < 4; i++) {
        MOTOR_LZ_Enable(&chassis->param.joint_param[i]);
    }
    
    // 开启轮子电机
    for (int i = 0; i < 2; i++) {
        MOTOR_LK_MotorOn(&chassis->param.wheel_param[i]);
    }
    
    return DEVICE_OK;
}

/**
 * @brief 底盘松弛
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK
 */
int8_t Chassis_Relax(Chassis_t *chassis) {
    if (chassis == NULL) {
        return DEVICE_ERR_NULL;
    }
    
    // 关节电机松弛
    for (int i = 0; i < 4; i++) {
        MOTOR_LZ_Relax(&chassis->param.joint_param[i]);
    }
    
    // 轮子电机松弛
    for (int i = 0; i < 2; i++) {
        MOTOR_LK_Relax(&chassis->param.wheel_param[i]);
    }
    
    return DEVICE_OK;
}

/**
 * @brief 底盘控制处理
 * @param chassis 底盘实例
 * @return 成功返回DEVICE_OK
 */
int8_t Chassis_Ctrl(Chassis_t *chassis) {
    if (chassis == NULL) {
        return DEVICE_ERR_NULL;
    }
    
    // 处理CAN控制命令
    Chassis_ProcessCANCommands(chassis);
    
    // 如果没有收到关节控制命令，关节电机进入relax模式
    if (!joint_command_received) {
        for (int i = 0; i < 4; i++) {
            MOTOR_LZ_Relax(&chassis->param.joint_param[i]);
        }
    }
    
    // 如果没有收到轮子控制命令，轮子电机进入relax模式
    for (int i = 0; i < 2; i++) {
        if (!wheel_command_received[i]) {
            MOTOR_LK_Relax(&chassis->param.wheel_param[i]);
        }
    }
    
    // 发送反馈数据
    Chassis_SendJointFeedback(chassis);
    Chassis_SendWheelFeedback(chassis);
    
    return DEVICE_OK;
}