/**
 * DreamStack ESP-NOW Transport — Implementation
 *
 * Handles ESP-NOW receive/send for binary signal frames,
 * and UDP listener for IR JSON push.
 */

#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "esp_now.h"
#include "lwip/sockets.h"

#include "ds_espnow.h"

static const char *TAG = "ds-espnow";

// ─── State ───
static ds_espnow_config_t s_config;
static uint8_t s_seq = 0;
static int s_udp_sock = -1;
static TaskHandle_t s_udp_task = NULL;

// ─── IR fragment reassembly ───
#define MAX_IR_SIZE 16384
static uint8_t s_ir_buf[MAX_IR_SIZE];
static size_t  s_ir_len = 0;
static uint8_t s_frag_received = 0;
static uint8_t s_frag_total = 0;
static uint8_t s_frag_seq = 0xFF;

// ─── ESP-NOW Receive Callback ───
static void espnow_recv_cb(const esp_now_recv_info_t *recv_info,
                            const uint8_t *data, int len) {
    if (len < 1) return;

    uint8_t type = data[0];

    switch (type) {
    case DS_NOW_SIG:
        if (len >= sizeof(ds_sig_frame_t)) {
            const ds_sig_frame_t *f = (const ds_sig_frame_t *)data;
            if (s_config.on_signal) {
                s_config.on_signal(f->signal_id, f->value);
            }
        }
        break;

    case DS_NOW_SIG_BATCH:
        if (len >= sizeof(ds_sig_batch_t)) {
            const ds_sig_batch_t *b = (const ds_sig_batch_t *)data;
            const ds_sig_entry_t *entries = (const ds_sig_entry_t *)(data + sizeof(ds_sig_batch_t));
            size_t expected = sizeof(ds_sig_batch_t) + b->count * sizeof(ds_sig_entry_t);
            if (len >= expected && s_config.on_signal) {
                for (int i = 0; i < b->count; i++) {
                    s_config.on_signal(entries[i].id, entries[i].val);
                }
            }
        }
        break;

    case DS_NOW_PING: {
        // Respond with pong
        ds_heartbeat_t pong = { .type = DS_NOW_PONG, .seq = data[1] };
        esp_now_send(recv_info->src_addr, (const uint8_t *)&pong, sizeof(pong));
        break;
    }

    case DS_NOW_PONG:
        ESP_LOGD(TAG, "Pong received (seq=%d)", data[1]);
        break;

    default:
        ESP_LOGW(TAG, "Unknown ESP-NOW frame type: 0x%02x", type);
        break;
    }
}

// ─── ESP-NOW Send Callback ───
static void espnow_send_cb(const uint8_t *mac_addr, esp_now_send_status_t status) {
    if (status != ESP_NOW_SEND_SUCCESS) {
        ESP_LOGW(TAG, "ESP-NOW send failed");
    }
}

// ─── UDP Listener Task ───
// Receives IR JSON push and fragmented IR over UDP
static void udp_listener_task(void *arg) {
    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port = htons(DS_UDP_PORT),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };

    s_udp_sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (s_udp_sock < 0) {
        ESP_LOGE(TAG, "Failed to create UDP socket");
        vTaskDelete(NULL);
        return;
    }

    if (bind(s_udp_sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "Failed to bind UDP port %d", DS_UDP_PORT);
        close(s_udp_sock);
        s_udp_sock = -1;
        vTaskDelete(NULL);
        return;
    }

    ESP_LOGI(TAG, "UDP listener on port %d", DS_UDP_PORT);

    uint8_t buf[1500];
    while (1) {
        int len = recvfrom(s_udp_sock, buf, sizeof(buf), 0, NULL, NULL);
        if (len < 4) continue;

        // Check magic
        if (buf[0] != 0xD5 || buf[1] != 0x7A) continue;

        uint8_t frame_type = buf[2];

        if (frame_type == DS_UDP_IR_PUSH) {
            // Non-fragmented IR push
            const ds_ir_push_t *hdr = (const ds_ir_push_t *)buf;
            size_t json_len = hdr->length;
            if (json_len <= len - sizeof(ds_ir_push_t) && json_len < MAX_IR_SIZE) {
                memcpy(s_ir_buf, buf + sizeof(ds_ir_push_t), json_len);
                s_ir_buf[json_len] = '\0';
                ESP_LOGI(TAG, "IR push received (%zu bytes)", json_len);
                if (s_config.on_ir_push) {
                    s_config.on_ir_push((const char *)s_ir_buf, json_len);
                }
            }
        } else if (frame_type == DS_UDP_IR_FRAG) {
            // Fragmented IR push
            if (len < sizeof(ds_ir_frag_t)) continue;
            const ds_ir_frag_t *frag = (const ds_ir_frag_t *)buf;
            size_t payload_len = len - sizeof(ds_ir_frag_t);
            const uint8_t *payload = buf + sizeof(ds_ir_frag_t);

            // New fragment group?
            if (frag->seq != s_frag_seq) {
                s_frag_seq = frag->seq;
                s_frag_received = 0;
                s_frag_total = frag->frag_total;
                s_ir_len = 0;
            }

            // Append fragment (assume ordered delivery)
            if (s_ir_len + payload_len < MAX_IR_SIZE) {
                memcpy(s_ir_buf + s_ir_len, payload, payload_len);
                s_ir_len += payload_len;
                s_frag_received++;

                if (s_frag_received >= s_frag_total) {
                    s_ir_buf[s_ir_len] = '\0';
                    ESP_LOGI(TAG, "IR reassembled (%zu bytes, %d frags)",
                             s_ir_len, s_frag_total);
                    if (s_config.on_ir_push) {
                        s_config.on_ir_push((const char *)s_ir_buf, s_ir_len);
                    }
                }
            }
        }
    }
}

// ─── Public API ───

esp_err_t ds_espnow_init(const ds_espnow_config_t *config) {
    s_config = *config;
    if (s_config.channel == 0) s_config.channel = DS_ESPNOW_CHANNEL;

    // Initialize ESP-NOW
    esp_err_t ret = esp_now_init();
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "ESP-NOW init failed: %s", esp_err_to_name(ret));
        return ret;
    }

    esp_now_register_recv_cb(espnow_recv_cb);
    esp_now_register_send_cb(espnow_send_cb);

    // Add hub as peer
    esp_now_peer_info_t peer = {
        .channel = s_config.channel,
        .ifidx = WIFI_IF_STA,
        .encrypt = false,
    };
    memcpy(peer.peer_addr, s_config.hub_mac, 6);
    ret = esp_now_add_peer(&peer);
    if (ret != ESP_OK) {
        ESP_LOGW(TAG, "Add peer failed (may already exist): %s", esp_err_to_name(ret));
    }

    ESP_LOGI(TAG, "ESP-NOW initialized (ch=%d, hub=%02x:%02x:%02x:%02x:%02x:%02x)",
             s_config.channel,
             s_config.hub_mac[0], s_config.hub_mac[1], s_config.hub_mac[2],
             s_config.hub_mac[3], s_config.hub_mac[4], s_config.hub_mac[5]);

    // Start UDP listener
    xTaskCreate(udp_listener_task, "ds_udp", 4096, NULL, 5, &s_udp_task);

    return ESP_OK;
}

esp_err_t ds_espnow_send_action(uint8_t node_id, uint8_t action) {
    ds_action_frame_t frame = {
        .type = DS_NOW_ACTION,
        .node_id = node_id,
        .action = action,
        .seq = s_seq++,
    };
    return esp_now_send(s_config.hub_mac, (const uint8_t *)&frame, sizeof(frame));
}

esp_err_t ds_espnow_send_touch(uint8_t node_id, uint8_t event,
                                uint16_t x, uint16_t y) {
    ds_touch_now_t frame = {
        .type = DS_NOW_TOUCH,
        .node_id = node_id,
        .event = event,
        .seq = s_seq++,
        .x = x,
        .y = y,
    };
    return esp_now_send(s_config.hub_mac, (const uint8_t *)&frame, sizeof(frame));
}

esp_err_t ds_espnow_send_ping(void) {
    ds_heartbeat_t ping = { .type = DS_NOW_PING, .seq = s_seq++ };
    return esp_now_send(s_config.hub_mac, (const uint8_t *)&ping, sizeof(ping));
}

void ds_espnow_deinit(void) {
    if (s_udp_task) {
        vTaskDelete(s_udp_task);
        s_udp_task = NULL;
    }
    if (s_udp_sock >= 0) {
        close(s_udp_sock);
        s_udp_sock = -1;
    }
    esp_now_deinit();
    ESP_LOGI(TAG, "ESP-NOW deinitialized");
}