// scenario_mesh — see include/scenario_mesh.h for the design notes.
//
// Tasks 4 & 5 of the firmware-scenario-hotload spec: ESP-NOW frame protocol
// (chunk + reassembly) for relaying Runtime 3 IR to WiFi-disabled peers.

#include "scenario_mesh.h"

#include <string.h>

#include "esp_log.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/semphr.h"
#include "freertos/task.h"

static const char *TAG = "scenario_mesh";

static const uint8_t kBroadcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};

// ─── alias -> MAC peer table (minimal local registry) ───────────────────────

#define MESH_ALIAS_MAX  32
#define MESH_PEERS_MAX  8

typedef struct {
    bool    used;
    char    alias[MESH_ALIAS_MAX];
    uint8_t mac[6];
} mesh_peer_t;

static mesh_peer_t s_peers[MESH_PEERS_MAX];

// ─── send-side ack synchronization ──────────────────────────────────────────

static SemaphoreHandle_t s_send_done;   // given by on_sent for each frame
static volatile esp_now_send_status_t s_last_send_status;
static SemaphoreHandle_t s_send_lock;   // serializes scenario_mesh_send calls

// ─── receive-side reassembly ────────────────────────────────────────────────
//
// One in-flight reassembly slot per distinct (sender MAC, total) tuple. A
// single slot is enough in practice — the master relays to one peer at a time
// and frames arrive in order — but we keep a small array so two senders (or a
// retried transfer) don't clobber each other.

#define MESH_REASM_SLOTS  2

typedef struct {
    bool     used;
    uint8_t  src[6];
    uint16_t total;
    uint16_t received;        // count of distinct frames stored
    bool    *seen;            // [total] frame-arrival bitmap (heap)
    char    *buf;             // [total * PAYLOAD_MAX + 1] (heap)
    size_t   buf_len;         // running max byte offset written + payload
    uint32_t last_tick;       // for staleness eviction
} mesh_reasm_t;

static mesh_reasm_t              s_reasm[MESH_REASM_SLOTS];
static SemaphoreHandle_t         s_reasm_lock;
static scenario_mesh_apply_cb_t  s_apply_cb;

// A completed reassembly is handed off to a worker task rather than applied in
// the Wi-Fi recv-callback context: the apply does filesystem I/O (and triggers
// a reboot) which must not run on the Wi-Fi stack's callback.
typedef struct {
    char  *buf;     // heap, NUL-terminated, ownership transferred to the task
    size_t len;
} mesh_apply_job_t;

static QueueHandle_t s_apply_queue;

static void apply_worker_task(void *arg) {
    (void) arg;
    mesh_apply_job_t job;
    for (;;) {
        if (xQueueReceive(s_apply_queue, &job, portMAX_DELAY) != pdTRUE) {
            continue;
        }
        if (s_apply_cb) {
            esp_err_t aerr = s_apply_cb(job.buf, job.len);
            if (aerr != ESP_OK) {
                ESP_LOGW(TAG, "apply_cb returned %s", esp_err_to_name(aerr));
            }
        } else {
            ESP_LOGW(TAG, "no apply_cb registered — dropping scenario");
        }
        free(job.buf);
    }
}

// ─── peer table helpers ─────────────────────────────────────────────────────

esp_err_t scenario_mesh_register_peer(const char *alias, const uint8_t mac[6]) {
    if (!alias || !mac) return ESP_ERR_INVALID_ARG;

    int free_slot = -1;
    for (int i = 0; i < MESH_PEERS_MAX; i++) {
        if (s_peers[i].used && strncmp(s_peers[i].alias, alias,
                                       MESH_ALIAS_MAX) == 0) {
            free_slot = i;  // update existing
            break;
        }
        if (!s_peers[i].used && free_slot < 0) free_slot = i;
    }
    if (free_slot < 0) {
        ESP_LOGE(TAG, "peer table full, cannot register \"%s\"", alias);
        return ESP_ERR_NO_MEM;
    }

    s_peers[free_slot].used = true;
    strncpy(s_peers[free_slot].alias, alias, MESH_ALIAS_MAX - 1);
    s_peers[free_slot].alias[MESH_ALIAS_MAX - 1] = '\0';
    memcpy(s_peers[free_slot].mac, mac, 6);

    // Add (or refresh) the ESP-NOW peer entry. esp_now_add_peer fails with
    // ESP_ERR_ESPNOW_EXIST if already present — treat that as success.
    esp_now_peer_info_t pi = {0};
    memcpy(pi.peer_addr, mac, 6);
    pi.channel = 0;        // current channel
    pi.ifidx   = WIFI_IF_STA;
    pi.encrypt = false;
    esp_err_t err = esp_now_add_peer(&pi);
    if (err == ESP_ERR_ESPNOW_EXIST) {
        err = ESP_OK;
    } else if (err != ESP_OK) {
        ESP_LOGW(TAG, "esp_now_add_peer(%s) failed: %s",
                 alias, esp_err_to_name(err));
    }

    ESP_LOGI(TAG, "peer \"%s\" -> %02x:%02x:%02x:%02x:%02x:%02x",
             alias, mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    return err;
}

esp_err_t scenario_mesh_mac_for_alias(const char *alias, uint8_t mac_out[6]) {
    if (!alias || !mac_out) return ESP_ERR_INVALID_ARG;
    for (int i = 0; i < MESH_PEERS_MAX; i++) {
        if (s_peers[i].used &&
            strncmp(s_peers[i].alias, alias, MESH_ALIAS_MAX) == 0) {
            memcpy(mac_out, s_peers[i].mac, 6);
            return ESP_OK;
        }
    }
    return ESP_ERR_NOT_FOUND;
}

// ─── ESP-NOW callbacks (Wi-Fi task context — keep short) ────────────────────

static void on_sent(const uint8_t *mac, esp_now_send_status_t status) {
    (void) mac;
    s_last_send_status = status;
    if (s_send_done) {
        BaseType_t hp = pdFALSE;
        xSemaphoreGiveFromISR(s_send_done, &hp);
        if (hp) portYIELD_FROM_ISR();
    }
}

// Locate (or allocate) the reassembly slot for this sender/total. Caller holds
// s_reasm_lock. Returns NULL on allocation failure.
static mesh_reasm_t *reasm_slot_for(const uint8_t src[6], uint16_t total) {
    mesh_reasm_t *free_slot = NULL;
    mesh_reasm_t *oldest    = NULL;
    for (int i = 0; i < MESH_REASM_SLOTS; i++) {
        mesh_reasm_t *r = &s_reasm[i];
        if (r->used && r->total == total && memcmp(r->src, src, 6) == 0) {
            return r;
        }
        if (!r->used && !free_slot) free_slot = r;
        if (r->used && (!oldest || r->last_tick < oldest->last_tick)) oldest = r;
    }

    // New transfer. Reuse a free slot, else evict the oldest in-flight one.
    mesh_reasm_t *r = free_slot ? free_slot : oldest;
    if (!r) return NULL;
    if (r->used) {
        ESP_LOGW(TAG, "evicting stale reassembly (%u/%u frames)",
                 r->received, r->total);
        free(r->seen);
        free(r->buf);
    }
    memset(r, 0, sizeof(*r));

    if (total == 0 || (size_t) total * SCENARIO_MESH_PAYLOAD_MAX > SCENARIO_MESH_MAX_BLOB) {
        ESP_LOGW(TAG, "reject transfer: implausible total=%u", total);
        return NULL;
    }
    r->seen = calloc(total, sizeof(bool));
    r->buf  = malloc((size_t) total * SCENARIO_MESH_PAYLOAD_MAX + 1);
    if (!r->seen || !r->buf) {
        free(r->seen);
        free(r->buf);
        ESP_LOGE(TAG, "OOM allocating reassembly for total=%u", total);
        return NULL;
    }
    r->used  = true;
    r->total = total;
    memcpy(r->src, src, 6);
    return r;
}

static void on_recv(const esp_now_recv_info_t *info,
                    const uint8_t *data, int len) {
    if (!info || !data || len < SCENARIO_MESH_HEADER_BYTES) return;
    if (len > SCENARIO_MESH_FRAME_MAX) return;

    // Header: seq:u16, total:u16 (little-endian).
    uint16_t seq   = (uint16_t) (data[0] | (data[1] << 8));
    uint16_t total = (uint16_t) (data[2] | (data[3] << 8));
    const uint8_t *payload = data + SCENARIO_MESH_HEADER_BYTES;
    int payload_len = len - SCENARIO_MESH_HEADER_BYTES;
    if (seq >= total) return;  // malformed

    if (!s_reasm_lock) return;
    // Run reassembly off the Wi-Fi callback by doing the bookkeeping under a
    // mutex here; the (potentially slow) apply is deferred to a short task so
    // we never block the Wi-Fi stack inside the recv callback.
    char  *complete_buf = NULL;
    size_t complete_len = 0;

    xSemaphoreTake(s_reasm_lock, portMAX_DELAY);
    mesh_reasm_t *r = reasm_slot_for(info->src_addr, total);
    if (r) {
        if (!r->seen[seq]) {
            r->seen[seq] = true;
            r->received++;
            size_t off = (size_t) seq * SCENARIO_MESH_PAYLOAD_MAX;
            memcpy(r->buf + off, payload, payload_len);
            // Track the highest end offset so the final length is exact even
            // though only the last frame is short.
            if (off + payload_len > r->buf_len) r->buf_len = off + payload_len;
        }
        r->last_tick = (uint32_t) xTaskGetTickCount();

        if (r->received == r->total) {
            r->buf[r->buf_len] = '\0';
            complete_buf = r->buf;
            complete_len = r->buf_len;
            free(r->seen);
            memset(r, 0, sizeof(*r));   // releases the slot; buf handed off
        }
    }
    xSemaphoreGive(s_reasm_lock);

    if (complete_buf) {
        ESP_LOGI(TAG, "reassembled scenario: %u bytes from "
                      "%02x:%02x:%02x:%02x:%02x:%02x",
                 (unsigned) complete_len,
                 info->src_addr[0], info->src_addr[1], info->src_addr[2],
                 info->src_addr[3], info->src_addr[4], info->src_addr[5]);
        // Hand off to the worker task — never touch the filesystem (or reboot)
        // from the Wi-Fi recv-callback context.
        mesh_apply_job_t job = { .buf = complete_buf, .len = complete_len };
        if (!s_apply_queue ||
            xQueueSend(s_apply_queue, &job, 0) != pdTRUE) {
            ESP_LOGW(TAG, "apply queue full/unavailable — dropping scenario");
            free(complete_buf);
        }
    }
}

// ─── init ───────────────────────────────────────────────────────────────────

esp_err_t scenario_mesh_init(scenario_mesh_apply_cb_t apply_cb) {
    s_apply_cb = apply_cb;

    if (!s_send_done)  s_send_done  = xSemaphoreCreateBinary();
    if (!s_send_lock)  s_send_lock  = xSemaphoreCreateMutex();
    if (!s_reasm_lock) s_reasm_lock = xSemaphoreCreateMutex();
    if (!s_send_done || !s_send_lock || !s_reasm_lock) {
        return ESP_ERR_NO_MEM;
    }

    // Receiver path: spin up the apply worker (and its job queue) so completed
    // reassemblies are applied off the Wi-Fi callback. A pure sender (master
    // relay with no apply_cb) skips this to save RAM.
    if (apply_cb && !s_apply_queue) {
        s_apply_queue = xQueueCreate(2, sizeof(mesh_apply_job_t));
        if (!s_apply_queue) return ESP_ERR_NO_MEM;
        BaseType_t ok = xTaskCreate(apply_worker_task, "scn_mesh_apply",
                                    4096, NULL, tskIDLE_PRIORITY + 2, NULL);
        if (ok != pdPASS) {
            vQueueDelete(s_apply_queue);
            s_apply_queue = NULL;
            return ESP_ERR_NO_MEM;
        }
    }

    // esp_now_init() requires Wi-Fi to be started already (the caller brings up
    // STA/AP). A second init returns ESP_ERR_ESPNOW_INTERNAL on some IDF lines;
    // since no other component in this tree owns ESP-NOW yet a hard failure is
    // genuinely fatal, but we keep the relay endpoint optional at the call site.
    esp_err_t err = esp_now_init();
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_now_init: %s", esp_err_to_name(err));
        return err;
    }

    esp_now_register_send_cb(on_sent);
    esp_now_register_recv_cb(on_recv);

    // Broadcast peer (handy for future fan-out; unicast peers are added on
    // demand by scenario_mesh_register_peer).
    esp_now_peer_info_t bcast = {0};
    memcpy(bcast.peer_addr, kBroadcast, 6);
    bcast.channel = 0;
    bcast.ifidx   = WIFI_IF_STA;
    bcast.encrypt = false;
    esp_err_t berr = esp_now_add_peer(&bcast);
    if (berr != ESP_OK && berr != ESP_ERR_ESPNOW_EXIST) {
        ESP_LOGW(TAG, "esp_now_add_peer(broadcast): %s",
                 esp_err_to_name(berr));
    }

    ESP_LOGI(TAG, "scenario_mesh ready (apply_cb=%s)",
             apply_cb ? "set" : "none");
    return ESP_OK;
}

// ─── send (chunk + per-frame ack) ───────────────────────────────────────────

esp_err_t scenario_mesh_send(const uint8_t dest_mac[6],
                             const char *data, size_t len) {
    if (!dest_mac || !data || len == 0) return ESP_ERR_INVALID_ARG;
    if (len > SCENARIO_MESH_MAX_BLOB) return ESP_ERR_INVALID_SIZE;

    size_t total = (len + SCENARIO_MESH_PAYLOAD_MAX - 1) /
                   SCENARIO_MESH_PAYLOAD_MAX;
    if (total == 0 || total > 0xFFFF) return ESP_ERR_INVALID_SIZE;

    // Serialize: the single send-done semaphore is shared across frames.
    if (xSemaphoreTake(s_send_lock, portMAX_DELAY) != pdTRUE) {
        return ESP_FAIL;
    }

    esp_err_t result = ESP_OK;
    uint8_t frame[SCENARIO_MESH_FRAME_MAX];

    for (size_t seq = 0; seq < total; seq++) {
        size_t off   = seq * SCENARIO_MESH_PAYLOAD_MAX;
        size_t chunk = len - off;
        if (chunk > SCENARIO_MESH_PAYLOAD_MAX) chunk = SCENARIO_MESH_PAYLOAD_MAX;

        frame[0] = (uint8_t) (seq & 0xFF);
        frame[1] = (uint8_t) ((seq >> 8) & 0xFF);
        frame[2] = (uint8_t) (total & 0xFF);
        frame[3] = (uint8_t) ((total >> 8) & 0xFF);
        memcpy(frame + SCENARIO_MESH_HEADER_BYTES, data + off, chunk);

        // Drain any stale ack from a previous frame, then send + await ack.
        xSemaphoreTake(s_send_done, 0);
        s_last_send_status = ESP_NOW_SEND_FAIL;

        esp_err_t serr = esp_now_send(dest_mac, frame,
                                      SCENARIO_MESH_HEADER_BYTES + chunk);
        if (serr != ESP_OK) {
            ESP_LOGW(TAG, "esp_now_send frame %u/%u: %s",
                     (unsigned) seq, (unsigned) total, esp_err_to_name(serr));
            result = serr;
            break;
        }

        if (xSemaphoreTake(s_send_done,
                           pdMS_TO_TICKS(SCENARIO_MESH_ACK_TIMEOUT_MS))
                != pdTRUE) {
            ESP_LOGW(TAG, "ack timeout on frame %u/%u",
                     (unsigned) seq, (unsigned) total);
            result = ESP_ERR_TIMEOUT;
            break;
        }
        if (s_last_send_status != ESP_NOW_SEND_SUCCESS) {
            ESP_LOGW(TAG, "frame %u/%u not acked by peer",
                     (unsigned) seq, (unsigned) total);
            result = ESP_ERR_TIMEOUT;  // surfaced as a skipped peer
            break;
        }
    }

    xSemaphoreGive(s_send_lock);

    if (result == ESP_OK) {
        ESP_LOGI(TAG, "sent scenario: %u bytes in %u frames",
                 (unsigned) len, (unsigned) total);
    }
    return result;
}