#include <errno.h>
#include <inttypes.h>
#include <string.h>
#include <unistd.h>

#include "config.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_timer.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/event_groups.h"
#include "freertos/semphr.h"
#include "freertos/task.h"
#include "lwip/inet.h"
#include "lwip/sockets.h"
#include "nvs_flash.h"

#include "guguji_protocol.h"
#include "robot_config.h"
#include "rs00_mit.h"
#include "sensors.h"
#include "status_outputs.h"

static const char *TAG = "guguji_bridge";
#if !GUGUJI_WIFI_USE_AP
static const int WIFI_CONNECTED_BIT = BIT0;
#endif

static SemaphoreHandle_t g_command_mutex;
static SemaphoreHandle_t g_feedback_mutex;
static SemaphoreHandle_t g_sensor_mutex;
static SemaphoreHandle_t g_peer_mutex;
#if !GUGUJI_WIFI_USE_AP
static EventGroupHandle_t g_wifi_event_group;
#endif

static guguji_command_packet_t g_latest_command;
static bool g_have_command = false;
static uint32_t g_last_command_ms = 0;
static uint32_t g_latest_command_sequence = 0;

static rs00_feedback_t g_feedback[GUGUJI_JOINT_COUNT];
static guguji_sensor_state_t g_sensor_state;
static guguji_robot_state_t g_robot_state = GUGUJI_STATE_DISARMED;
static volatile uint32_t g_latest_fault_mask = GUGUJI_FAULT_IMU_OFFLINE | GUGUJI_FAULT_MAG_OFFLINE;
static struct sockaddr_storage g_last_peer_addr;
static socklen_t g_last_peer_len = 0;
static bool g_have_peer = false;

static uint32_t now_ms(void)
{
    return (uint32_t)(esp_timer_get_time() / 1000ULL);
}

static float clampf_local(float value, float min_value, float max_value)
{
    if (value < min_value) {
        return min_value;
    }
    if (value > max_value) {
        return max_value;
    }
    return value;
}

static int find_joint_by_motor_id(uint8_t motor_id)
{
    for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
        if (GUGUJI_JOINTS[i].can_id == motor_id) {
            return i;
        }
    }
    return -1;
}

static float joint_to_motor_angle(int index, float joint_angle)
{
    const guguji_joint_config_t *joint = &GUGUJI_JOINTS[index];
    return joint_angle * joint->direction + joint->zero_offset_rad;
}

static float motor_to_joint_angle(int index, float motor_angle)
{
    const guguji_joint_config_t *joint = &GUGUJI_JOINTS[index];
    return (motor_angle - joint->zero_offset_rad) * joint->direction;
}

static void stop_all_motors(void)
{
    for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
        rs00_mit_send_stop(GUGUJI_JOINTS[i].can_id);
    }
}

static void clear_all_motor_faults(void)
{
    for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
        rs00_mit_send_clear_fault(GUGUJI_JOINTS[i].can_id);
    }
}

static void enable_all_motors(void)
{
    for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
        // RS00 MIT 协议中 0 表示 MIT 运控模式；上电默认也是 MIT，这里再写一次便于调试。
        rs00_mit_send_set_mode(GUGUJI_JOINTS[i].can_id, 0);
        vTaskDelay(pdMS_TO_TICKS(2));
        rs00_mit_send_enable(GUGUJI_JOINTS[i].can_id);
        vTaskDelay(pdMS_TO_TICKS(2));
    }
}

#if !GUGUJI_WIFI_USE_AP
static void wifi_event_handler(
    void *arg,
    esp_event_base_t event_base,
    int32_t event_id,
    void *event_data)
{
    (void)arg;
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        ESP_LOGW(TAG, "WiFi 断开，尝试重连");
        esp_wifi_connect();
        xEventGroupClearBits(g_wifi_event_group, WIFI_CONNECTED_BIT);
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
        ESP_LOGI(TAG, "WiFi 已连接，IP=" IPSTR, IP2STR(&event->ip_info.ip));
        xEventGroupSetBits(g_wifi_event_group, WIFI_CONNECTED_BIT);
    }
}
#endif

static void wifi_init(void)
{
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

#if GUGUJI_WIFI_USE_AP
    esp_netif_create_default_wifi_ap();
    wifi_config_t wifi_config = {
        .ap = {
            .ssid = GUGUJI_WIFI_AP_SSID,
            .ssid_len = strlen(GUGUJI_WIFI_AP_SSID),
            .password = GUGUJI_WIFI_AP_PASSWORD,
            .channel = 6,
            .max_connection = 2,
            .authmode = WIFI_AUTH_WPA_WPA2_PSK,
        },
    };
    if (strlen(GUGUJI_WIFI_AP_PASSWORD) == 0) {
        wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    }
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_LOGI(TAG, "WiFi AP 已启动: ssid=%s", GUGUJI_WIFI_AP_SSID);
#else
    g_wifi_event_group = xEventGroupCreate();
    esp_netif_create_default_wifi_sta();
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, NULL));
    wifi_config_t wifi_config = {
        .sta = {
            .ssid = GUGUJI_WIFI_STA_SSID,
            .password = GUGUJI_WIFI_STA_PASSWORD,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    xEventGroupWaitBits(g_wifi_event_group, WIFI_CONNECTED_BIT, pdFALSE, pdTRUE, portMAX_DELAY);
#endif
}

static void udp_receive_task(void *arg)
{
    (void)arg;
    bool peer_logged = false;
    uint32_t last_logged_flags = UINT32_MAX;
    const int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0) {
        ESP_LOGE(TAG, "创建 UDP socket 失败: errno=%d", errno);
        vTaskDelete(NULL);
    }

    struct sockaddr_in listen_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(GUGUJI_UDP_COMMAND_PORT),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };
    if (bind(sock, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) < 0) {
        ESP_LOGE(TAG, "绑定 UDP 端口失败: errno=%d", errno);
        close(sock);
        vTaskDelete(NULL);
    }
    ESP_LOGI(TAG, "UDP 命令端口已监听: %d", GUGUJI_UDP_COMMAND_PORT);

    while (true) {
        guguji_command_packet_t packet = {0};
        struct sockaddr_storage source_addr = {0};
        socklen_t source_len = sizeof(source_addr);
        const int received = recvfrom(
            sock,
            &packet,
            sizeof(packet),
            0,
            (struct sockaddr *)&source_addr,
            &source_len);
        if (received < 0) {
            continue;
        }
        if (!guguji_validate_command_packet(&packet, (size_t)received)) {
            ESP_LOGW(TAG, "收到无效 UDP 命令包，长度=%d", received);
            continue;
        }

        xSemaphoreTake(g_command_mutex, portMAX_DELAY);
        memcpy(&g_latest_command, &packet, sizeof(packet));
        g_have_command = true;
        g_last_command_ms = now_ms();
        g_latest_command_sequence = packet.sequence;
        xSemaphoreGive(g_command_mutex);

        xSemaphoreTake(g_peer_mutex, portMAX_DELAY);
        memcpy(&g_last_peer_addr, &source_addr, source_len);
        g_last_peer_len = source_len;
        g_have_peer = true;
        xSemaphoreGive(g_peer_mutex);

        if (!peer_logged && source_addr.ss_family == AF_INET) {
            char addr_str[INET_ADDRSTRLEN] = {0};
            const struct sockaddr_in *source_ipv4 = (const struct sockaddr_in *)&source_addr;
            inet_ntoa_r(source_ipv4->sin_addr, addr_str, sizeof(addr_str));
            ESP_LOGI(TAG, "收到主机命令: %s:%d", addr_str, ntohs(source_ipv4->sin_port));
            peer_logged = true;
        }
        if (packet.flags != last_logged_flags) {
            ESP_LOGI(TAG, "主机命令 flags=0x%08" PRIX32 " seq=%" PRIu32, packet.flags, packet.sequence);
            last_logged_flags = packet.flags;
        }
    }
}

static void can_feedback_task(void *arg)
{
    (void)arg;
    while (true) {
        rs00_feedback_t feedback = {0};
        if (!rs00_mit_poll_feedback(&feedback, 10)) {
            continue;
        }
        const int joint_index = find_joint_by_motor_id(feedback.motor_id);
        if (joint_index < 0) {
            continue;
        }

        xSemaphoreTake(g_feedback_mutex, portMAX_DELAY);
        g_feedback[joint_index] = feedback;
        xSemaphoreGive(g_feedback_mutex);
    }
}

static void sensor_task(void *arg)
{
    (void)arg;
    uint32_t last_ms = now_ms();
    while (true) {
        const uint32_t current_ms = now_ms();
        const float dt_s = (float)(current_ms - last_ms) * 0.001f;
        last_ms = current_ms;

        xSemaphoreTake(g_sensor_mutex, portMAX_DELAY);
        guguji_sensors_update(&g_sensor_state, dt_s > 0.0f ? dt_s : 0.02f);
        xSemaphoreGive(g_sensor_mutex);

        vTaskDelay(pdMS_TO_TICKS(10));
    }
}

static void control_task(void *arg)
{
    (void)arg;
    bool motors_enabled = false;
    uint32_t last_clear_fault_sequence = UINT32_MAX;

    while (true) {
        guguji_command_packet_t command = {0};
        bool have_command = false;
        uint32_t command_age_ms = UINT32_MAX;

        xSemaphoreTake(g_command_mutex, portMAX_DELAY);
        have_command = g_have_command;
        if (have_command) {
            command = g_latest_command;
            command_age_ms = now_ms() - g_last_command_ms;
        }
        xSemaphoreGive(g_command_mutex);

        if (!have_command || command_age_ms > GUGUJI_COMMAND_TIMEOUT_MS) {
            if (motors_enabled) {
                stop_all_motors();
                motors_enabled = false;
            }
            g_robot_state = have_command ? GUGUJI_STATE_TIMEOUT : GUGUJI_STATE_DISARMED;
            vTaskDelay(pdMS_TO_TICKS(GUGUJI_CONTROL_PERIOD_MS));
            continue;
        }

        if ((command.flags & GUGUJI_COMMAND_FLAG_CLEAR_FAULTS) != 0 &&
            command.sequence != last_clear_fault_sequence) {
            clear_all_motor_faults();
            last_clear_fault_sequence = command.sequence;
        }

        if ((command.flags & GUGUJI_COMMAND_FLAG_ESTOP) != 0) {
            stop_all_motors();
            motors_enabled = false;
            g_robot_state = GUGUJI_STATE_ESTOP;
            vTaskDelay(pdMS_TO_TICKS(GUGUJI_CONTROL_PERIOD_MS));
            continue;
        }

        if ((command.flags & GUGUJI_COMMAND_FLAG_DISARM) != 0 ||
            (command.flags & GUGUJI_COMMAND_FLAG_ARM) == 0) {
            if (motors_enabled) {
                stop_all_motors();
                motors_enabled = false;
            }
            g_robot_state = GUGUJI_STATE_DISARMED;
            vTaskDelay(pdMS_TO_TICKS(GUGUJI_CONTROL_PERIOD_MS));
            continue;
        }

        if (!motors_enabled) {
            enable_all_motors();
            motors_enabled = true;
            ESP_LOGI(TAG, "电机已使能");
        }

        g_robot_state = GUGUJI_STATE_ARMED;
        for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
            const guguji_joint_config_t *joint = &GUGUJI_JOINTS[i];
            const float joint_target = clampf_local(
                command.positions_rad[i],
                joint->lower_limit_rad,
                joint->upper_limit_rad);
            const float motor_position = joint_to_motor_angle(i, joint_target);
            const float motor_velocity = command.velocities_rad_s[i] * joint->direction;
            const float motor_torque = command.torque_ff_nm[i] * joint->direction;
            rs00_mit_send_control(
                joint->can_id,
                motor_position,
                motor_velocity,
                command.kp[i],
                command.kd[i],
                motor_torque);
        }

        vTaskDelay(pdMS_TO_TICKS(GUGUJI_CONTROL_PERIOD_MS));
    }
}

static void telemetry_task(void *arg)
{
    (void)arg;
    const int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0) {
        ESP_LOGE(TAG, "创建遥测 UDP socket 失败: errno=%d", errno);
        vTaskDelete(NULL);
    }

    while (true) {
        struct sockaddr_storage peer_addr = {0};
        socklen_t peer_len = 0;
        bool have_peer = false;

        xSemaphoreTake(g_peer_mutex, portMAX_DELAY);
        have_peer = g_have_peer;
        if (have_peer) {
            memcpy(&peer_addr, &g_last_peer_addr, g_last_peer_len);
            peer_len = g_last_peer_len;
        }
        xSemaphoreGive(g_peer_mutex);

        if (!have_peer) {
            vTaskDelay(pdMS_TO_TICKS(GUGUJI_TELEMETRY_PERIOD_MS));
            continue;
        }

        guguji_telemetry_packet_t packet = {0};
        packet.robot_state = (uint8_t)g_robot_state;
        packet.sequence = g_latest_command_sequence;
        packet.uptime_ms = now_ms();

        xSemaphoreTake(g_command_mutex, portMAX_DELAY);
        packet.last_command_age_ms = g_have_command ? (now_ms() - g_last_command_ms) : UINT32_MAX;
        xSemaphoreGive(g_command_mutex);

        uint32_t fault_mask = 0;
        const uint32_t telemetry_now_ms = packet.uptime_ms;
        xSemaphoreTake(g_feedback_mutex, portMAX_DELAY);
        for (int i = 0; i < GUGUJI_JOINT_COUNT; ++i) {
            const rs00_feedback_t feedback = g_feedback[i];
            packet.joint_position_rad[i] = motor_to_joint_angle(i, feedback.position_rad);
            packet.joint_velocity_rad_s[i] = feedback.velocity_rad_s * GUGUJI_JOINTS[i].direction;
            packet.joint_torque_nm[i] = feedback.torque_nm * GUGUJI_JOINTS[i].direction;
            packet.motor_temperature_c[i] = feedback.temperature_c;
            if (feedback.fault) {
                fault_mask |= (1u << i);
            }
            if (feedback.last_feedback_ms == 0 ||
                telemetry_now_ms - feedback.last_feedback_ms > GUGUJI_MOTOR_FEEDBACK_TIMEOUT_MS) {
                fault_mask |= (1u << (8 + i));
            }
        }
        xSemaphoreGive(g_feedback_mutex);

        xSemaphoreTake(g_sensor_mutex, portMAX_DELAY);
        memcpy(packet.accel_m_s2, g_sensor_state.accel_m_s2, sizeof(packet.accel_m_s2));
        memcpy(packet.gyro_rad_s, g_sensor_state.gyro_rad_s, sizeof(packet.gyro_rad_s));
        memcpy(packet.mag_u_t, g_sensor_state.mag_u_t, sizeof(packet.mag_u_t));
        memcpy(packet.rpy_rad, g_sensor_state.rpy_rad, sizeof(packet.rpy_rad));
        if (!g_sensor_state.imu_ok) {
            fault_mask |= GUGUJI_FAULT_IMU_OFFLINE;
        }
        if (!g_sensor_state.mag_ok) {
            fault_mask |= GUGUJI_FAULT_MAG_OFFLINE;
        }
        xSemaphoreGive(g_sensor_mutex);

        g_latest_fault_mask = fault_mask;
        packet.fault_mask = fault_mask;

        guguji_finalize_telemetry_packet(&packet);
        sendto(sock, &packet, sizeof(packet), 0, (struct sockaddr *)&peer_addr, peer_len);
        vTaskDelay(pdMS_TO_TICKS(GUGUJI_TELEMETRY_PERIOD_MS));
    }
}

static void status_task(void *arg)
{
    (void)arg;
    while (true) {
        bool have_peer = false;
        xSemaphoreTake(g_peer_mutex, portMAX_DELAY);
        have_peer = g_have_peer;
        xSemaphoreGive(g_peer_mutex);

        const uint32_t fault_mask = g_latest_fault_mask;
        const uint32_t motor_faults = fault_mask & GUGUJI_FAULT_MOTOR_MASK;
        const uint32_t warning_faults = fault_mask & (
            GUGUJI_FAULT_FEEDBACK_TIMEOUT_MASK |
            GUGUJI_FAULT_IMU_OFFLINE |
            GUGUJI_FAULT_MAG_OFFLINE);
        guguji_status_mode_t mode = GUGUJI_STATUS_DISARMED;

        if (motor_faults != 0) {
            mode = GUGUJI_STATUS_FAULT;
        } else if (g_robot_state == GUGUJI_STATE_ESTOP) {
            mode = GUGUJI_STATUS_ESTOP;
        } else if (g_robot_state == GUGUJI_STATE_TIMEOUT) {
            mode = GUGUJI_STATUS_TIMEOUT;
        } else if (!have_peer) {
            mode = GUGUJI_STATUS_WAITING_HOST;
        } else if (g_robot_state == GUGUJI_STATE_ARMED) {
            mode = GUGUJI_STATUS_ARMED;
        } else if (warning_faults != 0) {
            mode = GUGUJI_STATUS_SENSOR_WARN;
        } else {
            mode = GUGUJI_STATUS_DISARMED;
        }

        guguji_status_outputs_apply(mode, now_ms());
        vTaskDelay(pdMS_TO_TICKS(80));
    }
}

void app_main(void)
{
    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ESP_ERROR_CHECK(nvs_flash_init());
    } else {
        ESP_ERROR_CHECK(err);
    }

    g_command_mutex = xSemaphoreCreateMutex();
    g_feedback_mutex = xSemaphoreCreateMutex();
    g_sensor_mutex = xSemaphoreCreateMutex();
    g_peer_mutex = xSemaphoreCreateMutex();
    memset(&g_sensor_state, 0, sizeof(g_sensor_state));
    memset(&g_feedback, 0, sizeof(g_feedback));

    ESP_ERROR_CHECK(guguji_status_outputs_init());
    wifi_init();
    ESP_ERROR_CHECK(guguji_sensors_init());
    ESP_ERROR_CHECK(rs00_mit_init(GUGUJI_TWAI_TX_GPIO, GUGUJI_TWAI_RX_GPIO));

    xTaskCreatePinnedToCore(udp_receive_task, "udp_rx", 4096, NULL, 6, NULL, 0);
    xTaskCreatePinnedToCore(control_task, "control", 4096, NULL, 8, NULL, 1);
    xTaskCreatePinnedToCore(can_feedback_task, "can_rx", 4096, NULL, 7, NULL, 1);
    xTaskCreatePinnedToCore(sensor_task, "sensors", 4096, NULL, 5, NULL, 0);
    xTaskCreatePinnedToCore(telemetry_task, "telemetry", 4096, NULL, 5, NULL, 0);
    xTaskCreatePinnedToCore(status_task, "status", 4096, NULL, 4, NULL, 0);

    ESP_LOGI(TAG, "guguji ESP32S3 real-robot bridge started");
}
