// espnow_slave.c — ESP-NOW slave implementation (shared by all V3 puzzle nodes)
// Requires Wi-Fi to be initialised (WIFI_MODE_STA, no AP connection needed).
#include "espnow_slave.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "esp_mac.h"          // MACSTR / MAC2STR
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/task.h"
#include <string.h>
#include <stdlib.h>

static const char *TAG = "espnow_slave";

#define RECV_QUEUE_LEN 8

typedef struct {
    uint8_t mac[6];
    uint8_t data[250];
    int     data_len;
} recv_item_t;

static QueueHandle_t         s_recv_queue;
static uint8_t               s_master_mac[6];
static uint8_t               s_puzzle_id;
static espnow_cmd_callback_t s_callback;

// ---------------------------------------------------------------------------
// Scenario hot-load reassembly (see espnow_slave.h). Demuxed inside
// espnow_slave_process(); runs in task context, never the Wi-Fi ISR.
// ---------------------------------------------------------------------------
#define SCENARIO_PAYLOAD_MAX  236                   // 240 ESP-NOW max - 4 header
#define SCENARIO_MAX_BYTES    (64 * 1024)           // matches master cap
#define SCENARIO_MAX_FRAMES   ((SCENARIO_MAX_BYTES + SCENARIO_PAYLOAD_MAX - 1) \
                                / SCENARIO_PAYLOAD_MAX)
#define SCENARIO_TIMEOUT_MS   5000

typedef struct {
    bool      active;
    uint8_t   src[6];
    uint16_t  total;        // expected frame count
    uint16_t  count;        // distinct frames received
    size_t    tail;         // highest payload end offset = reassembled length
    uint8_t  *buf;          // heap, total * SCENARIO_PAYLOAD_MAX
    uint8_t  *seen;         // heap bitmap, (total + 7) / 8 bytes
    uint32_t  last_ms;
} scenario_rx_t;

static scenario_rx_t              s_scn;
static espnow_scenario_callback_t s_scn_cb;

static inline uint32_t now_ms(void)
{
    return (uint32_t)(xTaskGetTickCount() * portTICK_PERIOD_MS);
}

static void scenario_reset(void)
{
    free(s_scn.buf);
    free(s_scn.seen);
    memset(&s_scn, 0, sizeof(s_scn));
}

// True if this frame belongs to the scenario stream: either a continuation of
// the active transfer (matching sender MAC), or the first frame of a new one
// (seq==0, total>=1). MSG_* commands (data[0] in 0x01..0x08) never match.
static bool scenario_is_frame(const recv_item_t *it)
{
    if (s_scn.active && memcmp(it->mac, s_scn.src, 6) == 0)
        return true;
    if (it->data_len >= 4 && it->data[0] == 0x00 && it->data[1] == 0x00) {
        uint16_t total = (uint16_t)it->data[2] | ((uint16_t)it->data[3] << 8);
        return total >= 1;
    }
    return false;
}

static void scenario_feed(const recv_item_t *it)
{
    uint16_t seq   = (uint16_t)it->data[0] | ((uint16_t)it->data[1] << 8);
    uint16_t total = (uint16_t)it->data[2] | ((uint16_t)it->data[3] << 8);
    const uint8_t *payload = it->data + 4;
    int payload_len = it->data_len - 4;
    if (payload_len < 0) return;

    // Begin (or restart from a new sender) on the first frame.
    if (seq == 0 && (!s_scn.active || memcmp(it->mac, s_scn.src, 6) != 0)) {
        scenario_reset();
        if (total == 0 || total > SCENARIO_MAX_FRAMES) {
            ESP_LOGW(TAG, "scenario rx: bad/oversized total %u — dropped", total);
            return;
        }
        s_scn.buf  = calloc(total, SCENARIO_PAYLOAD_MAX);
        s_scn.seen = calloc((total + 7) / 8, 1);
        if (!s_scn.buf || !s_scn.seen) {
            ESP_LOGE(TAG, "scenario rx: OOM for %u frames — dropped", total);
            scenario_reset();
            return;
        }
        s_scn.active = true;
        s_scn.total  = total;
        memcpy(s_scn.src, it->mac, 6);
    }

    if (!s_scn.active) return;          // stray non-first frame, no live session
    if (total != s_scn.total) return;   // inconsistent header, ignore frame
    if (seq >= s_scn.total) return;
    if (payload_len > SCENARIO_PAYLOAD_MAX) payload_len = SCENARIO_PAYLOAD_MAX;

    s_scn.last_ms = now_ms();

    if (!(s_scn.seen[seq / 8] & (1u << (seq % 8)))) {
        size_t off = (size_t)seq * SCENARIO_PAYLOAD_MAX;
        memcpy(s_scn.buf + off, payload, (size_t)payload_len);
        s_scn.seen[seq / 8] |= (uint8_t)(1u << (seq % 8));
        s_scn.count++;
        size_t end = off + (size_t)payload_len;
        if (end > s_scn.tail) s_scn.tail = end;
    }

    if (s_scn.count == s_scn.total) {
        ESP_LOGI(TAG, "scenario reassembled: %u frames, %u bytes",
                 s_scn.total, (unsigned)s_scn.tail);
        if (s_scn_cb) {
            s_scn_cb(s_scn.src, s_scn.buf, s_scn.tail);
        } else {
            ESP_LOGI(TAG, "no scenario consumer on this node — dropped");
        }
        scenario_reset();
    }
}

// ---------------------------------------------------------------------------
// Internal callbacks (called from Wi-Fi ISR context)
// ---------------------------------------------------------------------------

static void on_recv(const esp_now_recv_info_t *info,
                    const uint8_t *data, int len)
{
    recv_item_t item;
    memcpy(item.mac, info->src_addr, 6);
    int copy_len = len < (int)sizeof(item.data) ? len : (int)sizeof(item.data);
    memcpy(item.data, data, copy_len);
    item.data_len = copy_len;
    xQueueSendFromISR(s_recv_queue, &item, NULL);
}

static void on_sent(const uint8_t *mac, esp_now_send_status_t status)
{
    if (status != ESP_NOW_SEND_SUCCESS) {
        ESP_LOGW(TAG, "Send failed to " MACSTR, MAC2STR(mac));
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

esp_err_t espnow_slave_init(const uint8_t master_mac[6], uint8_t puzzle_id)
{
    s_puzzle_id = puzzle_id;
    memcpy(s_master_mac, master_mac, 6);

    s_recv_queue = xQueueCreate(RECV_QUEUE_LEN, sizeof(recv_item_t));
    if (!s_recv_queue) {
        ESP_LOGE(TAG, "Failed to create recv queue");
        return ESP_ERR_NO_MEM;
    }

    ESP_ERROR_CHECK(esp_now_init());
    ESP_ERROR_CHECK(esp_now_register_recv_cb(on_recv));
    ESP_ERROR_CHECK(esp_now_register_send_cb(on_sent));

    // Register master as peer
    esp_now_peer_info_t peer = {};
    memcpy(peer.peer_addr, master_mac, 6);
    peer.channel = 0;      // use current Wi-Fi channel
    peer.encrypt = false;
    ESP_ERROR_CHECK(esp_now_add_peer(&peer));

    ESP_LOGI(TAG, "Puzzle P%d slave ready, master=" MACSTR,
             puzzle_id, MAC2STR(master_mac));
    return ESP_OK;
}

esp_err_t espnow_slave_notify_solved(const uint8_t code_fragment[4],
                                     uint32_t solve_time_ms)
{
    // Wire format: [MSG_TYPE(1)] [PUZZLE_ID(1)] [CODE(4)] [SOLVE_MS(4)]
    uint8_t buf[10];
    buf[0] = (uint8_t)MSG_PUZZLE_SOLVED;
    buf[1] = s_puzzle_id;
    memcpy(buf + 2, code_fragment, 4);
    buf[6] = (uint8_t)(solve_time_ms >> 24);
    buf[7] = (uint8_t)(solve_time_ms >> 16);
    buf[8] = (uint8_t)(solve_time_ms >>  8);
    buf[9] = (uint8_t)(solve_time_ms      );
    return esp_now_send(s_master_mac, buf, sizeof(buf));
}

esp_err_t espnow_slave_request_hint(void)
{
    uint8_t buf[2] = { (uint8_t)MSG_HINT_REQUEST, s_puzzle_id };
    return esp_now_send(s_master_mac, buf, sizeof(buf));
}

esp_err_t espnow_slave_send_status(const puzzle_status_t *status)
{
    uint8_t buf[1 + sizeof(puzzle_status_t)];
    buf[0] = (uint8_t)MSG_STATUS_REPLY;
    memcpy(buf + 1, status, sizeof(puzzle_status_t));
    return esp_now_send(s_master_mac, buf, sizeof(buf));
}

void espnow_slave_register_callback(espnow_cmd_callback_t cb)
{
    s_callback = cb;
}

void espnow_slave_register_scenario_callback(espnow_scenario_callback_t cb)
{
    s_scn_cb = cb;
}

void espnow_slave_process(void)
{
    // Abandon a half-received scenario whose sender went silent.
    if (s_scn.active && (now_ms() - s_scn.last_ms) > SCENARIO_TIMEOUT_MS) {
        ESP_LOGW(TAG, "scenario rx timeout (%u/%u frames) — dropped",
                 s_scn.count, s_scn.total);
        scenario_reset();
    }

    recv_item_t item;
    while (xQueueReceive(s_recv_queue, &item, 0) == pdTRUE) {
        if (item.data_len < 1) continue;
        if (scenario_is_frame(&item)) {
            scenario_feed(&item);
            continue;
        }
        espnow_msg_type_t type = (espnow_msg_type_t)item.data[0];
        if (s_callback) {
            s_callback(type,
                       item.data + 1,
                       (size_t)(item.data_len - 1));
        }
    }
}
